From 20d8b61aa3f64491afe594fa4dd840a93bf3dd42 Mon Sep 17 00:00:00 2001 From: shangzy Date: Fri, 8 May 2026 17:27:59 +0800 Subject: [PATCH] feat: add user token breakdown data --- lib/queries.test.ts | 81 ++++++++++++++++++++++++++++++++++++++- lib/queries.ts | 93 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 1 deletion(-) diff --git a/lib/queries.test.ts b/lib/queries.test.ts index 8bc138a..bc536b8 100644 --- a/lib/queries.test.ts +++ b/lib/queries.test.ts @@ -16,11 +16,22 @@ mock.module("./db", () => ({ query: queryMock, })); -const { getTrends } = await import("./queries"); +const { getTrends, getUserDetail } = await import("./queries"); describe("getTrends", () => { beforeEach(() => { queryMock.mockClear(); + queryMock.mockImplementation(async () => [ + { + date: "2026-04-01 13:00:00", + calls: 1, + prompt_tokens: 10, + completion_tokens: 20, + cache_creation_tokens: 3, + cache_read_tokens: 4, + quota: 100, + }, + ]); }); test("adds optional detail filters to the trend query", async () => { @@ -44,3 +55,71 @@ describe("getTrends", () => { expect(trends[0].date).toBe("2026-04-01 13:00"); }); }); + +describe("getUserDetail", () => { + beforeEach(() => { + queryMock.mockClear(); + }); + + test("returns token breakdown with nested model rows for user details", async () => { + queryMock.mockImplementation(async (sql: string) => { + if (sql.includes("token_models AS")) { + return [ + { token_name: "prod-key", model: "claude-sonnet-4", calls: 3, tokens: 100, cache_creation: 5, cache_read: 7, quota: 50 }, + { token_name: "", model: "gpt-4o", calls: 1, tokens: 20, cache_creation: 0, cache_read: 0, quota: 10 }, + ]; + } + + if (sql.includes("GROUP BY token_name")) { + return [ + { token_name: "prod-key", calls: 3, tokens: 100, cache_creation: 5, cache_read: 7, quota: 50 }, + { token_name: "", calls: 1, tokens: 20, cache_creation: 0, cache_read: 0, quota: 10 }, + ]; + } + + if (sql.includes("GROUP BY model")) { + return []; + } + + if (sql.includes("SELECT id, display_name FROM users")) { + return []; + } + + if (sql.includes("SELECT id FROM users WHERE username")) { + return []; + } + + if (sql.includes("SELECT COUNT(*)::int as calls")) { + return [ + { calls: 4, prompt: 120, completion: 0, cache_creation: 5, cache_read: 7, quota: 60 }, + ]; + } + + return []; + }); + + const detail = await getUserDetail("token-user", 301, 401); + + expect(detail.tokens).toEqual([ + { + name: "prod-key", + calls: 3, + total_tokens: 112, + quota: 50, + models: [ + { name: "claude-sonnet-4", calls: 3, total_tokens: 112, quota: 50 }, + ], + }, + { + name: "", + calls: 1, + total_tokens: 20, + quota: 10, + models: [ + { name: "gpt-4o", calls: 1, total_tokens: 20, quota: 10 }, + ], + }, + ]); + expect(queryMock.mock.calls.some(([sql]) => String(sql).includes("token_name"))).toBe(true); + }); +}); diff --git a/lib/queries.ts b/lib/queries.ts index f87399e..c6b273e 100644 --- a/lib/queries.ts +++ b/lib/queries.ts @@ -33,6 +33,8 @@ const CACHE_READ = `COALESCE( THEN (other::jsonb->>'cache_tokens')::bigint END, 0)`; +const TOKEN_NAME = `COALESCE(NULLIF(BTRIM(token_name), ''), '')`; + // ── 数据时间边界 ──────────────────────────────────────────────── export async function getDateRange(): Promise<{ minDate: string; maxDate: string }> { @@ -366,6 +368,10 @@ export interface DetailBreakdown { quota: number; } +export interface TokenDetailBreakdown extends DetailBreakdown { + models: DetailBreakdown[]; +} + export function getUserDetail( username: string, startTs?: number, @@ -406,6 +412,83 @@ export function getUserDetail( params2 ); + const params3: (string | number | boolean | null)[] = []; + const where3 = timeWhere(params3, startTs, endTs); + params3.push(username); + + const tokens = await query( + `SELECT ${TOKEN_NAME} as token_name, + COUNT(*)::int as calls, + COALESCE(SUM(prompt_tokens + completion_tokens),0)::bigint as tokens, + COALESCE(SUM(${CACHE_CREATION}),0)::bigint as cache_creation, + COALESCE(SUM(${CACHE_READ}),0)::bigint as cache_read, + COALESCE(SUM(quota),0)::bigint as quota + FROM logs WHERE ${where3} AND username = $${params3.length} + GROUP BY token_name + ORDER BY COALESCE(SUM(prompt_tokens),0) + COALESCE(SUM(completion_tokens),0) + COALESCE(SUM(${CACHE_CREATION}),0) + COALESCE(SUM(${CACHE_READ}),0) DESC, token_name ASC + LIMIT 50`, + params3 + ); + + const params4: (string | number | boolean | null)[] = []; + const where4 = timeWhere(params4, startTs, endTs); + params4.push(username); + + const tokenModels = await query( + `WITH filtered_logs AS ( + SELECT ${TOKEN_NAME} as token_name, + ${REAL_MODEL} as model, + prompt_tokens, + completion_tokens, + ${CACHE_CREATION} as cache_creation, + ${CACHE_READ} as cache_read, + quota + FROM logs WHERE ${where4} AND username = $${params4.length} + ), + top_tokens AS ( + SELECT token_name + FROM filtered_logs + GROUP BY token_name + ORDER BY COALESCE(SUM(prompt_tokens),0) + COALESCE(SUM(completion_tokens),0) + COALESCE(SUM(cache_creation),0) + COALESCE(SUM(cache_read),0) DESC, token_name ASC + LIMIT 50 + ), + token_models AS ( + SELECT filtered_logs.token_name, + filtered_logs.model, + COUNT(*)::int as calls, + COALESCE(SUM(prompt_tokens + completion_tokens),0)::bigint as tokens, + COALESCE(SUM(cache_creation),0)::bigint as cache_creation, + COALESCE(SUM(cache_read),0)::bigint as cache_read, + COALESCE(SUM(quota),0)::bigint as quota + FROM filtered_logs + INNER JOIN top_tokens ON filtered_logs.token_name = top_tokens.token_name + GROUP BY filtered_logs.token_name, filtered_logs.model + ), + ranked AS ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY token_name ORDER BY tokens + cache_creation + cache_read DESC, model ASC) as rn + FROM token_models + ) + SELECT token_name, model, calls, tokens, cache_creation, cache_read, quota + FROM ranked + WHERE rn <= 10 + ORDER BY token_name ASC, tokens + cache_creation + cache_read DESC, model ASC`, + params4 + ); + + const modelsByToken = new Map(); + for (const row of tokenModels) { + const tokenName = String(row.token_name ?? ""); + const rows = modelsByToken.get(tokenName) ?? []; + rows.push({ + name: String(row.model || "(unknown)"), + calls: Number(row.calls), + total_tokens: Number(row.tokens) + Number(row.cache_creation) + Number(row.cache_read), + quota: Number(row.quota), + }); + modelsByToken.set(tokenName, rows); + } + const displayNames = await getDisplayNames(); const userRow = await query( "SELECT id FROM users WHERE username = $1 LIMIT 1", @@ -430,6 +513,16 @@ export function getUserDetail( total_tokens: Number(m.tokens) + Number(m.cache_creation) + Number(m.cache_read), quota: Number(m.quota), })), + tokens: tokens.map((token): TokenDetailBreakdown => { + const name = String(token.token_name ?? ""); + return { + name, + calls: Number(token.calls), + total_tokens: Number(token.tokens) + Number(token.cache_creation) + Number(token.cache_read), + quota: Number(token.quota), + models: modelsByToken.get(name) ?? [], + }; + }), }; }); }