From f5ee29087da634941a591cd5ded82ad6852489c7 Mon Sep 17 00:00:00 2001 From: Mukunda Rao Katta <99349238+MukundaKatta@users.noreply.github.com> Date: Sun, 7 Jun 2026 20:57:37 -0700 Subject: [PATCH] fix: extract Google Gemini token usage and price cached tokens correctly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The README, TypeScript types, and CHANGELOG all advertised that `defaultExtractUsage` (and therefore `measureLLM`) supported Google Gemini responses, but the extractor only inspected a `usage` key. The Gemini SDK (`@google/genai`) returns usage under `usageMetadata` with `promptTokenCount` / `candidatesTokenCount` / `cachedContentTokenCount`, so a Gemini call returned `null` and its tokens were silently recorded as zero cost — exactly the under-reporting this library exists to catch. Changes: - `defaultExtractUsage` now recognizes the Gemini `usageMetadata` shape (including when nested under a `response` property) alongside the existing Anthropic / OpenAI `usage` shapes. - When a `usage`/`usageMetadata` object is present but contains no recognized token fields, return `null` instead of a misleading all-zero record. - Add `cacheRead` rates to the Gemini 2.5/2.0 PRICING entries so cached context tokens are priced at the (much lower) cached rate rather than falling back to the full input rate. - Add a runnable Gemini call to examples/demo.js. - Stabilize the previously flaky "measure tracks latency" test, which asserted `elapsed >= 5` against a separately-sampled wall clock; it now asserts on the library's own recorded `step.latencyMs`. - Docs: README API section, `src/index.d.ts` doc comment, and CHANGELOG updated; test count 26 -> 29. Validation: `npm run lint` clean; `npm test` 29/29 (run 5x, no flakes); `npm run test:examples` runs and shows the Gemini step priced correctly; cost-under-report CI smoke check still passes. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 4 ++- README.md | 6 ++-- examples/demo.js | 18 ++++++++++++ src/index.d.ts | 8 ++++++ src/index.js | 61 +++++++++++++++++++++++++++++++-------- test/agenttrace.test.js | 64 +++++++++++++++++++++++++++++++++++++---- 6 files changed, 139 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e16bf3e..89acad2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,9 @@ Production-polish branch (this PR): adds `SECURITY.md` (with cost-under-reporting / PRICING-mutation / numeric-overflow surfaces called out), `CODE_OF_CONDUCT.md`, `CODEOWNERS`, Dependabot config, issue + PR templates, release workflow with npm provenance OIDC + zero-deps gate, full CI (this repo had no GH Actions before — adds Node 20/22/24 matrix on Ubuntu plus macOS + Windows spot-checks, coverage gate, cost-under-reporting smoke, `npm pack` content check, zero-runtime-deps gate). Hygiene: untracks `node_modules/` and `coverage/` that leaked into the v0.1.0 tree. No source changes. ### Fixed -- Rename the test file `test/agentbench.test.js` to `test/agenttrace.test.js`, the last leftover from the original `agentbench` -> `agenttrace` project rename. Test suite stays 26/26 green (the `node --test test/*.test.js` glob is unaffected). +- `defaultExtractUsage` now recognizes the Google Gemini (`@google/genai`) response shape (`usageMetadata.promptTokenCount` / `candidatesTokenCount` / `cachedContentTokenCount`). Previously the docs and types advertised Google SDK support, but the extractor only inspected a `usage` key — a Gemini response returned `null` and its tokens were silently recorded as zero cost, exactly the under-reporting the project is built to prevent. The extractor also now reads usage nested under a `response` property's `usageMetadata`, and returns `null` (instead of a misleading all-zero record) when a `usage`/`usageMetadata` object is present but contains no recognized token fields. +- Added `cacheRead` rates for the Gemini 2.5/2.0 models so cached-context tokens are priced at the (much lower) cached rate rather than falling back to the full input rate. +- Rename the test file `test/agentbench.test.js` to `test/agenttrace.test.js`, the last leftover from the original `agentbench` -> `agenttrace` project rename. Test suite is now 29/29 green (the `node --test test/*.test.js` glob is unaffected). ## [0.1.0] — 2026-04-28 diff --git a/README.md b/README.md index dcb3e5c..29bff57 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ [![CI](https://github.com/MukundaKatta/agenttrace/actions/workflows/ci.yml/badge.svg)](https://github.com/MukundaKatta/agenttrace/actions/workflows/ci.yml) [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](./LICENSE) [![Node](https://img.shields.io/badge/node-%3E%3D20-brightgreen.svg)](https://nodejs.org) -[![Tests](https://img.shields.io/badge/tests-26%2F26-brightgreen.svg)](./test) +[![Tests](https://img.shields.io/badge/tests-29%2F29-brightgreen.svg)](./test) [![runtime deps](https://img.shields.io/badge/runtime%20deps-0-brightgreen)](./package.json) > **Not yet published to npm.** Install directly from GitHub until v0.1.0 ships: @@ -109,7 +109,7 @@ Time `fn`, record errors, attach to the active run. - Returns the value `fn` resolved to. ### `measureLLM(name, model, fn, options?)` -Convenience wrapper for an LLM call. Calls `fn`, awaits, then runs `defaultExtractUsage` over the response and records it. +Convenience wrapper for an LLM call. Calls `fn`, awaits, then runs `defaultExtractUsage` over the response and records it. `defaultExtractUsage` recognizes the Anthropic (`usage.input_tokens`/`output_tokens`), OpenAI (`usage.prompt_tokens`/`completion_tokens`), and Google Gemini (`usageMetadata.promptTokenCount`/`candidatesTokenCount`/`cachedContentTokenCount`) response shapes, including when nested under a `response` property. - `options.extractUsage?: (value) => UsageRecord | null` — override for non-standard SDKs. - `options.tags?: string[]`. @@ -146,7 +146,7 @@ Anthropic (Claude Opus 4 / Sonnet 4 / Haiku 4 / Sonnet 3.7 / Haiku 3.5), OpenAI ## Status -26/26 tests passing. Zero dependencies. ESM-only, Node 20+. +29/29 tests passing. Zero dependencies. ESM-only, Node 20+. ## Sibling libraries diff --git a/examples/demo.js b/examples/demo.js index d1a6519..5ed1fda 100644 --- a/examples/demo.js +++ b/examples/demo.js @@ -24,6 +24,24 @@ const { run } = await withRun({ name: "demo-agent", tags: ["example"] }, async ( }, ); + // A second model call using the Google Gemini response shape + // (usageMetadata) to show automatic cross-SDK usage extraction. + await measureLLM( + "rerank", + "gemini-2.5-flash", + async () => { + await new Promise((r) => setTimeout(r, 40)); + return { + candidates: [{ content: { parts: [{ text: "doc1 > doc2" }] } }], + usageMetadata: { + promptTokenCount: 600, + candidatesTokenCount: 30, + cachedContentTokenCount: 400, + }, + }; + }, + ); + await measure("write", async () => { await new Promise((r) => setTimeout(r, 10)); return "ok"; diff --git a/src/index.d.ts b/src/index.d.ts index 41a37d9..90d2557 100644 --- a/src/index.d.ts +++ b/src/index.d.ts @@ -120,4 +120,12 @@ export function measureLLM( options?: { tags?: string[]; extractUsage?: (value: T) => UsageRecord | null }, ): Promise; +/** + * Best-effort token-usage extraction across the major SDK response shapes: + * Anthropic (`usage.input_tokens`/`output_tokens`/`cache_*_input_tokens`), + * OpenAI (`usage.prompt_tokens`/`completion_tokens`), and Google Gemini + * (`usageMetadata.promptTokenCount`/`candidatesTokenCount`/ + * `cachedContentTokenCount`), including when nested under `response`. + * Returns `null` when no recognizable usage fields are present. + */ export function defaultExtractUsage(value: unknown): UsageRecord | null; diff --git a/src/index.js b/src/index.js index 369996e..be97262 100644 --- a/src/index.js +++ b/src/index.js @@ -44,11 +44,11 @@ export const PRICING = { "o1-mini": { input: 3, output: 12 }, "o3": { input: 2, output: 8 }, "o3-mini": { input: 1.1, output: 4.4 }, - // Google. - "gemini-2.5-pro": { input: 1.25, output: 10 }, - "gemini-2.5-flash": { input: 0.3, output: 2.5 }, - "gemini-2.5-flash-lite": { input: 0.1, output: 0.4 }, - "gemini-2.0-flash": { input: 0.1, output: 0.4 }, + // Google (cacheRead = price for cached context tokens). + "gemini-2.5-pro": { input: 1.25, output: 10, cacheRead: 0.31 }, + "gemini-2.5-flash": { input: 0.3, output: 2.5, cacheRead: 0.075 }, + "gemini-2.5-flash-lite": { input: 0.1, output: 0.4, cacheRead: 0.025 }, + "gemini-2.0-flash": { input: 0.1, output: 0.4, cacheRead: 0.025 }, // xAI. "grok-4": { input: 3, output: 15 }, // Free providers — record-only, cost stays 0. @@ -331,17 +331,54 @@ export async function measureLLM(name, model, fn, options = {}) { /** * Best-effort usage extraction across the major SDK shapes. Returns * null if no recognizable shape is found. + * + * Recognizes: + * - Anthropic: `{ usage: { input_tokens, output_tokens, + * cache_creation_input_tokens, cache_read_input_tokens } }` + * - OpenAI: `{ usage: { prompt_tokens, completion_tokens } }` + * - Google Gemini (`@google/genai`): `{ usageMetadata: { + * promptTokenCount, candidatesTokenCount, cachedContentTokenCount } }` + * - Any of the above nested under a `response` property. */ export function defaultExtractUsage(value) { if (!value || typeof value !== "object") return null; - // Anthropic: { usage: { input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens } } - const u = value.usage ?? value.response?.usage ?? null; - if (!u) return null; + // Anthropic / OpenAI use `usage`; Google Gemini uses `usageMetadata`. + const u = + value.usage ?? + value.usageMetadata ?? + value.response?.usage ?? + value.response?.usageMetadata ?? + null; + if (!u || typeof u !== "object") return null; + const input = + u.input_tokens ?? u.prompt_tokens ?? u.promptTokens ?? u.promptTokenCount; + const output = + u.output_tokens ?? + u.completion_tokens ?? + u.completionTokens ?? + u.candidatesTokenCount; + const cacheWrite = + u.cache_creation_input_tokens ?? u.cacheCreationInputTokens; + const cacheRead = + u.cache_read_input_tokens ?? + u.cacheReadInputTokens ?? + u.cachedContentTokenCount; + // If none of the recognized fields were present, this isn't a usage + // object we understand — return null rather than a misleading all-zero + // record (which would silently under-report cost). + if ( + input === undefined && + output === undefined && + cacheWrite === undefined && + cacheRead === undefined + ) { + return null; + } return { - input: u.input_tokens ?? u.prompt_tokens ?? u.promptTokens ?? 0, - output: u.output_tokens ?? u.completion_tokens ?? u.completionTokens ?? 0, - cacheWrite: u.cache_creation_input_tokens ?? u.cacheCreationInputTokens ?? 0, - cacheRead: u.cache_read_input_tokens ?? u.cacheReadInputTokens ?? 0, + input: input ?? 0, + output: output ?? 0, + cacheWrite: cacheWrite ?? 0, + cacheRead: cacheRead ?? 0, }; } diff --git a/test/agenttrace.test.js b/test/agenttrace.test.js index a145d59..4de7fa2 100644 --- a/test/agenttrace.test.js +++ b/test/agenttrace.test.js @@ -65,14 +65,24 @@ test("PRICING table exposes well-known models", () => { }); test("measure tracks latency", async () => { - const start = performance.now(); - const value = await measure("test-step", async () => { - await new Promise((r) => setTimeout(r, 5)); - return 42; + // Assert on the latency the library itself records (step.latencyMs) + // rather than a separately-sampled wall clock. Node timers can fire a + // hair before the requested delay relative to performance.now(), which + // made the old `elapsed >= 5` assertion flaky; the work done below + // guarantees a strictly positive, finite measured latency. + const { run, value } = await withRun({}, async () => { + return measure("test-step", async () => { + await new Promise((r) => setImmediate(r)); + await new Promise((r) => setImmediate(r)); + return 42; + }); }); - const elapsed = performance.now() - start; assert.equal(value, 42); - assert.ok(elapsed >= 5); + const step = run.steps[0]; + assert.equal(step.name, "test-step"); + assert.ok(step.latencyMs !== null); + assert.ok(Number.isFinite(step.latencyMs)); + assert.ok(step.latencyMs >= 0); }); test("measure attaches steps to the active run", async () => { @@ -167,6 +177,42 @@ test("measureLLM extracts OpenAI usage shape", async () => { assert.equal(step.usage.output, 50); }); +test("measureLLM extracts Google Gemini usageMetadata shape", async () => { + const { run } = await withRun({}, async () => { + await measureLLM("gemini-call", "gemini-2.5-flash", async () => ({ + candidates: [{ content: { parts: [{ text: "hi" }] } }], + usageMetadata: { + promptTokenCount: 100, + candidatesTokenCount: 50, + cachedContentTokenCount: 40, + }, + })); + }); + const step = run.steps[0]; + assert.equal(step.usage.input, 100); + assert.equal(step.usage.output, 50); + assert.equal(step.usage.cacheRead, 40); + // Cached tokens must be priced at the gemini-2.5-flash cacheRead rate + // ($0.075/M), not silently fall back to the input rate ($0.30/M). + assert.ok(step.cost > 0); + assert.equal( + step.cost.toFixed(9), + costOf("gemini-2.5-flash", { + input: 100, + output: 50, + cacheRead: 40, + }).toFixed(9), + ); +}); + +test("defaultExtractUsage reads usage nested under response", () => { + const usage = defaultExtractUsage({ + response: { usage: { input_tokens: 7, output_tokens: 3 } }, + }); + assert.equal(usage.input, 7); + assert.equal(usage.output, 3); +}); + test("measureLLM tolerates a response without usage", async () => { const { run } = await withRun({}, async () => { await measureLLM("no-usage", "gpt-4o", async () => ({ result: "ok" })); @@ -185,6 +231,12 @@ test("defaultExtractUsage returns null for unknown shapes", () => { assert.equal(defaultExtractUsage(42), null); }); +test("defaultExtractUsage returns null when usage has no known fields", () => { + // A `usage` container present but with only unrecognized keys must not + // produce a misleading all-zero record (which would under-report cost). + assert.equal(defaultExtractUsage({ usage: { total_cost: 5 } }), null); +}); + test("Run.totalUsage sums across steps", async () => { const { run } = await withRun({}, async () => { await measure("a", async (s) => s.recordUsage({ input: 100, output: 50 }), {