146 lines
4.1 KiB
TypeScript
146 lines
4.1 KiB
TypeScript
import { describe, expect, it } from "vitest";
|
|
import { parseUsage } from "../src/llm/client.js";
|
|
import {
|
|
recordTokenUsage,
|
|
summarizeTokenUsage,
|
|
} from "../src/stats/token-store.js";
|
|
import { TokenTrackingProvider, type LlmTrackingContext } from "../src/llm/token-tracker.js";
|
|
import type { CompleteResult, LlmProvider } from "../src/llm/client.js";
|
|
|
|
describe("parseUsage", () => {
|
|
it("parses standard OpenAI usage", () => {
|
|
const usage = parseUsage({
|
|
prompt_tokens: 100,
|
|
completion_tokens: 50,
|
|
total_tokens: 150,
|
|
});
|
|
expect(usage).toEqual({
|
|
promptTokens: 100,
|
|
completionTokens: 50,
|
|
totalTokens: 150,
|
|
});
|
|
});
|
|
|
|
it("parses OpenAI cached_tokens via prompt_tokens_details", () => {
|
|
const usage = parseUsage({
|
|
prompt_tokens: 1000,
|
|
completion_tokens: 20,
|
|
total_tokens: 1020,
|
|
prompt_tokens_details: { cached_tokens: 800 },
|
|
});
|
|
expect(usage?.cachedTokens).toBe(800);
|
|
expect(usage?.cacheMissTokens).toBeUndefined();
|
|
});
|
|
|
|
it("parses DeepSeek cache hit and miss fields", () => {
|
|
const usage = parseUsage({
|
|
prompt_tokens: 500,
|
|
completion_tokens: 30,
|
|
total_tokens: 530,
|
|
prompt_cache_hit_tokens: 400,
|
|
prompt_cache_miss_tokens: 100,
|
|
});
|
|
expect(usage?.cachedTokens).toBe(400);
|
|
expect(usage?.cacheMissTokens).toBe(100);
|
|
});
|
|
|
|
it("returns undefined for empty usage", () => {
|
|
expect(parseUsage(null)).toBeUndefined();
|
|
expect(parseUsage({})).toBeUndefined();
|
|
});
|
|
});
|
|
|
|
describe("TokenTrackingProvider", () => {
|
|
it("stores pending usage with cache fields on context", async () => {
|
|
const inner: LlmProvider = {
|
|
async complete(): Promise<CompleteResult> {
|
|
return {
|
|
content: "ok",
|
|
model: "test-model",
|
|
usage: {
|
|
promptTokens: 10,
|
|
completionTokens: 5,
|
|
totalTokens: 15,
|
|
cachedTokens: 8,
|
|
cacheMissTokens: 2,
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
const ctx: LlmTrackingContext = { sessionId: "sess-1" };
|
|
const provider = new TokenTrackingProvider(inner, () => ctx);
|
|
await provider.complete([], { caller: "worker:write-rules" });
|
|
|
|
expect(ctx.pendingUsage).toMatchObject({
|
|
totalTokens: 15,
|
|
cachedTokens: 8,
|
|
cacheMissTokens: 2,
|
|
caller: "worker:write-rules",
|
|
model: "test-model",
|
|
});
|
|
expect(ctx.pendingUsage?.at).toBeTruthy();
|
|
expect(ctx.pendingUsage?.recordId).toBeTruthy();
|
|
});
|
|
|
|
it("stores pending reasoning when usage is absent", async () => {
|
|
const inner: LlmProvider = {
|
|
async complete(): Promise<CompleteResult> {
|
|
return {
|
|
content: "ok",
|
|
model: "test-model",
|
|
reasoning: " chain of thought ",
|
|
};
|
|
},
|
|
};
|
|
|
|
const ctx: LlmTrackingContext = { sessionId: "sess-2" };
|
|
const provider = new TokenTrackingProvider(inner, () => ctx);
|
|
await provider.complete([], { caller: "worker:test" });
|
|
|
|
expect(ctx.pendingUsage).toBeUndefined();
|
|
expect(ctx.pendingReasoning).toBe("chain of thought");
|
|
});
|
|
});
|
|
|
|
describe("summarizeTokenUsage", () => {
|
|
it("aggregates cache totals in byCallerDetailed", () => {
|
|
const sessionId = `test-${Date.now()}`;
|
|
recordTokenUsage({
|
|
sessionId,
|
|
caller: "main_agent",
|
|
model: "m",
|
|
promptTokens: 100,
|
|
completionTokens: 10,
|
|
totalTokens: 110,
|
|
cachedTokens: 50,
|
|
});
|
|
recordTokenUsage({
|
|
sessionId,
|
|
caller: "worker:write-rules",
|
|
model: "m",
|
|
promptTokens: 200,
|
|
completionTokens: 20,
|
|
totalTokens: 220,
|
|
cachedTokens: 100,
|
|
cacheMissTokens: 100,
|
|
});
|
|
|
|
const updated = summarizeTokenUsage({ sessionId, limit: 10 });
|
|
expect(updated.totalCalls).toBe(2);
|
|
expect(updated.totalCached).toBe(150);
|
|
expect(updated.totalCacheMiss).toBe(100);
|
|
expect(updated.byCallerDetailed["main_agent"]).toMatchObject({
|
|
totalTokens: 110,
|
|
cachedTokens: 50,
|
|
calls: 1,
|
|
});
|
|
expect(updated.byCallerDetailed["worker:write-rules"]).toMatchObject({
|
|
totalTokens: 220,
|
|
cachedTokens: 100,
|
|
cacheMissTokens: 100,
|
|
calls: 1,
|
|
});
|
|
});
|
|
});
|