diff --git a/.changeset/gemini-cache-reporting-exclusive.md b/.changeset/gemini-cache-reporting-exclusive.md new file mode 100644 index 0000000000..231ae16c4d --- /dev/null +++ b/.changeset/gemini-cache-reporting-exclusive.md @@ -0,0 +1,5 @@ +--- +'@posthog/ai': patch +--- + +fix(gemini): declare Gemini's cache accounting model on generations with cache reads, so ingestion prices cached tokens from `$ai_cache_reporting_exclusive` instead of inferring it from the token counts diff --git a/packages/ai/src/captureAiGeneration.ts b/packages/ai/src/captureAiGeneration.ts index 57e0bca0da..e11d88c8e1 100644 --- a/packages/ai/src/captureAiGeneration.ts +++ b/packages/ai/src/captureAiGeneration.ts @@ -4,7 +4,7 @@ import { uuidv7, ErrorTracking as CoreErrorTracking, toJsonSafeValue } from '@po import { version } from '../package.json' import type { TokenUsage } from './types' import { stringifyError } from './serializeError' -import { AIEvent, CostOverride, getTokensSource, withPrivacyMode } from './utils' +import { AIEvent, CostOverride, getTokensSource, hasTokenOverrides, withPrivacyMode } from './utils' import { warnIfPostHogAiGateway } from './gatewayWarning' /** @@ -161,10 +161,26 @@ export const captureAiGeneration = async (client: PostHog, options: CaptureAiGen } } + // The caller's own token counts override the SDK-derived ones further down, via the + // `options.properties` spread. + const tokensOverridden = hasTokenOverrides(options.properties) + const additionalTokenValues = { ...(usage.reasoningTokens ? { $ai_reasoning_tokens: usage.reasoningTokens } : {}), ...(usage.cacheReadInputTokens ? { $ai_cache_read_input_tokens: usage.cacheReadInputTokens } : {}), ...(usage.cacheCreationInputTokens ? { $ai_cache_creation_input_tokens: usage.cacheCreationInputTokens } : {}), + // Checked against undefined rather than truthiness, because false is the meaningful + // value here and a truthiness guard would drop it. + // + // Dropped entirely when the caller overrides the token counts: the flag describes how + // the SDK-derived counts relate to each other, so against passthrough counts it can be + // wrong in the expensive direction. Declaring inclusive over counts that are actually + // exclusive makes ingestion subtract the cache pool that was never in the input. A + // caller who knows their own accounting model can still pass + // `$ai_cache_reporting_exclusive` themselves, and that value wins. + ...(usage.cacheReportingExclusive !== undefined && !tokensOverridden + ? { $ai_cache_reporting_exclusive: usage.cacheReportingExclusive } + : {}), ...(usage.webSearchCount ? { $ai_web_search_count: usage.webSearchCount } : {}), ...(usage.rawUsage ? { $ai_usage: usage.rawUsage } : {}), } diff --git a/packages/ai/src/gemini/index.ts b/packages/ai/src/gemini/index.ts index 12ef5783fc..06877a9aca 100644 --- a/packages/ai/src/gemini/index.ts +++ b/packages/ai/src/gemini/index.ts @@ -80,6 +80,11 @@ export class WrappedModels { (metadata as GenerateContentResponseUsageMetadata & { thoughtsTokenCount?: number })?.thoughtsTokenCount ?? 0, cacheReadInputTokens: metadata?.cachedContentTokenCount ?? 0, + // Gemini counts cachedContentTokenCount inside promptTokenCount, so declare the + // accounting model rather than leaving ingestion to infer it. Under explicit + // context caching the two counts come from separate measurements and can disagree + // by a few percent, which makes inference from the counts alone unreliable. + ...(metadata?.cachedContentTokenCount ? { cacheReportingExclusive: false } : {}), webSearchCount: calculateGoogleWebSearchCount(response), rawUsage: metadata, }, @@ -196,6 +201,9 @@ export class WrappedModels { (metadata as GenerateContentResponseUsageMetadata & { thoughtsTokenCount?: number }).thoughtsTokenCount ?? 0, cacheReadInputTokens: metadata.cachedContentTokenCount ?? 0, + // See the non-streaming path: Gemini counts cachedContentTokenCount inside + // promptTokenCount, so the accounting model is declared rather than inferred. + ...(metadata.cachedContentTokenCount ? { cacheReportingExclusive: false } : {}), webSearchCount: usage.webSearchCount, rawUsage: metadata, } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 33b8bd98bf..a6ef560c47 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -95,6 +95,10 @@ export interface TokenUsage { reasoningTokens?: unknown // Use unknown since various providers return different types cacheReadInputTokens?: unknown // Use unknown for provider flexibility cacheCreationInputTokens?: unknown // Use unknown for provider flexibility + // Whether cache tokens are counted separately from inputTokens. Providers that report + // them as a subset of inputTokens set this false. Left undefined when the provider's + // accounting model is not known, in which case ingestion infers it from the counts. + cacheReportingExclusive?: boolean webSearchCount?: number // Count of web search queries/calls used rawUsage?: unknown // Raw provider usage metadata for backend processing } diff --git a/packages/ai/src/utils.ts b/packages/ai/src/utils.ts index 9c4a54f53d..cfbe4a6b3a 100644 --- a/packages/ai/src/utils.ts +++ b/packages/ai/src/utils.ts @@ -28,11 +28,16 @@ const TOKEN_PROPERTY_KEYS = new Set([ '$ai_reasoning_tokens', ]) +/** + * Whether the caller supplied their own token counts, which override the ones the SDK + * derived from the provider response. + */ +export function hasTokenOverrides(posthogProperties?: Record): boolean { + return !!posthogProperties && Object.keys(posthogProperties).some((key) => TOKEN_PROPERTY_KEYS.has(key)) +} + export function getTokensSource(posthogProperties?: Record): string { - if (posthogProperties && Object.keys(posthogProperties).some((key) => TOKEN_PROPERTY_KEYS.has(key))) { - return 'passthrough' - } - return 'sdk' + return hasTokenOverrides(posthogProperties) ? 'passthrough' : 'sdk' } // limit large outputs by truncating to 200kb (approx 200k bytes) diff --git a/packages/ai/tests/gemini.test.ts b/packages/ai/tests/gemini.test.ts index b028050552..f5484e46c1 100644 --- a/packages/ai/tests/gemini.test.ts +++ b/packages/ai/tests/gemini.test.ts @@ -675,6 +675,96 @@ describe('PostHogGemini - Jest test suite', () => { ]) }) + describe('Cache reporting', () => { + test('declares inclusive cache reporting when cached tokens are present', async () => { + mockGeminiResponse = { + text: 'Cached answer', + candidates: [{ content: { parts: [{ text: 'Cached answer' }] }, finishReason: 'STOP' }], + usageMetadata: { + promptTokenCount: 23000, + candidatesTokenCount: 8, + cachedContentTokenCount: 25000, + }, + } + ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse) + + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Test', + posthogDistinctId: 'test-id', + }) + + const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0] + expect(properties['$ai_cache_read_input_tokens']).toBe(25000) + expect(properties['$ai_cache_reporting_exclusive']).toBe(false) + }) + + test('drops the cache reporting flag when the caller overrides token counts', async () => { + mockGeminiResponse = { + text: 'Cached answer', + candidates: [{ content: { parts: [{ text: 'Cached answer' }] }, finishReason: 'STOP' }], + usageMetadata: { + promptTokenCount: 23000, + candidatesTokenCount: 8, + cachedContentTokenCount: 25000, + }, + } + ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse) + + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Test', + posthogDistinctId: 'test-id', + posthogProperties: { $ai_input_tokens: 400, $ai_cache_read_input_tokens: 25000 }, + }) + + const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0] + expect(properties['$ai_tokens_source']).toBe('passthrough') + expect(properties['$ai_input_tokens']).toBe(400) + expect(properties).not.toHaveProperty('$ai_cache_reporting_exclusive') + }) + + test('keeps an explicit cache reporting flag from the caller', async () => { + mockGeminiResponse = { + text: 'Cached answer', + candidates: [{ content: { parts: [{ text: 'Cached answer' }] }, finishReason: 'STOP' }], + usageMetadata: { + promptTokenCount: 23000, + candidatesTokenCount: 8, + cachedContentTokenCount: 25000, + }, + } + ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse) + + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Test', + posthogDistinctId: 'test-id', + posthogProperties: { + $ai_input_tokens: 400, + $ai_cache_read_input_tokens: 25000, + $ai_cache_reporting_exclusive: true, + }, + }) + + const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0] + expect(properties['$ai_cache_reporting_exclusive']).toBe(true) + }) + + test('omits the cache reporting flag when no tokens were cached', async () => { + ;(client as any).client.models.generateContent = jest.fn().mockResolvedValue(mockGeminiResponse) + + await client.models.generateContent({ + model: 'gemini-2.0-flash-001', + contents: 'Test', + posthogDistinctId: 'test-id', + }) + + const { properties } = (mockPostHogClient.capture as jest.Mock).mock.calls[0][0] + expect(properties).not.toHaveProperty('$ai_cache_reporting_exclusive') + }) + }) + describe('Web Search Tracking', () => { test('should detect grounding metadata (binary detection)', async () => { mockGeminiResponse = {