pieced-portal/src/lib/litellm.ts

const LITELLM_URL =
  process.env.LITELLM_INTERNAL_URL ?? "http://litellm.inference.svc:4000";
const LITELLM_MASTER_KEY = process.env.LITELLM_MASTER_KEY!;

async function litellmFetch(path: string, init?: RequestInit) {
  const res = await fetch(`${LITELLM_URL}${path}`, {
    ...init,
    headers: {
      Authorization: `Bearer ${LITELLM_MASTER_KEY}`,
      "Content-Type": "application/json",
      ...init?.headers,
    },
  });
  if (!res.ok) {
    throw new Error(`LiteLLM ${path}: ${res.status} ${await res.text()}`);
  }
  return res.json();
}

export async function getTeamInfo(teamId: string) {
  return litellmFetch(`/team/info?team_id=${encodeURIComponent(teamId)}`);
}

export async function getTeamSpendLogs(
  teamId: string,
  startDate?: string,
  endDate?: string
) {
  const params = new URLSearchParams({ team_id: teamId });
  if (startDate) params.set("start_date", startDate);
  if (endDate) params.set("end_date", endDate);
  return litellmFetch(`/global/spend/logs?${params}`);
}

export async function getTeamSpendLogsV2(
  teamId: string,
  startDate: string,
  endDate: string,
  page: number = 1,
  pageSize: number = 100
) {
  const params = new URLSearchParams({
    team_id: teamId,
    start_date: `${startDate} 00:00:00`,
    end_date: `${endDate} 23:59:59`,
    page: String(page),
    page_size: String(pageSize),
  });
  return litellmFetch(`/spend/logs/v2?${params}`);
}

/**
 * Get all teams registered in LiteLLM.
 * Returns team_id, spend, max_budget, etc.
 */
export async function listTeams(): Promise<any[]> {
  const data = await litellmFetch("/team/list");
  // LiteLLM returns either an array or { data: [...] }
  return Array.isArray(data) ? data : data?.data ?? data?.teams ?? [];
}

/**
 * Get LiteLLM health status.
 */
export async function getLitellmHealth(): Promise<{
  healthy: boolean;
  details?: any;
}> {
  try {
    const data = await litellmFetch("/health");
    return { healthy: true, details: data };
  } catch (e: any) {
    return { healthy: false, details: e.message };
  }
}

/**
 * Get global spend across all teams for the current month.
 */
export async function getGlobalSpend(): Promise<number> {
  try {
    const data = await litellmFetch("/global/spend");
    // LiteLLM returns { spend: number } or similar
    if (typeof data === "number") return data;
    return data?.spend ?? data?.total_spend ?? 0;
  } catch {
    return 0;
  }
}

/**
 * Fetch per-team spend as a map: teamId → spend (CHF).
 * Uses /team/list which includes current spend per team.
 *
 * Since Slice 2, a "team" is the company-level budget shared across all
 * tenants of the same ZITADEL org. So this map gives company totals, not
 * per-tenant spend. For per-tenant attribution, use {@link getPerKeySpend}.
 */
export async function getPerTeamSpend(): Promise<Map<string, number>> {
  const teams = await listTeams();
  const map = new Map<string, number>();
  for (const team of teams) {
    const id = team.team_id ?? team.id;
    const spend = team.spend ?? 0;
    if (id) map.set(id, spend);
  }
  return map;
}

/**
 * Fetch per-virtual-key spend as a map: keyAlias → spend (CHF).
 *
 * Since Slice 2, each PiecedTenant CR owns one virtual key under its
 * org's team, with `key_alias = tenant.metadata.name`. Filtering by the
 * key alias is how we get genuinely per-tenant spend.
 *
 * Implementation
 * --------------
 * Calls `/key/list?return_full_object=true&include_team_keys=true`,
 * which returns objects with `spend` and `key_alias`. Older LiteLLM
 * builds may return raw token strings instead — we degrade gracefully
 * to an empty map in that case rather than throwing, since the admin
 * health page should still render even if per-tenant numbers are
 * temporarily unavailable.
 *
 * @returns Map<keyAlias, spend>. May be empty if the LiteLLM build
 *          doesn't expose key-alias info; callers must handle that.
 */
export async function getPerKeySpend(): Promise<Map<string, number>> {
  const map = new Map<string, number>();
  try {
    const data = await litellmFetch(
      "/key/list?return_full_object=true&include_team_keys=true"
    );

    // Response shape: { keys: [ { key_alias, spend, token, ... } ] }
    // or sometimes { data: [...] }, or raw arrays. Be tolerant.
    const keys: any[] = Array.isArray(data?.keys)
      ? data.keys
      : Array.isArray(data?.data)
        ? data.data
        : Array.isArray(data)
          ? data
          : [];

    for (const k of keys) {
      // Skip raw-string entries from older API shapes — we can't attribute them.
      if (typeof k !== "object" || k === null) continue;
      const alias = k.key_alias ?? k.keyAlias;
      if (typeof alias !== "string" || !alias) continue;
      const spend =
        typeof k.spend === "number" ? k.spend : Number(k.spend) || 0;
      map.set(alias, spend);
    }
  } catch (e) {
    console.warn("getPerKeySpend failed, returning empty map:", e);
  }
  return map;
}