From 9c81e6840debd56716c67f57d807bec32a7c76fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Silva=20Gouv=C3=AAa?= Date: Thu, 23 Jul 2026 10:33:45 -0300 Subject: [PATCH 1/7] fix(ai): preserve Anthropic cache-write TTL breakdowns --- .changeset/pretty-corners-peel.md | 5 + packages/ai/src/anthropic/index.ts | 11 +- packages/ai/src/langchain/callbacks.ts | 119 +++++++++++---- packages/ai/tests/anthropic.test.ts | 50 ++++++ packages/ai/tests/callbacks.test.ts | 204 +++++++++++++++++++++++++ 5 files changed, 358 insertions(+), 31 deletions(-) create mode 100644 .changeset/pretty-corners-peel.md 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 3cf8b2c858..cf3d2db970 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 ('tee' in value) { const [stream1, stream2] = value.tee() ;(async () => { @@ -189,14 +189,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) { @@ -211,7 +214,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 a248580d40..72dc72d47c 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -13,6 +13,9 @@ import { sanitizeLangChain } from '../sanitization' import { stringifyError } from '../serializeError' import { warnIfPostHogAiGateway } from '../gatewayWarning' +const isRecord = (value: unknown): value is Record => + value !== null && typeof value === 'object' && !Array.isArray(value) + interface SpanMetadata { /** Name of the trace/span (e.g. chain name) */ name: string @@ -500,6 +503,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 } @@ -676,7 +686,34 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { return undefined } - private _parseUsageModel(usage: any, provider?: string, model?: string): [number, number, Record] { + private _extractCacheCreationTtlBreakdown(cacheCreation: unknown): [number, number] | undefined { + if (!isRecord(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, + ] + return breakdown[0] + breakdown[1] > 0 ? breakdown : undefined + } + + 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'], @@ -725,6 +762,16 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { additionalTokenData.cacheWriteInputTokens = usage.input_token_details.cache_creation } + const cacheCreationTtl = + this._extractCacheCreationTtlBreakdown(usage.cache_creation) ?? + this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation) + 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 @@ -790,7 +837,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) { @@ -802,40 +849,58 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { } private parseUsage(response: LLMResult, provider?: string, model?: string): [number, number, Record] { - let llmUsage: [number, number, Record] = [0, 0, {}] + 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 ??= messageUsage ?? generationInfo.usage_metadata + + const messageResponseMetadata = + message && + typeof message === 'object' && + 'response_metadata' in message && + isRecord(message.response_metadata) + ? message.response_metadata + : undefined + const generationResponseMetadata = isRecord(generationInfo.response_metadata) + ? generationInfo.response_metadata + : undefined + rawGenerationUsage ??= messageResponseMetadata?.usage ?? generationResponseMetadata?.usage + fallbackGenerationUsage ??= + messageResponseMetadata?.['amazon-bedrock-invocationMetrics'] ?? + generationResponseMetadata?.['amazon-bedrock-invocationMetrics'] ?? + generationInfo.usage_metadata + } + } + + const isAnthropic = provider?.toLowerCase() === 'anthropic' || model?.toLowerCase().includes('anthropic') === true + if (isAnthropic && 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) + const llmUsage = response.llmOutput[key] + 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) { - // Check other paths for usage information - if (genChunk.generationInfo?.usage_metadata) { - llmUsage = this._parseUsageModel(genChunk.generationInfo.usage_metadata, provider, model) - return llmUsage - } - - const messageChunk = genChunk.generationInfo ?? {} - const responseMetadata = messageChunk.response_metadata ?? {} - const chunkUsage = - responseMetadata['usage'] ?? - responseMetadata['amazon-bedrock-invocationMetrics'] ?? - messageChunk.usage_metadata - if (chunkUsage) { - llmUsage = this._parseUsageModel(chunkUsage, provider, model) - return llmUsage - } - } - } + if (rawGenerationUsage) { + return this._parseUsageModel(rawGenerationUsage, provider, model, false, rawGenerationUsage) + } + if (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 ddfd270618..764d25c816 100644 --- a/packages/ai/tests/anthropic.test.ts +++ b/packages/ai/tests/anthropic.test.ts @@ -17,6 +17,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 } @@ -159,6 +163,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 ? { @@ -808,6 +813,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('Error Handling', () => { diff --git a/packages/ai/tests/callbacks.test.ts b/packages/ai/tests/callbacks.test.ts index a30d691d38..9cdadfb25e 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 = { @@ -365,6 +366,91 @@ 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 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, @@ -575,6 +661,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 () => { @@ -633,6 +721,122 @@ 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 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 not subtract cache_creation_input_tokens for non-Anthropic providers', async () => { const serialized = { lc: 1, From 193ca9126c61821784fe9655ada63c4c2e163e5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Silva=20Gouv=C3=AAa?= Date: Thu, 23 Jul 2026 10:56:42 -0300 Subject: [PATCH 2/7] fix(ai): preserve Bedrock cache TTL breakdowns --- packages/ai/src/langchain/callbacks.ts | 67 +++++++- packages/ai/tests/callbacks.test.ts | 205 +++++++++++++++++++++++++ 2 files changed, 270 insertions(+), 2 deletions(-) diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 72dc72d47c..274b8879d7 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -707,6 +707,51 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { return breakdown[0] + breakdown[1] > 0 ? 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 (!isRecord(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, @@ -764,7 +809,15 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { const cacheCreationTtl = this._extractCacheCreationTtlBreakdown(usage.cache_creation) ?? - this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation) + this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation) ?? + 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 @@ -871,7 +924,17 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { const generationResponseMetadata = isRecord(generationInfo.response_metadata) ? generationInfo.response_metadata : undefined - rawGenerationUsage ??= messageResponseMetadata?.usage ?? generationResponseMetadata?.usage + const messageStreamMetadata = isRecord(messageResponseMetadata?.metadata) + ? messageResponseMetadata.metadata + : undefined + const generationStreamMetadata = isRecord(generationResponseMetadata?.metadata) + ? generationResponseMetadata.metadata + : undefined + rawGenerationUsage ??= + messageResponseMetadata?.usage ?? + messageStreamMetadata?.usage ?? + generationResponseMetadata?.usage ?? + generationStreamMetadata?.usage fallbackGenerationUsage ??= messageResponseMetadata?.['amazon-bedrock-invocationMetrics'] ?? generationResponseMetadata?.['amazon-bedrock-invocationMetrics'] ?? diff --git a/packages/ai/tests/callbacks.test.ts b/packages/ai/tests/callbacks.test.ts index 9cdadfb25e..b41fe78cc4 100644 --- a/packages/ai/tests/callbacks.test.ts +++ b/packages/ai/tests/callbacks.test.ts @@ -837,6 +837,211 @@ describe('LangChainCallbackHandler', () => { 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) + }) + + 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, From 5cecc5978efd080591d685770fe814971d38e2d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Gouv=C3=AAa?= <64869705+gouveags@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:09:10 -0300 Subject: [PATCH 3/7] fix(ai): address LangChain usage fallbacks --- packages/ai/src/langchain/callbacks.ts | 46 +++++++---- packages/ai/tests/callbacks.test.ts | 102 +++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 15 deletions(-) diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 274b8879d7..47ddbbe403 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -12,9 +12,7 @@ import { BaseMessage } from '@langchain/core/messages' import { sanitizeLangChain } from '../sanitization' import { stringifyError } from '../serializeError' import { warnIfPostHogAiGateway } from '../gatewayWarning' - -const isRecord = (value: unknown): value is Record => - value !== null && typeof value === 'object' && !Array.isArray(value) +import { isObject } from '../typeGuards' interface SpanMetadata { /** Name of the trace/span (e.g. chain name) */ @@ -686,8 +684,11 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { return undefined } - private _extractCacheCreationTtlBreakdown(cacheCreation: unknown): [number, number] | undefined { - if (!isRecord(cacheCreation)) { + private _extractCacheCreationTtlBreakdown( + cacheCreation: unknown, + aggregateValues: unknown[] + ): [number, number] | undefined { + if (!isObject(cacheCreation)) { return undefined } @@ -704,7 +705,11 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { typeof cache5m === 'number' ? cache5m : 0, typeof cache1h === 'number' ? cache1h : 0, ] - return breakdown[0] + breakdown[1] > 0 ? breakdown : undefined + 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( @@ -719,7 +724,7 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { let cache1h = 0 for (const detail of cacheDetails) { - if (!isRecord(detail)) { + if (!isObject(detail)) { continue } @@ -807,9 +812,18 @@ 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) ?? - this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation) ?? + this._extractCacheCreationTtlBreakdown(usage.cache_creation, directCacheCreationAggregates) ?? + this._extractCacheCreationTtlBreakdown(rawUsage?.cache_creation, directCacheCreationAggregates) ?? this._extractBedrockCacheCreationTtlBreakdown(usage.cacheDetails, [ usage.cacheWriteInputTokens, additionalTokenData.cacheWriteInputTokens, @@ -918,16 +932,16 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { message && typeof message === 'object' && 'response_metadata' in message && - isRecord(message.response_metadata) + isObject(message.response_metadata) ? message.response_metadata : undefined - const generationResponseMetadata = isRecord(generationInfo.response_metadata) + const generationResponseMetadata = isObject(generationInfo.response_metadata) ? generationInfo.response_metadata : undefined - const messageStreamMetadata = isRecord(messageResponseMetadata?.metadata) + const messageStreamMetadata = isObject(messageResponseMetadata?.metadata) ? messageResponseMetadata.metadata : undefined - const generationStreamMetadata = isRecord(generationResponseMetadata?.metadata) + const generationStreamMetadata = isObject(generationResponseMetadata?.metadata) ? generationResponseMetadata.metadata : undefined rawGenerationUsage ??= @@ -950,9 +964,11 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { const llmUsageKeys = ['token_usage', 'usage', 'tokenUsage'] if (response.llmOutput != null) { - const key = llmUsageKeys.find((k) => response.llmOutput?.[k] != null) - if (key) { + for (const key of llmUsageKeys) { const llmUsage = response.llmOutput[key] + if (!isObject(llmUsage) || Object.keys(llmUsage).length === 0) { + continue + } return this._parseUsageModel(llmUsage, provider, model, key !== 'usage', llmUsage) } } diff --git a/packages/ai/tests/callbacks.test.ts b/packages/ai/tests/callbacks.test.ts index b41fe78cc4..06c59c17be 100644 --- a/packages/ai/tests/callbacks.test.ts +++ b/packages/ai/tests/callbacks.test.ts @@ -407,6 +407,50 @@ describe('LangChainCallbackHandler', () => { 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 prefer top-level usage for non-Anthropic providers', async () => { const serialized = { lc: 1, @@ -779,6 +823,64 @@ describe('LangChainCallbackHandler', () => { 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, From 572f7c27dd6433579d7c7a81efc7cdcbc467ae66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Gouv=C3=AAa?= <64869705+gouveags@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:23:09 -0300 Subject: [PATCH 4/7] fix(ai): handle empty Anthropic usage metadata --- packages/ai/src/langchain/callbacks.ts | 2 +- packages/ai/tests/callbacks.test.ts | 39 ++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 47ddbbe403..892206689d 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -957,7 +957,7 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { } const isAnthropic = provider?.toLowerCase() === 'anthropic' || model?.toLowerCase().includes('anthropic') === true - if (isAnthropic && normalizedGenerationUsage) { + if (isAnthropic && isObject(normalizedGenerationUsage) && Object.keys(normalizedGenerationUsage).length > 0) { return this._parseUsageModel(normalizedGenerationUsage, provider, model, true, rawGenerationUsage) } diff --git a/packages/ai/tests/callbacks.test.ts b/packages/ai/tests/callbacks.test.ts index 06c59c17be..cb1a90e672 100644 --- a/packages/ai/tests/callbacks.test.ts +++ b/packages/ai/tests/callbacks.test.ts @@ -451,6 +451,45 @@ describe('LangChainCallbackHandler', () => { 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, From 388ef5f2e89ea8e7539bbde9c2944c278bceb71d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Gouv=C3=AAa?= <64869705+gouveags@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:37:48 -0300 Subject: [PATCH 5/7] fix(ai): skip empty LangChain usage sources --- packages/ai/src/langchain/callbacks.ts | 37 ++-- packages/ai/tests/callbacks.test.ts | 223 +++++++++++++++++++++++++ 2 files changed, 248 insertions(+), 12 deletions(-) diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 892206689d..8bead0d47a 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -916,6 +916,11 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { } private parseUsage(response: LLMResult, provider?: string, model?: string): [number, number, Record] { + 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 @@ -926,7 +931,11 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { const message = 'message' in genChunk ? genChunk.message : undefined const messageUsage = message && typeof message === 'object' && 'usage_metadata' in message ? message.usage_metadata : undefined - normalizedGenerationUsage ??= messageUsage ?? generationInfo.usage_metadata + normalizedGenerationUsage = firstNonEmptyUsage( + normalizedGenerationUsage, + messageUsage, + generationInfo.usage_metadata + ) const messageResponseMetadata = message && @@ -944,20 +953,24 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { const generationStreamMetadata = isObject(generationResponseMetadata?.metadata) ? generationResponseMetadata.metadata : undefined - rawGenerationUsage ??= - messageResponseMetadata?.usage ?? - messageStreamMetadata?.usage ?? - generationResponseMetadata?.usage ?? + rawGenerationUsage = firstNonEmptyUsage( + rawGenerationUsage, + messageResponseMetadata?.usage, + messageStreamMetadata?.usage, + generationResponseMetadata?.usage, generationStreamMetadata?.usage - fallbackGenerationUsage ??= - messageResponseMetadata?.['amazon-bedrock-invocationMetrics'] ?? - generationResponseMetadata?.['amazon-bedrock-invocationMetrics'] ?? + ) + 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 && isObject(normalizedGenerationUsage) && Object.keys(normalizedGenerationUsage).length > 0) { + if (isAnthropic && isNonEmptyUsage(normalizedGenerationUsage)) { return this._parseUsageModel(normalizedGenerationUsage, provider, model, true, rawGenerationUsage) } @@ -966,17 +979,17 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { if (response.llmOutput != null) { for (const key of llmUsageKeys) { const llmUsage = response.llmOutput[key] - if (!isObject(llmUsage) || Object.keys(llmUsage).length === 0) { + if (!isNonEmptyUsage(llmUsage)) { continue } return this._parseUsageModel(llmUsage, provider, model, key !== 'usage', llmUsage) } } - if (rawGenerationUsage) { + if (isNonEmptyUsage(rawGenerationUsage)) { return this._parseUsageModel(rawGenerationUsage, provider, model, false, rawGenerationUsage) } - if (fallbackGenerationUsage) { + if (isNonEmptyUsage(fallbackGenerationUsage)) { return this._parseUsageModel(fallbackGenerationUsage, provider, model) } diff --git a/packages/ai/tests/callbacks.test.ts b/packages/ai/tests/callbacks.test.ts index cb1a90e672..5f83759496 100644 --- a/packages/ai/tests/callbacks.test.ts +++ b/packages/ai/tests/callbacks.test.ts @@ -1102,6 +1102,229 @@ describe('LangChainCallbackHandler', () => { 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', From 6d20779e39435702b096ea210b5718665859c700 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Gouv=C3=AAa?= <64869705+gouveags@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:21:55 -0300 Subject: [PATCH 6/7] chore(ai): align conflicted files with current main Align the two conflicted AI files before updating the branch, so the main sync can be created through GitHub with a verified signature. The final TTL resolution follows in the next signed commit. --- packages/ai/src/anthropic/index.ts | 41 ++-- packages/ai/src/langchain/callbacks.ts | 285 +++++++------------------ 2 files changed, 95 insertions(+), 231 deletions(-) diff --git a/packages/ai/src/anthropic/index.ts b/packages/ai/src/anthropic/index.ts index cf3d2db970..015029d758 100644 --- a/packages/ai/src/anthropic/index.ts +++ b/packages/ai/src/anthropic/index.ts @@ -1,4 +1,4 @@ -import AnthropicOriginal, { APIPromise } from '@anthropic-ai/sdk' +import AnthropicOriginal, { APIPromise, ClientOptions } from '@anthropic-ai/sdk' import { PostHog } from 'posthog-node' import { formatResponseAnthropic, @@ -18,18 +18,19 @@ type Message = AnthropicOriginal.Messages.Message type RawMessageStreamEvent = AnthropicOriginal.Messages.RawMessageStreamEvent type MessageCreateParamsBase = AnthropicOriginal.Messages.MessageCreateParams type RequestOptions = AnthropicOriginal.RequestOptions -import type { Stream } from '@anthropic-ai/sdk/streaming' +import { Stream } from '@anthropic-ai/sdk/streaming' import { sanitizeAnthropic } from '../sanitization' +import { preserveProviderPromise } from '../providerPromise' +import { monitoredStreamTee } from '../stream' interface ToolInProgress { block: FormattedFunctionCall inputString: string } -interface MonitoringAnthropicConfig { +interface MonitoringAnthropicConfig extends ClientOptions { apiKey: string posthog: PostHog - baseURL?: string } export class PostHogAnthropic extends AnthropicOriginal { @@ -70,10 +71,9 @@ export class WrappedMessages extends AnthropicOriginal.Messages { const { providerParams: anthropicParams, posthogParams } = extractPosthogParams(body) const startTime = Date.now() - const parentPromise = super.create(anthropicParams, options) - if (anthropicParams.stream) { - return parentPromise.then((value) => { + const parentPromise = super.create(anthropicParams, options) + const wrappedPromise = parentPromise.then((value) => { let accumulatedContent = '' const contentBlocks: FormattedContentItem[] = [] const toolsInProgress: Map = new Map() @@ -95,9 +95,12 @@ export class WrappedMessages extends AnthropicOriginal.Messages { cacheReadInputTokens: 0, webSearchCount: 0, } - let rawUsage: Record = {} - if ('tee' in value) { - const [stream1, stream2] = value.tee() + let lastRawUsage: unknown + if (Symbol.asyncIterator in value) { + const [stream1, stream2] = monitoredStreamTee>( + value as Stream, + (iterator, controller) => new Stream(iterator, controller) + ) ;(async () => { try { for await (const chunk of stream1) { @@ -189,17 +192,14 @@ export class WrappedMessages extends AnthropicOriginal.Messages { } if (chunk.type == 'message_start') { - rawUsage = { ...chunk.message.usage } + lastRawUsage = 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) { - rawUsage = { - ...rawUsage, - ...Object.fromEntries(Object.entries(chunk.usage).filter(([, value]) => value != null)), - } + lastRawUsage = chunk.usage 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 +214,7 @@ export class WrappedMessages extends AnthropicOriginal.Messages { } } } - usage.rawUsage = rawUsage + usage.rawUsage = lastRawUsage const latency = (Date.now() - startTime) / 1000 const timeToFirstToken = firstTokenTime !== undefined ? (firstTokenTime - startTime) / 1000 : undefined @@ -279,8 +279,11 @@ export class WrappedMessages extends AnthropicOriginal.Messages { return stream2 } return value - }) as APIPromise> + }) + + return preserveProviderPromise(parentPromise, wrappedPromise, { requestIdHeader: 'request-id' }) } else { + const parentPromise = super.create(anthropicParams, options) const wrappedPromise = parentPromise.then( async (result) => { if ('content' in result) { @@ -331,9 +334,9 @@ export class WrappedMessages extends AnthropicOriginal.Messages { }) throw error } - ) as APIPromise + ) - return wrappedPromise + return preserveProviderPromise(parentPromise, wrappedPromise, { requestIdHeader: 'request-id' }) } } } diff --git a/packages/ai/src/langchain/callbacks.ts b/packages/ai/src/langchain/callbacks.ts index 79be3a09e5..69b07c8834 100644 --- a/packages/ai/src/langchain/callbacks.ts +++ b/packages/ai/src/langchain/callbacks.ts @@ -1,4 +1,4 @@ -import { PostHog } from 'posthog-node' +import { EventMessage, PostHog } from 'posthog-node' import { withPrivacyMode, getModelParams, toContentString } from '../utils' import { BaseCallbackHandler } from '@langchain/core/callbacks/base' import { version } from '../../package.json' @@ -12,7 +12,14 @@ 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, +// and future subclasses) exposes a prototype getter `is_bubble_up` that returns true, which the +// LangGraph runtime itself uses to distinguish control flow from real failures. The getter reads +// as undefined on ordinary Errors and works across duplicated LangGraph package copies. +const isLangGraphControlFlow = (error: Error): boolean => + (error as Error & { is_bubble_up?: boolean }).is_bubble_up === true interface SpanMetadata { /** Name of the trace/span (e.g. chain name) */ @@ -378,6 +385,14 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { return parentRunId } + private _safeCapture(message: EventMessage): void { + try { + this.client.capture(message) + } catch { + // Telemetry delivery must never affect the LangChain callback lifecycle. + } + } + private _popRunAndCaptureTraceOrSpan( runId: string, parentRunId: string | undefined, @@ -425,12 +440,26 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { eventProperties['$process_person_profile'] = false } if (outputs instanceof Error) { - eventProperties['$ai_error'] = stringifyError(outputs) - eventProperties['$ai_is_error'] = true + if (isLangGraphControlFlow(outputs)) { + // GraphInterrupt carries the pending interrupts (e.g. the question posed to a human). + // Surface them under the same `__interrupt__` key LangGraph hands back to the caller, + // so an interrupted span stays distinguishable from a node that returned nothing. + const interrupts = (outputs as Error & { interrupts?: unknown }).interrupts + if (interrupts !== undefined) { + eventProperties['$ai_output_state'] = withPrivacyMode( + this.client, + this.privacyMode, + sanitizeLangChain({ __interrupt__: interrupts }) + ) + } + } else { + eventProperties['$ai_error'] = stringifyError(outputs) + eventProperties['$ai_is_error'] = true + } } else if (outputs !== undefined) { eventProperties['$ai_output_state'] = withPrivacyMode(this.client, this.privacyMode, sanitizeLangChain(outputs)) } - this.client.capture({ + this._safeCapture({ distinctId: this.distinctId ? this.distinctId.toString() : runId, event: eventName, properties: eventProperties, @@ -501,13 +530,6 @@ 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 } @@ -553,7 +575,7 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { eventProperties['$process_person_profile'] = false } - this.client.capture({ + this._safeCapture({ distinctId: this.distinctId ? this.distinctId.toString() : traceId, event: '$ai_generation', properties: eventProperties, @@ -667,110 +689,20 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { return undefined } const gen = lastGeneration[0] + const messageResponseMetadata = (gen as any).message?.response_metadata + const generationResponseMetadata = gen.generationInfo?.response_metadata + const stopReason = + messageResponseMetadata?.finish_reason || + messageResponseMetadata?.stop_reason || + gen.generationInfo?.finish_reason || + generationResponseMetadata?.stop_reason || + generationResponseMetadata?.finish_reason || + gen.generationInfo?.stop_reason - // Check generationInfo for finish_reason (OpenAI format) - if (gen.generationInfo?.finish_reason) { - return String(gen.generationInfo.finish_reason) - } - - // Check generationInfo for response_metadata.stop_reason (Anthropic format) - if (gen.generationInfo?.response_metadata?.stop_reason) { - return String(gen.generationInfo.response_metadata.stop_reason) - } - - // Check message response_metadata for finish_reason (common LangChain format) - if (gen.generationInfo?.response_metadata?.finish_reason) { - return String(gen.generationInfo.response_metadata.finish_reason) - } - - // Check for stop_reason directly in generationInfo - if (gen.generationInfo?.stop_reason) { - return String(gen.generationInfo.stop_reason) - } - - return undefined - } - - 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] + return stopReason != null ? String(stopReason) : undefined } - private _parseUsageModel( - usage: any, - provider?: string, - model?: string, - inputIncludesCacheTokens = true, - rawUsage?: any - ): [number, number, Record] { + private _parseUsageModel(usage: any, provider?: string, model?: string): [number, number, Record] { const conversionList: Array<[string, 'input' | 'output']> = [ ['promptTokens', 'input'], ['completionTokens', 'output'], @@ -819,33 +751,6 @@ 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 @@ -911,7 +816,7 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { isAnthropic = true } - if (isAnthropic && inputIncludesCacheTokens && parsedUsage.input) { + if (isAnthropic && parsedUsage.input) { const cacheTokens = (additionalTokenData.cacheReadInputTokens || 0) + (additionalTokenData.cacheWriteInputTokens || 0) if (cacheTokens > 0) { @@ -923,83 +828,39 @@ export class LangChainCallbackHandler extends BaseCallbackHandler { } private parseUsage(response: LLMResult, provider?: string, model?: string): [number, number, Record] { - 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) - } - + let llmUsage: [number, number, Record] = [0, 0, {}] const llmUsageKeys = ['token_usage', 'usage', 'tokenUsage'] if (response.llmOutput != null) { - for (const key of llmUsageKeys) { - const llmUsage = response.llmOutput[key] - if (!isNonEmptyUsage(llmUsage)) { - continue - } - return this._parseUsageModel(llmUsage, provider, model, key !== 'usage', llmUsage) + const key = llmUsageKeys.find((k) => response.llmOutput?.[k] != null) + if (key) { + llmUsage = this._parseUsageModel(response.llmOutput[key], provider, model) } } - if (isNonEmptyUsage(rawGenerationUsage)) { - return this._parseUsageModel(rawGenerationUsage, provider, model, false, rawGenerationUsage) - } - if (isNonEmptyUsage(fallbackGenerationUsage)) { - return this._parseUsageModel(fallbackGenerationUsage, provider, model) + // 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 + } + } + } } - return [0, 0, {}] + return llmUsage } } From 62340941cb1c55d03cc2ec95f9a7ecfa72b498e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Gouv=C3=AAa?= <64869705+gouveags@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:22:43 -0300 Subject: [PATCH 7/7] fix(ai): resolve cache TTL handling after main sync Keep current main streaming and callback behavior while restoring the reviewed Anthropic and Bedrock cache TTL extraction, including the normalized generation fallback required by current LangChain result shapes. --- packages/ai/src/anthropic/index.ts | 11 +- packages/ai/src/langchain/callbacks.ts | 217 +++++++++++++++++++++---- 2 files changed, 196 insertions(+), 32 deletions(-) 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, {}] } }