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}`);
}

/**
 * Fetch one page of spend logs for a team, optionally narrowed to a
 * single virtual key by alias.
 *
 * Slice 2 / Bug 19 context
 * ------------------------
 * Teams in LiteLLM are now org-scoped (one team per org), and each
 * tenant in the org has its own virtual key with `key_alias = tenant
 * CR name`. Without `keyAlias`, this returns the full team's spend —
 * which mingles every tenant in the org. The portal's per-tenant
 * usage view passes `keyAlias` to filter server-side via LiteLLM's
 * native `key_alias` query param. Confirmed available on the
 * `/spend/logs/v2` endpoint via OpenAPI introspection — no need to
 * page-and-post-filter as the previous slice did.
 *
 * Why this matters
 * ----------------
 * Previous implementation fetched all team pages, then post-filtered
 * by alias in JS. Two problems: (1) at any reasonable scale this is
 * O(team_total) memory per request even when only one tenant's data
 * is needed; (2) more importantly, when called from the customer
 * dashboard without an explicit alias, the route's "pick the first
 * visible tenant" fallback meant both Acme tenants showed identical
 * numbers — the alias used was always the first tenant in the
 * visible list, regardless of which tenant page was being viewed.
 *
 * The route layer above is responsible for resolving the tenant
 * identity correctly and passing the right alias here. This
 * function's only job is to pass it through to LiteLLM.
 */
export async function getTeamSpendLogsV2(
  teamId: string,
  startDate: string,
  endDate: string,
  page: number = 1,
  pageSize: number = 100,
  keyAlias?: string | null
) {
  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),
  });
  if (keyAlias) {
    params.set("key_alias", keyAlias);
  }
  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;
}