diff --git a/.changeset/langchain-js-token-total-metrics.md b/.changeset/langchain-js-token-total-metrics.md new file mode 100644 index 000000000..548b13b62 --- /dev/null +++ b/.changeset/langchain-js-token-total-metrics.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix: Emit canonical LangChain JS total token metrics diff --git a/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/log-payloads.json b/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/log-payloads.json index d916a030c..c1dbe44e6 100644 --- a/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/log-payloads.json +++ b/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/log-payloads.json @@ -283,6 +283,7 @@ "prompt_cached_tokens": "", "prompt_tokens": "", "start": 0, + "tokens": "", "total_tokens": "" }, "output": { diff --git a/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.json b/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.json index ed1dba52c..136d4d5f5 100644 --- a/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.json +++ b/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.json @@ -116,6 +116,7 @@ "completion_tokens", "prompt_cached_tokens", "prompt_tokens", + "tokens", "total_tokens" ] } diff --git a/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.txt b/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.txt index 2e3c03b9b..7dced8fe2 100644 --- a/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.txt +++ b/e2e/scenarios/langgraph-auto-instrumentation/__snapshots__/span-tree.txt @@ -116,6 +116,7 @@ span_tree: │ "completion_tokens", │ "prompt_cached_tokens", │ "prompt_tokens", + │ "tokens", │ "total_tokens" │ ] └── sayBye [task] diff --git a/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.json b/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.json index 1bfa61b77..e1ea02b4f 100644 --- a/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.json +++ b/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.json @@ -163,6 +163,7 @@ "completion_tokens": 2, "prompt_cached_tokens": 0, "prompt_tokens": 12, + "tokens": 14, "total_tokens": 14 } } @@ -433,6 +434,7 @@ "completion_tokens": 2, "prompt_cached_tokens": 0, "prompt_tokens": 18, + "tokens": 20, "total_tokens": 20 } } @@ -761,6 +763,7 @@ "prompt_cached_tokens": 0, "prompt_tokens": 22, "time_to_first_token": 0, + "tokens": 28, "total_tokens": 28 } } @@ -996,6 +999,7 @@ "completion_tokens": 16, "prompt_cached_tokens": 0, "prompt_tokens": 72, + "tokens": 88, "total_tokens": 88 } } @@ -1259,6 +1263,7 @@ "completion_tokens": 21, "prompt_cached_tokens": 0, "prompt_tokens": 76, + "tokens": 97, "total_tokens": 97 } }, @@ -1533,6 +1538,7 @@ "completion_tokens": 11, "prompt_cached_tokens": 0, "prompt_tokens": 106, + "tokens": 117, "total_tokens": 117 } } diff --git a/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.txt b/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.txt index 8558e7123..06687a93b 100644 --- a/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.txt +++ b/e2e/scenarios/wrap-langchain-js-traces/__snapshots__/span-tree.txt @@ -162,6 +162,7 @@ span_tree: │ "completion_tokens": 2, │ "prompt_cached_tokens": 0, │ "prompt_tokens": 12, + │ "tokens": 14, │ "total_tokens": 14 │ } ├── langchain-chain-operation @@ -564,6 +565,7 @@ span_tree: │ "completion_tokens": 2, │ "prompt_cached_tokens": 0, │ "prompt_tokens": 18, + │ "tokens": 20, │ "total_tokens": 20 │ } ├── langchain-stream-operation @@ -735,6 +737,7 @@ span_tree: │ "prompt_cached_tokens": 0, │ "prompt_tokens": 22, │ "time_to_first_token": 0, + │ "tokens": 28, │ "total_tokens": 28 │ } ├── langchain-tool-operation @@ -962,6 +965,7 @@ span_tree: │ "completion_tokens": 16, │ "prompt_cached_tokens": 0, │ "prompt_tokens": 72, + │ "tokens": 88, │ "total_tokens": 88 │ } └── langchain-tool-result-operation @@ -1217,6 +1221,7 @@ span_tree: │ "completion_tokens": 21, │ "prompt_cached_tokens": 0, │ "prompt_tokens": 76, + │ "tokens": 97, │ "total_tokens": 97 │ } └── ChatOpenAI [llm] @@ -1487,5 +1492,6 @@ span_tree: "completion_tokens": 11, "prompt_cached_tokens": 0, "prompt_tokens": 106, + "tokens": 117, "total_tokens": 117 } diff --git a/js/src/wrappers/langchain/callback-handler.test.ts b/js/src/wrappers/langchain/callback-handler.test.ts new file mode 100644 index 000000000..45995d5e6 --- /dev/null +++ b/js/src/wrappers/langchain/callback-handler.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, it, vi } from "vitest"; +import { BraintrustLangChainCallbackHandler } from "./callback-handler"; + +function createHarness() { + const logs: unknown[] = []; + const span = { + log: vi.fn((event: unknown) => { + logs.push(event); + }), + end: vi.fn(), + }; + const parent = { + startSpan: vi.fn(() => span), + }; + const handler = new BraintrustLangChainCallbackHandler({ + parent: parent as never, + }); + + return { handler, logs, parent, span }; +} + +async function finishChatModelRun(output: unknown) { + const harness = createHarness(); + + await harness.handler.handleChatModelStart( + { name: "ChatOpenAI" }, + [[{ role: "user", content: "hello" }]], + "run-1", + ); + await harness.handler.handleLLMEnd(output as never, "run-1"); + + const endLog = harness.logs.at(-1) as { + metrics?: Record; + }; + return { ...harness, endLog }; +} + +describe("BraintrustLangChainCallbackHandler metrics", () => { + it("synthesizes tokens from message usage metadata prompt and completion counts", async () => { + const { endLog } = await finishChatModelRun({ + generations: [ + [ + { + message: { + usage_metadata: { + input_tokens: 10, + output_tokens: 2, + }, + }, + }, + ], + ], + }); + + expect(endLog.metrics).toEqual({ + prompt_tokens: 10, + completion_tokens: 2, + tokens: 12, + }); + }); + + it("prefers explicit total tokens from message usage metadata", async () => { + const { endLog } = await finishChatModelRun({ + generations: [ + [ + { + message: { + usage_metadata: { + input_tokens: 10, + output_tokens: 2, + total_tokens: 99, + }, + }, + }, + ], + ], + }); + + expect(endLog.metrics).toEqual({ + total_tokens: 99, + prompt_tokens: 10, + completion_tokens: 2, + tokens: 99, + }); + }); + + it("synthesizes tokens from llmOutput token usage prompt and completion counts", async () => { + const { endLog } = await finishChatModelRun({ + llmOutput: { + tokenUsage: { + promptTokens: 10, + completionTokens: 2, + }, + }, + }); + + expect(endLog.metrics).toEqual({ + prompt_tokens: 10, + completion_tokens: 2, + tokens: 12, + }); + }); + + it("preserves cache metrics from message usage metadata", async () => { + const { endLog } = await finishChatModelRun({ + generations: [ + [ + { + message: { + usage_metadata: { + input_tokens: 10, + output_tokens: 2, + input_token_details: { + cache_creation: 4, + cache_read: 3, + }, + }, + }, + }, + ], + ], + }); + + expect(endLog.metrics).toEqual({ + prompt_tokens: 10, + completion_tokens: 2, + prompt_cache_creation_tokens: 4, + prompt_cached_tokens: 3, + tokens: 12, + }); + }); +}); diff --git a/js/src/wrappers/langchain/callback-handler.ts b/js/src/wrappers/langchain/callback-handler.ts index e5612735d..a36ac9cff 100644 --- a/js/src/wrappers/langchain/callback-handler.ts +++ b/js/src/wrappers/langchain/callback-handler.ts @@ -438,6 +438,23 @@ function cleanObject(obj: Record): Record { ) as Record; } +function normalizeTokenMetrics( + obj: Record, +): Record { + const metrics = cleanObject(obj); + + if (metrics.total_tokens !== undefined) { + metrics.tokens = metrics.total_tokens; + } else if ( + metrics.prompt_tokens !== undefined && + metrics.completion_tokens !== undefined + ) { + metrics.tokens = metrics.prompt_tokens + metrics.completion_tokens; + } + + return metrics; +} + function walkGenerations( response: LangChainLLMResult, ): Record[] { @@ -497,7 +514,7 @@ function getMetricsFromResponse( } const inputTokenDetails = usageMetadata.input_token_details; - return cleanObject({ + return normalizeTokenMetrics({ total_tokens: usageMetadata.total_tokens, prompt_tokens: usageMetadata.input_tokens, completion_tokens: usageMetadata.output_tokens, @@ -517,7 +534,7 @@ function getMetricsFromResponse( ? llmOutput.estimatedTokens : {}; - return cleanObject({ + return normalizeTokenMetrics({ total_tokens: tokenUsage.totalTokens, prompt_tokens: tokenUsage.promptTokens, completion_tokens: tokenUsage.completionTokens,