diff --git a/.changeset/pretty-corners-peel.md b/.changeset/pretty-corners-peel.md new file mode 100644 index 0000000000..d59636c314 --- /dev/null +++ b/.changeset/pretty-corners-peel.md @@ -0,0 +1,5 @@ +--- +'@posthog/ai': patch +--- + +Preserve Anthropic cache creation TTL breakdowns in streaming and LangChain generation events. diff --git a/packages/ai/src/anthropic/index.ts b/packages/ai/src/anthropic/index.ts index 015029d758..3ad1030112 100644 --- a/packages/ai/src/anthropic/index.ts +++ b/packages/ai/src/anthropic/index.ts @@ -95,7 +95,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages { cacheReadInputTokens: 0, webSearchCount: 0, } - let lastRawUsage: unknown + let rawUsage: Record = {} if (Symbol.asyncIterator in value) { const [stream1, stream2] = monitoredStreamTee>( value as Stream, @@ -192,14 +192,17 @@ export class WrappedMessages extends AnthropicOriginal.Messages { } if (chunk.type == 'message_start') { - lastRawUsage = chunk.message.usage + rawUsage = { ...chunk.message.usage } usage.inputTokens = chunk.message.usage.input_tokens ?? 0 usage.cacheCreationInputTokens = chunk.message.usage.cache_creation_input_tokens ?? 0 usage.cacheReadInputTokens = chunk.message.usage.cache_read_input_tokens ?? 0 usage.webSearchCount = chunk.message.usage.server_tool_use?.web_search_requests ?? 0 } if ('usage' in chunk) { - lastRawUsage = chunk.usage + rawUsage = { + ...rawUsage, + ...Object.fromEntries(Object.entries(chunk.usage).filter(([, value]) => value != null)), + } usage.outputTokens = chunk.usage.output_tokens ?? 0 // Update web search count if present in delta if (chunk.usage.server_tool_use?.web_search_requests !== undefined) { @@ -214,7 +217,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages { } } } - usage.rawUsage = lastRawUsage + usage.rawUsage = rawUsage const latency = (Date.now() - startTime) / 1000 const timeToFirstToken = firstTokenTime !== undefined ? (firstTokenTime - startTime) / 1000 : undefined diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 69b07c8834..8dd0b31fd2 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -12,6 +12,7 @@ import { BaseMessage } from '@langchain/core/messages' import { sanitizeLangChain } from '../sanitization' import { stringifyError } from '../serializeError' import { warnIfPostHogAiGateway } from '../gatewayWarning' +import { isObject } from '../typeGuards' // Mirror LangGraph's isGraphBubbleUp guard without adding LangGraph as a dependency. Every // LangGraph control-flow exception (GraphInterrupt, NodeInterrupt, ParentCommand, GraphDrained, @@ -530,6 +531,13 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { if (additionalTokenData.cacheWriteInputTokens) { eventProperties['$ai_cache_creation_input_tokens'] = additionalTokenData.cacheWriteInputTokens } + if ( + additionalTokenData.cacheWrite5mInputTokens !== undefined && + additionalTokenData.cacheWrite1hInputTokens !== undefined + ) { + eventProperties['$ai_cache_creation_5m_input_tokens'] = additionalTokenData.cacheWrite5mInputTokens + eventProperties['$ai_cache_creation_1h_input_tokens'] = additionalTokenData.cacheWrite1hInputTokens + } if (additionalTokenData.reasoningTokens) { eventProperties['$ai_reasoning_tokens'] = additionalTokenData.reasoningTokens } @@ -702,7 +710,86 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { return stopReason != null ? String(stopReason) : undefined } - private _parseUsageModel(usage: any, provider?: string, model?: string): [number, number, Record] { + private _extractCacheCreationTtlBreakdown( + cacheCreation: unknown, + aggregateValues: unknown[] + ): [number, number] | undefined { + if (!isObject(cacheCreation)) { + return undefined + } + + const { ephemeral_5m_input_tokens: cache5m, ephemeral_1h_input_tokens: cache1h } = cacheCreation + const providedValues = [cache5m, cache1h].filter((value) => value != null) + if ( + providedValues.length === 0 || + !providedValues.every((value) => typeof value === 'number' && Number.isFinite(value) && value >= 0) + ) { + return undefined + } + + const breakdown: [number, number] = [ + typeof cache5m === 'number' ? cache5m : 0, + typeof cache1h === 'number' ? cache1h : 0, + ] + const total = breakdown[0] + breakdown[1] + const validAggregates = aggregateValues.filter( + (value): value is number => typeof value === 'number' && Number.isFinite(value) && value >= 0 + ) + return total > 0 && !validAggregates.some((aggregate) => aggregate !== total) ? breakdown : undefined + } + + private _extractBedrockCacheCreationTtlBreakdown( + cacheDetails: unknown, + aggregateValues: unknown[] + ): [number, number] | undefined { + if (!Array.isArray(cacheDetails)) { + return undefined + } + + let cache5m = 0 + let cache1h = 0 + + for (const detail of cacheDetails) { + if (!isObject(detail)) { + continue + } + + const ttl = typeof detail.ttl === 'string' ? detail.ttl.toLowerCase() : undefined + const inputTokens = detail.inputTokens + if ( + (ttl !== '5m' && ttl !== 't5m' && ttl !== '1h' && ttl !== 't1h') || + typeof inputTokens !== 'number' || + !Number.isFinite(inputTokens) || + inputTokens < 0 + ) { + continue + } + + if (ttl === '5m' || ttl === 't5m') { + cache5m += inputTokens + } else { + cache1h += inputTokens + } + } + + const total = cache5m + cache1h + const validAggregates = aggregateValues.filter( + (value): value is number => typeof value === 'number' && Number.isFinite(value) && value >= 0 + ) + if (total === 0 || validAggregates.some((aggregate) => aggregate !== total)) { + return undefined + } + + return [cache5m, cache1h] + } + + private _parseUsageModel( + usage: any, + provider?: string, + model?: string, + inputIncludesCacheTokens = true, + rawUsage?: any + ): [number, number, Record] { const conversionList: Array<[string, 'input' | 'output']> = [ ['promptTokens', 'input'], ['completionTokens', 'output'], @@ -751,6 +838,33 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { additionalTokenData.cacheWriteInputTokens = usage.input_token_details.cache_creation } + const directCacheCreationAggregates = [ + usage.cache_creation_input_tokens, + usage.input_token_details?.cache_creation, + usage.cacheWriteInputTokens, + rawUsage?.cache_creation_input_tokens, + rawUsage?.input_token_details?.cache_creation, + rawUsage?.cacheWriteInputTokens, + additionalTokenData.cacheWriteInputTokens, + ] + const cacheCreationTtl = + this._extractCacheCreationTtlBreakdown(usage.cache_creation, directCacheCreationAggregates) ?? + this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation, directCacheCreationAggregates) ?? + this._extractBedrockCacheCreationTtlBreakdown(usage.cacheDetails, [ + usage.cacheWriteInputTokens, + additionalTokenData.cacheWriteInputTokens, + ]) ?? + this._extractBedrockCacheCreationTtlBreakdown(rawUsage?.cacheDetails, [ + rawUsage?.cacheWriteInputTokens, + additionalTokenData.cacheWriteInputTokens, + ]) + if (cacheCreationTtl) { + const [cacheWrite5mInputTokens, cacheWrite1hInputTokens] = cacheCreationTtl + additionalTokenData.cacheWrite5mInputTokens = cacheWrite5mInputTokens + additionalTokenData.cacheWrite1hInputTokens = cacheWrite1hInputTokens + additionalTokenData.cacheWriteInputTokens = cacheWrite5mInputTokens + cacheWrite1hInputTokens + } + // Check for reasoning tokens in various formats if (usage.completion_tokens_details?.reasoning_tokens != null) { additionalTokenData.reasoningTokens = usage.completion_tokens_details.reasoning_tokens @@ -816,7 +930,7 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { isAnthropic = true } - if (isAnthropic && parsedUsage.input) { + if (isAnthropic && inputIncludesCacheTokens && parsedUsage.input) { const cacheTokens = (additionalTokenData.cacheReadInputTokens || 0) + (additionalTokenData.cacheWriteInputTokens || 0) if (cacheTokens > 0) { @@ -828,39 +942,86 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { } private parseUsage(response: LLMResult, provider?: string, model?: string): [number, number, Record] { - let llmUsage: [number, number, Record] = [0, 0, {}] + const isNonEmptyUsage = (usage: unknown): usage is Record => + isObject(usage) && Object.keys(usage).length > 0 + const firstNonEmptyUsage = (...candidates: unknown[]): Record | undefined => + candidates.find(isNonEmptyUsage) + + let normalizedGenerationUsage: any + let rawGenerationUsage: any + let fallbackGenerationUsage: any + + for (const generation of response.generations ?? []) { + for (const genChunk of generation) { + const generationInfo = genChunk.generationInfo ?? {} + const message = 'message' in genChunk ? genChunk.message : undefined + const messageUsage = + message && typeof message === 'object' && 'usage_metadata' in message ? message.usage_metadata : undefined + normalizedGenerationUsage = firstNonEmptyUsage( + normalizedGenerationUsage, + messageUsage, + generationInfo.usage_metadata + ) + + const messageResponseMetadata = + message && + typeof message === 'object' && + 'response_metadata' in message && + isObject(message.response_metadata) + ? message.response_metadata + : undefined + const generationResponseMetadata = isObject(generationInfo.response_metadata) + ? generationInfo.response_metadata + : undefined + const messageStreamMetadata = isObject(messageResponseMetadata?.metadata) + ? messageResponseMetadata.metadata + : undefined + const generationStreamMetadata = isObject(generationResponseMetadata?.metadata) + ? generationResponseMetadata.metadata + : undefined + rawGenerationUsage = firstNonEmptyUsage( + rawGenerationUsage, + messageResponseMetadata?.usage, + messageStreamMetadata?.usage, + generationResponseMetadata?.usage, + generationStreamMetadata?.usage + ) + fallbackGenerationUsage = firstNonEmptyUsage( + fallbackGenerationUsage, + messageResponseMetadata?.['amazon-bedrock-invocationMetrics'], + generationResponseMetadata?.['amazon-bedrock-invocationMetrics'], + generationInfo.usage_metadata + ) + } + } + + const isAnthropic = provider?.toLowerCase() === 'anthropic' || model?.toLowerCase().includes('anthropic') === true + if (isAnthropic && isNonEmptyUsage(normalizedGenerationUsage)) { + return this._parseUsageModel(normalizedGenerationUsage, provider, model, true, rawGenerationUsage) + } + const llmUsageKeys = ['token_usage', 'usage', 'tokenUsage'] if (response.llmOutput != null) { - const key = llmUsageKeys.find((k) => response.llmOutput?.[k] != null) - if (key) { - llmUsage = this._parseUsageModel(response.llmOutput[key], provider, model) + for (const key of llmUsageKeys) { + const llmUsage = response.llmOutput[key] + if (!isNonEmptyUsage(llmUsage)) { + continue + } + return this._parseUsageModel(llmUsage, provider, model, key !== 'usage', llmUsage) } } - // If top-level usage info was not found, try checking the generations. - if (llmUsage[0] === 0 && llmUsage[1] === 0 && response.generations) { - for (const generation of response.generations) { - for (const genChunk of generation) { - const message = (genChunk as any).message ?? {} - const messageResponseMetadata = message.response_metadata ?? {} - const generationInfo = genChunk.generationInfo ?? {} - const generationResponseMetadata = generationInfo.response_metadata ?? {} - const chunkUsage = - message.usage_metadata ?? - messageResponseMetadata['usage'] ?? - messageResponseMetadata['amazon-bedrock-invocationMetrics'] ?? - generationInfo.usage_metadata ?? - generationResponseMetadata['usage'] ?? - generationResponseMetadata['amazon-bedrock-invocationMetrics'] - if (chunkUsage) { - llmUsage = this._parseUsageModel(chunkUsage, provider, model) - return llmUsage - } - } - } + if (isNonEmptyUsage(normalizedGenerationUsage)) { + return this._parseUsageModel(normalizedGenerationUsage, provider, model, true, rawGenerationUsage) + } + if (isNonEmptyUsage(rawGenerationUsage)) { + return this._parseUsageModel(rawGenerationUsage, provider, model, false, rawGenerationUsage) + } + if (isNonEmptyUsage(fallbackGenerationUsage)) { + return this._parseUsageModel(fallbackGenerationUsage, provider, model) } - return llmUsage + return [0, 0, {}] } } diff --git a/packages/ai/tests/anthropic.test.ts b/packages/ai/tests/anthropic.test.ts index fa0cb0ab6b..db7413e1ac 100644 --- a/packages/ai/tests/anthropic.test.ts +++ b/packages/ai/tests/anthropic.test.ts @@ -18,6 +18,10 @@ interface MockAnthropicResponseOptions { output_tokens: number cache_creation_input_tokens?: number cache_read_input_tokens?: number + cache_creation?: { + ephemeral_5m_input_tokens: number + ephemeral_1h_input_tokens: number + } server_tool_use?: { web_search_requests?: number } @@ -160,6 +164,7 @@ const createMockStreamChunks = (options: MockAnthropicResponseOptions = {}): Moc input_tokens: options.usage?.input_tokens || 20, cache_creation_input_tokens: options.usage?.cache_creation_input_tokens || 0, cache_read_input_tokens: options.usage?.cache_read_input_tokens || 0, + ...(options.usage?.cache_creation ? { cache_creation: options.usage.cache_creation } : {}), output_tokens: 0, ...(options.usage?.server_tool_use?.web_search_requests ? { @@ -804,6 +809,51 @@ describe('PostHogAnthropic', () => { cacheReadInputTokens: 5, }) }) + + it('should retain cache creation TTL usage after the final streaming delta', async () => { + mockStreamChunks = createMockStreamChunks({ + content: 'Streaming response', + usage: { + input_tokens: 75, + output_tokens: 40, + cache_creation_input_tokens: 30, + cache_read_input_tokens: 5, + cache_creation: { + ephemeral_5m_input_tokens: 10, + ephemeral_1h_input_tokens: 20, + }, + }, + }) + + const stream = await client.messages.create({ + model: 'claude-3-opus-20240229', + messages: [{ role: 'user', content: 'Hello' }], + max_tokens: 100, + stream: true, + posthogDistinctId: 'test-user-123', + }) + + for await (const _chunk of stream) { + // Consume the stream so the monitoring branch reaches message_delta. + } + await waitForAsyncCapture() + + const captureMock = mockPostHogClient.capture as jest.Mock + const [captureArgs] = captureMock.mock.calls + + expect(captureArgs[0].properties['$ai_usage']).toEqual({ + input_tokens: 75, + output_tokens: 40, + cache_creation_input_tokens: 30, + cache_read_input_tokens: 5, + cache_creation: { + ephemeral_5m_input_tokens: 10, + ephemeral_1h_input_tokens: 20, + }, + }) + expect(captureArgs[0].properties['$ai_cache_creation_input_tokens']).toBe(30) + expect(captureArgs[0].properties['$ai_output_tokens']).toBe(40) + }) }) describe('Telemetry failure isolation', () => { diff --git a/packages/ai/tests/callbacks.test.ts b/packages/ai/tests/callbacks.test.ts index 220c591425..9502122d85 100644 --- a/packages/ai/tests/callbacks.test.ts +++ b/packages/ai/tests/callbacks.test.ts @@ -1,6 +1,7 @@ import { LangChainCallbackHandler } from '../src/langchain/callbacks' import { PostHog } from 'posthog-node' import { AIMessage } from '@langchain/core/messages' +import type { ChatGeneration } from '@langchain/core/outputs' import { version } from '../package.json' const mockPostHogClient = { @@ -580,6 +581,174 @@ describe('LangChainCallbackHandler', () => { expect(captureCall[0].properties['$ai_cache_read_input_tokens']).toBe(50) }) + it('should use generation info usage when message response metadata is empty', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'llms', 'openai', 'OpenAI'], + kwargs: {}, + } + const runId = 'run_generation_info_usage_test' + handler.handleLLMStart(serialized, ['Use generation info'], runId, undefined, {}, undefined, { + ls_model_name: 'gpt-4', + ls_provider: 'openai', + }) + + const generation = { + text: 'Response with generation info usage.', + message: new AIMessage({ + content: 'Response with generation info usage.', + response_metadata: {}, + }), + generationInfo: { + response_metadata: { + usage: { + input_tokens: 21, + output_tokens: 9, + }, + }, + }, + } satisfies ChatGeneration + + handler.handleLLMEnd( + { + generations: [[generation]], + }, + runId + ) + + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureCall[0].properties['$ai_input_tokens']).toBe(21) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(9) + }) + + it('should ignore empty top-level usage and fall back to generation metadata', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'llms', 'openai', 'OpenAI'], + kwargs: {}, + } + const runId = 'run_empty_top_level_usage_test' + handler.handleLLMStart(serialized, ['Use generation info'], runId, undefined, {}, undefined, { + ls_model_name: 'gpt-4', + ls_provider: 'openai', + }) + + const generation = { + text: 'Response with generation info usage.', + message: new AIMessage({ + content: 'Response with generation info usage.', + response_metadata: {}, + }), + generationInfo: { + response_metadata: { + usage: { + input_tokens: 21, + output_tokens: 9, + }, + }, + }, + } satisfies ChatGeneration + + handler.handleLLMEnd( + { + generations: [[generation]], + llmOutput: { + tokenUsage: {}, + }, + }, + runId + ) + + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureCall[0].properties['$ai_input_tokens']).toBe(21) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(9) + }) + + it('should ignore empty Anthropic generation usage and fall back to raw generation usage', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'anthropic', 'ChatAnthropic'], + kwargs: {}, + } + const runId = 'run_empty_anthropic_generation_usage_test' + handler.handleLLMStart(serialized, ['Use raw generation usage'], runId, undefined, {}, undefined, { + ls_model_name: 'claude-sonnet-4-6', + ls_provider: 'anthropic', + }) + + const generation = { + text: 'Response with raw Anthropic usage.', + message: new AIMessage({ + content: 'Response with raw Anthropic usage.', + usage_metadata: {} as NonNullable, + response_metadata: { + usage: { + input_tokens: 21, + output_tokens: 9, + }, + }, + }), + } satisfies ChatGeneration + + handler.handleLLMEnd( + { + generations: [[generation]], + }, + runId + ) + + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureCall[0].properties['$ai_input_tokens']).toBe(21) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(9) + }) + + it('should prefer top-level usage for non-Anthropic providers', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'llms', 'openai', 'OpenAI'], + kwargs: {}, + } + const runId = 'run_top_level_usage_precedence_test' + handler.handleLLMStart(serialized, ['Use top-level usage'], runId, undefined, {}, undefined, { + ls_model_name: 'gpt-4', + ls_provider: 'openai', + }) + + const generation = { + text: 'Response with differing usage sources.', + message: new AIMessage({ + content: 'Response with differing usage sources.', + usage_metadata: { + input_tokens: 999, + output_tokens: 888, + total_tokens: 1887, + }, + }), + } satisfies ChatGeneration + + handler.handleLLMEnd( + { + generations: [[generation]], + llmOutput: { + tokenUsage: { + promptTokens: 100, + completionTokens: 20, + totalTokens: 120, + }, + }, + }, + runId + ) + + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureCall[0].properties['$ai_input_tokens']).toBe(100) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(20) + }) + it('should subtract cache_read_tokens from input_tokens for Anthropic provider', async () => { const serialized = { lc: 1, @@ -790,6 +959,8 @@ describe('LangChainCallbackHandler', () => { expect(captureCall[0].properties['$ai_input_tokens']).toBe(200) expect(captureCall[0].properties['$ai_output_tokens']).toBe(50) expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(800) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBeUndefined() + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBeUndefined() }) it('should subtract both cache_read and cache_creation tokens for Anthropic', async () => { @@ -848,6 +1019,608 @@ describe('LangChainCallbackHandler', () => { expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(500) }) + it('should preserve Anthropic cache creation TTLs from LangChain response metadata', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'anthropic', 'ChatAnthropic'], + kwargs: {}, + } + + const runId = 'run_anthropic_cache_ttl_test' + const metadata = { ls_model_name: 'claude-sonnet-4-6', ls_provider: 'anthropic' } + handler.handleLLMStart(serialized, ['Use the cached context'], runId, undefined, {}, undefined, metadata) + + const llmResult = { + generations: [ + [ + { + text: 'Response from Anthropic with cache creation TTLs.', + message: new AIMessage({ + content: 'Response from Anthropic with cache creation TTLs.', + response_metadata: { + usage: { + input_tokens: 18, + output_tokens: 50, + cache_creation_input_tokens: 300, + cache_read_input_tokens: 0, + cache_creation: { + ephemeral_5m_input_tokens: 100, + ephemeral_1h_input_tokens: 200, + }, + }, + }, + usage_metadata: { + input_tokens: 318, + output_tokens: 50, + total_tokens: 368, + input_token_details: { + cache_creation: 300, + cache_read: 0, + }, + }, + }), + }, + ], + ], + } + + handler.handleLLMEnd(llmResult, runId) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + + expect(captureCall[0].properties['$ai_input_tokens']).toBe(18) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(50) + expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(300) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBe(100) + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBe(200) + }) + + it('should keep the Anthropic aggregate when the direct TTL breakdown does not match', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'anthropic', 'ChatAnthropic'], + kwargs: {}, + } + + const runId = 'run_anthropic_cache_ttl_mismatch_test' + const metadata = { ls_model_name: 'claude-sonnet-4-6', ls_provider: 'anthropic' } + handler.handleLLMStart(serialized, ['Use the cached context'], runId, undefined, {}, undefined, metadata) + + const llmResult = { + generations: [ + [ + { + text: 'Response from Anthropic with mismatched cache creation TTLs.', + message: new AIMessage({ + content: 'Response from Anthropic with mismatched cache creation TTLs.', + response_metadata: { + usage: { + input_tokens: 18, + output_tokens: 50, + cache_creation_input_tokens: 300, + cache_read_input_tokens: 0, + cache_creation: { + ephemeral_5m_input_tokens: 100, + ephemeral_1h_input_tokens: 100, + }, + }, + }, + usage_metadata: { + input_tokens: 318, + output_tokens: 50, + total_tokens: 368, + input_token_details: { + cache_creation: 300, + cache_read: 0, + }, + }, + }), + }, + ], + ], + } + + handler.handleLLMEnd(llmResult, runId) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + + expect(captureCall[0].properties['$ai_input_tokens']).toBe(18) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(50) + expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(300) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBeUndefined() + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBeUndefined() + }) + + it('should not subtract cache creation twice from raw Anthropic usage', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'anthropic', 'ChatAnthropic'], + kwargs: {}, + } + + const runId = 'run_anthropic_raw_cache_ttl_test' + const metadata = { ls_model_name: 'claude-sonnet-4-6', ls_provider: 'anthropic' } + handler.handleLLMStart(serialized, ['Use the cached context'], runId, undefined, {}, undefined, metadata) + + const rawUsage = { + input_tokens: 18, + output_tokens: 50, + cache_creation_input_tokens: 300, + cache_read_input_tokens: 0, + cache_creation: { + ephemeral_5m_input_tokens: 100, + ephemeral_1h_input_tokens: 200, + }, + } + const llmResult = { + generations: [ + [ + { + text: 'Response from Anthropic with cache creation TTLs.', + message: new AIMessage({ + content: 'Response from Anthropic with cache creation TTLs.', + response_metadata: { usage: rawUsage }, + usage_metadata: { + input_tokens: 318, + output_tokens: 50, + total_tokens: 368, + input_token_details: { + cache_creation: 300, + cache_read: 0, + }, + }, + }), + }, + ], + ], + llmOutput: { usage: rawUsage }, + } + + handler.handleLLMEnd(llmResult, runId) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + + expect(captureCall[0].properties['$ai_input_tokens']).toBe(18) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(50) + expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(300) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBe(100) + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBe(200) + }) + + it('should preserve Anthropic cache creation TTLs from Bedrock Converse response metadata', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'bedrock', 'ChatBedrockConverse'], + kwargs: {}, + } + + const runId = 'run_bedrock_converse_cache_ttl_test' + const metadata = { + ls_model_name: 'us.anthropic.claude-sonnet-4-6-v1:0', + ls_provider: 'amazon_bedrock', + } + handler.handleLLMStart(serialized, ['Use the cached context'], runId, undefined, {}, undefined, metadata) + + const llmResult = { + generations: [ + [ + { + text: 'Response from Bedrock Converse with cache creation TTLs.', + message: new AIMessage({ + content: 'Response from Bedrock Converse with cache creation TTLs.', + response_metadata: { + usage: { + inputTokens: 18, + outputTokens: 50, + cacheWriteInputTokens: 300, + cacheReadInputTokens: 0, + cacheDetails: [ + { ttl: '5m', inputTokens: 100 }, + { ttl: '1h', inputTokens: 200 }, + ], + }, + }, + usage_metadata: { + input_tokens: 318, + output_tokens: 50, + total_tokens: 368, + input_token_details: { + cache_creation: 300, + cache_read: 0, + }, + }, + }), + }, + ], + ], + } + + handler.handleLLMEnd(llmResult, runId) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + + expect(captureCall[0].properties['$ai_input_tokens']).toBe(18) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(50) + expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(300) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBe(100) + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBe(200) + }) + + it('should preserve Anthropic cache creation TTLs from aggregated Bedrock Converse stream metadata', async () => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'bedrock', 'ChatBedrockConverse'], + kwargs: {}, + } + + const runId = 'run_bedrock_converse_stream_cache_ttl_test' + const metadata = { + ls_model_name: 'us.anthropic.claude-sonnet-4-6-v1:0', + ls_provider: 'amazon_bedrock', + } + handler.handleLLMStart(serialized, ['Use the cached context'], runId, undefined, {}, undefined, metadata) + + const llmResult = { + generations: [ + [ + { + text: 'Streamed response from Bedrock Converse with cache creation TTLs.', + message: new AIMessage({ + content: 'Streamed response from Bedrock Converse with cache creation TTLs.', + response_metadata: { + metadata: { + usage: { + inputTokens: 18, + outputTokens: 50, + cacheWriteInputTokens: 300, + cacheReadInputTokens: 0, + cacheDetails: [ + { ttl: '5m', inputTokens: 100 }, + { ttl: '1h', inputTokens: 200 }, + ], + }, + }, + }, + usage_metadata: { + input_tokens: 318, + output_tokens: 50, + total_tokens: 368, + input_token_details: { + cache_creation: 300, + cache_read: 0, + }, + }, + }), + }, + ], + ], + } + + handler.handleLLMEnd(llmResult, runId) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + + expect(captureCall[0].properties['$ai_input_tokens']).toBe(18) + expect(captureCall[0].properties['$ai_output_tokens']).toBe(50) + expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(300) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBe(100) + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBe(200) + }) + + describe('Bedrock usage source selection', () => { + const captureBedrockUsage = (generations: ChatGeneration[][], runId: string): Record => { + handler.handleLLMStart( + { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'bedrock', 'ChatBedrockConverse'], + kwargs: {}, + }, + ['Test Bedrock usage'], + runId, + undefined, + {}, + undefined, + { + ls_model_name: 'us.anthropic.claude-sonnet-4-6-v1:0', + ls_provider: 'amazon_bedrock', + } + ) + + handler.handleLLMEnd({ generations }, runId) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + return (mockPostHogClient.capture as jest.Mock).mock.calls[0][0].properties + } + + it('falls back to Bedrock invocation metrics when raw usage is empty', () => { + const properties = captureBedrockUsage( + [ + [ + { + text: 'Response from Bedrock.', + message: new AIMessage({ + content: 'Response from Bedrock.', + response_metadata: { + usage: {}, + 'amazon-bedrock-invocationMetrics': { + inputTokenCount: 21, + outputTokenCount: 9, + }, + }, + }), + }, + ], + ], + 'run_bedrock_empty_raw_usage' + ) + + expect(properties['$ai_input_tokens']).toBe(21) + expect(properties['$ai_output_tokens']).toBe(9) + }) + + it('uses later valid metadata when an earlier source in the same generation is empty', () => { + const properties = captureBedrockUsage( + [ + [ + { + text: 'Response from Bedrock.', + message: new AIMessage({ + content: 'Response from Bedrock.', + response_metadata: { + usage: {}, + metadata: { + usage: { + inputTokenCount: 34, + outputTokenCount: 12, + }, + }, + }, + }), + }, + ], + ], + 'run_bedrock_later_metadata_usage' + ) + + expect(properties['$ai_input_tokens']).toBe(34) + expect(properties['$ai_output_tokens']).toBe(12) + }) + + it('uses generation metadata when message usage metadata is empty', () => { + const properties = captureBedrockUsage( + [ + [ + { + text: 'Response from Bedrock.', + generationInfo: { + usage_metadata: { + input_tokens: 321, + output_tokens: 8, + total_tokens: 329, + input_token_details: { + cache_creation: 300, + cache_read: 0, + }, + }, + response_metadata: { + usage: { + cache_creation_input_tokens: 300, + cache_creation: { + ephemeral_5m_input_tokens: 100, + ephemeral_1h_input_tokens: 200, + }, + }, + }, + }, + message: new AIMessage({ + content: 'Response from Bedrock.', + usage_metadata: {} as NonNullable, + }), + }, + ], + ], + 'run_bedrock_generation_usage_metadata' + ) + + expect(properties['$ai_input_tokens']).toBe(21) + expect(properties['$ai_output_tokens']).toBe(8) + expect(properties['$ai_cache_creation_input_tokens']).toBe(300) + expect(properties['$ai_cache_creation_5m_input_tokens']).toBe(100) + expect(properties['$ai_cache_creation_1h_input_tokens']).toBe(200) + }) + + it('uses valid usage from a later generation when an earlier generation is empty', () => { + const properties = captureBedrockUsage( + [ + [ + { + text: 'Partial response from Bedrock.', + message: new AIMessage({ + content: 'Partial response from Bedrock.', + response_metadata: { usage: {} }, + }), + }, + ], + [ + { + text: 'Final response from Bedrock.', + message: new AIMessage({ + content: 'Final response from Bedrock.', + response_metadata: { + usage: { + inputTokenCount: 55, + outputTokenCount: 17, + }, + }, + }), + }, + ], + ], + 'run_bedrock_later_generation_usage' + ) + + expect(properties['$ai_input_tokens']).toBe(55) + expect(properties['$ai_output_tokens']).toBe(17) + }) + + it('uses valid Bedrock invocation metrics after an empty earlier fallback', () => { + const properties = captureBedrockUsage( + [ + [ + { + text: 'Partial response from Bedrock.', + message: new AIMessage({ + content: 'Partial response from Bedrock.', + response_metadata: { + 'amazon-bedrock-invocationMetrics': {}, + }, + }), + }, + ], + [ + { + text: 'Final response from Bedrock.', + message: new AIMessage({ + content: 'Final response from Bedrock.', + response_metadata: { + 'amazon-bedrock-invocationMetrics': { + inputTokenCount: 89, + outputTokenCount: 23, + }, + }, + }), + }, + ], + ], + 'run_bedrock_later_invocation_metrics' + ) + + expect(properties['$ai_input_tokens']).toBe(89) + expect(properties['$ai_output_tokens']).toBe(23) + }) + + it('keeps explicit zero-valued usage ahead of lower-priority fallback metrics', () => { + const properties = captureBedrockUsage( + [ + [ + { + text: 'Response from Bedrock.', + message: new AIMessage({ + content: 'Response from Bedrock.', + response_metadata: { + usage: { + inputTokenCount: 0, + outputTokenCount: 0, + }, + 'amazon-bedrock-invocationMetrics': { + inputTokenCount: 21, + outputTokenCount: 9, + }, + }, + }), + }, + ], + ], + 'run_bedrock_explicit_zero_usage' + ) + + expect(properties['$ai_input_tokens']).toBe(0) + expect(properties['$ai_output_tokens']).toBe(0) + }) + }) + + it.each([ + { + name: 'sums repeated buckets and accepts AWS SDK enum values', + aggregate: 300, + cacheDetails: [ + { ttl: 'T5M', inputTokens: 40 }, + { ttl: 't5m', inputTokens: 60 }, + { ttl: 'T1H', inputTokens: 200 }, + ], + expected5m: 100, + expected1h: 200, + }, + { + name: 'fills the missing bucket when one TTL is reported', + aggregate: 200, + cacheDetails: [{ ttl: '1h', inputTokens: 200 }], + expected5m: 0, + expected1h: 200, + }, + { + name: 'keeps aggregate-only fallback when valid details do not match it', + aggregate: 300, + cacheDetails: [ + { ttl: '5m', inputTokens: 100 }, + { ttl: '1h', inputTokens: 100 }, + { ttl: 'unknown', inputTokens: 100 }, + { ttl: '5m', inputTokens: -1 }, + { ttl: '1h', inputTokens: Number.NaN }, + ], + expected5m: undefined, + expected1h: undefined, + }, + ])('$name', async ({ aggregate, cacheDetails, expected5m, expected1h }) => { + const serialized = { + lc: 1, + type: 'constructor' as const, + id: ['langchain', 'chat_models', 'bedrock', 'ChatBedrockConverse'], + kwargs: {}, + } + const runId = `run_bedrock_cache_details_${aggregate}_${expected5m ?? 'fallback'}` + handler.handleLLMStart(serialized, ['Use the cached context'], runId, undefined, {}, undefined, { + ls_model_name: 'us.anthropic.claude-sonnet-4-6-v1:0', + ls_provider: 'amazon_bedrock', + }) + + const generation = { + text: 'Response from Bedrock Converse.', + message: new AIMessage({ + content: 'Response from Bedrock Converse.', + response_metadata: { + usage: { + inputTokens: 18, + outputTokens: 50, + cacheWriteInputTokens: aggregate, + cacheDetails, + }, + }, + usage_metadata: { + input_tokens: 18 + aggregate, + output_tokens: 50, + total_tokens: 68 + aggregate, + input_token_details: { cache_creation: aggregate }, + }, + }), + } satisfies ChatGeneration + + handler.handleLLMEnd( + { + generations: [[generation]], + }, + runId + ) + + expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1) + const [captureCall] = (mockPostHogClient.capture as jest.Mock).mock.calls + expect(captureCall[0].properties['$ai_input_tokens']).toBe(18) + expect(captureCall[0].properties['$ai_cache_creation_input_tokens']).toBe(aggregate) + expect(captureCall[0].properties['$ai_cache_creation_5m_input_tokens']).toBe(expected5m) + expect(captureCall[0].properties['$ai_cache_creation_1h_input_tokens']).toBe(expected1h) + }) + it('should not subtract cache_creation_input_tokens for non-Anthropic providers', async () => { const serialized = { lc: 1,