From 2dd93bdb6a350a565f3d013e5c6c551a957c42ba Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 20 Apr 2026 10:35:31 +0200 Subject: [PATCH 1/2] feat: Capture thinking with cohere Fixes https://github.com/braintrustdata/braintrust-sdk-javascript/issues/1845 --- .../cohere-v7-20-0.span-events.json | 54 ++++- .../cohere-v7-21-0.span-events.json | 54 ++++- .../__snapshots__/cohere-v7.span-events.json | 54 ++++- .../__snapshots__/cohere-v8.span-events.json | 54 ++++- .../cohere-instrumentation/assertions.ts | 90 +++++++- .../scenario.cohere-v7.mjs | 9 +- .../scenario.cohere-v7.ts | 9 +- .../cohere-instrumentation/scenario.impl.mjs | 58 ++++- .../cohere-instrumentation/scenario.test.ts | 10 + .../auto-instrumentations/configs/cohere.ts | 52 +++++ .../plugins/cohere-plugin.test.ts | 119 ++++++++++ .../instrumentation/plugins/cohere-plugin.ts | 206 +++++++++++++++++- js/src/vendor-sdk-types/cohere.ts | 10 + 13 files changed, 745 insertions(+), 34 deletions(-) diff --git a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-20-0.span-events.json b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-20-0.span-events.json index f31d4de70..72036b039 100644 --- a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-20-0.span-events.json +++ b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-20-0.span-events.json @@ -86,6 +86,48 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream-thinking" + }, + "metric_keys": [], + "name": "cohere-chat-stream-thinking-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "command-a-reasoning-08-2025", + "provider": "cohere", + "thinking": { + "tokenBudget": 128, + "type": "enabled" + } + }, + "metric_keys": [ + "completion_tokens", + "prompt_cached_tokens", + "prompt_tokens", + "reasoning_tokens", + "time_to_first_token", + "tokens" + ], + "name": "cohere.chatStream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -95,7 +137,7 @@ "metric_keys": [], "name": "cohere-embed-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -114,9 +156,9 @@ ], "name": "cohere.embed", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -129,7 +171,7 @@ "metric_keys": [], "name": "cohere-rerank-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -149,9 +191,9 @@ ], "name": "cohere.rerank", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-21-0.span-events.json b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-21-0.span-events.json index f31d4de70..72036b039 100644 --- a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-21-0.span-events.json +++ b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-21-0.span-events.json @@ -86,6 +86,48 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream-thinking" + }, + "metric_keys": [], + "name": "cohere-chat-stream-thinking-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "command-a-reasoning-08-2025", + "provider": "cohere", + "thinking": { + "tokenBudget": 128, + "type": "enabled" + } + }, + "metric_keys": [ + "completion_tokens", + "prompt_cached_tokens", + "prompt_tokens", + "reasoning_tokens", + "time_to_first_token", + "tokens" + ], + "name": "cohere.chatStream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -95,7 +137,7 @@ "metric_keys": [], "name": "cohere-embed-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -114,9 +156,9 @@ ], "name": "cohere.embed", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -129,7 +171,7 @@ "metric_keys": [], "name": "cohere-rerank-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -149,9 +191,9 @@ ], "name": "cohere.rerank", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7.span-events.json b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7.span-events.json index f31d4de70..72036b039 100644 --- a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7.span-events.json +++ b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7.span-events.json @@ -86,6 +86,48 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream-thinking" + }, + "metric_keys": [], + "name": "cohere-chat-stream-thinking-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "command-a-reasoning-08-2025", + "provider": "cohere", + "thinking": { + "tokenBudget": 128, + "type": "enabled" + } + }, + "metric_keys": [ + "completion_tokens", + "prompt_cached_tokens", + "prompt_tokens", + "reasoning_tokens", + "time_to_first_token", + "tokens" + ], + "name": "cohere.chatStream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -95,7 +137,7 @@ "metric_keys": [], "name": "cohere-embed-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -114,9 +156,9 @@ ], "name": "cohere.embed", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -129,7 +171,7 @@ "metric_keys": [], "name": "cohere-rerank-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -149,9 +191,9 @@ ], "name": "cohere.rerank", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v8.span-events.json b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v8.span-events.json index 9a11e3ffa..ac3ba97a0 100644 --- a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v8.span-events.json +++ b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v8.span-events.json @@ -86,6 +86,48 @@ ], "type": "llm" }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream-thinking" + }, + "metric_keys": [], + "name": "cohere-chat-stream-thinking-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "command-a-reasoning-08-2025", + "provider": "cohere", + "thinking": { + "tokenBudget": 128, + "type": "enabled" + } + }, + "metric_keys": [ + "completion_tokens", + "prompt_cached_tokens", + "prompt_tokens", + "reasoning_tokens", + "time_to_first_token", + "tokens" + ], + "name": "cohere.chatStream", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, { "has_input": false, "has_output": false, @@ -95,7 +137,7 @@ "metric_keys": [], "name": "cohere-embed-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -115,9 +157,9 @@ ], "name": "cohere.embed", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -130,7 +172,7 @@ "metric_keys": [], "name": "cohere-rerank-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -150,9 +192,9 @@ ], "name": "cohere.rerank", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" } diff --git a/e2e/scenarios/cohere-instrumentation/assertions.ts b/e2e/scenarios/cohere-instrumentation/assertions.ts index 3394006a9..25932f96c 100644 --- a/e2e/scenarios/cohere-instrumentation/assertions.ts +++ b/e2e/scenarios/cohere-instrumentation/assertions.ts @@ -13,6 +13,7 @@ import { ROOT_NAME, SCENARIO_NAME } from "./constants.mjs"; type RunCohereScenario = (harness: { runNodeScenarioDir: (options: { entry: string; + env?: Record; nodeArgs: string[]; runContext?: { variantKey: string }; scenarioDir: string; @@ -20,6 +21,7 @@ type RunCohereScenario = (harness: { }) => Promise; runScenarioDir: (options: { entry: string; + env?: Record; runContext?: { variantKey: string }; scenarioDir: string; timeoutMs: number; @@ -35,16 +37,23 @@ function findCohereSpan( return spans.find((candidate) => candidate.output !== undefined) ?? spans[0]; } -function buildSpanSummary(events: CapturedLogEvent[]): Json { +function buildSpanSummary( + events: CapturedLogEvent[], + supportsThinking: boolean, +): Json { const chatOperation = findLatestSpan(events, "cohere-chat-operation"); const chatStreamOperation = findLatestSpan( events, "cohere-chat-stream-operation", ); + const chatStreamThinkingOperation = findLatestSpan( + events, + "cohere-chat-stream-thinking-operation", + ); const embedOperation = findLatestSpan(events, "cohere-embed-operation"); const rerankOperation = findLatestSpan(events, "cohere-rerank-operation"); - return [ + const summaryEvents = [ findLatestSpan(events, ROOT_NAME), chatOperation, findCohereSpan(events, chatOperation?.span.id, "cohere.chat"), @@ -54,7 +63,22 @@ function buildSpanSummary(events: CapturedLogEvent[]): Json { findCohereSpan(events, embedOperation?.span.id, "cohere.embed"), rerankOperation, findCohereSpan(events, rerankOperation?.span.id, "cohere.rerank"), - ].map((event) => + ]; + + if (supportsThinking) { + summaryEvents.splice( + 5, + 0, + chatStreamThinkingOperation, + findCohereSpan( + events, + chatStreamThinkingOperation?.span.id, + "cohere.chatStream", + ), + ); + } + + return summaryEvents.map((event) => summarizeWrapperContract(event!, [ "document_count", "inputType", @@ -62,6 +86,7 @@ function buildSpanSummary(events: CapturedLogEvent[]): Json { "operation", "provider", "scenario", + "thinking", "topN", ]), ) as Json; @@ -71,6 +96,7 @@ export function defineCohereInstrumentationAssertions(options: { name: string; runScenario: RunCohereScenario; snapshotName: string; + supportsThinking: boolean; testFileUrl: string; timeoutMs: number; }): void { @@ -132,6 +158,60 @@ export function defineCohereInstrumentationAssertions(options: { expect(chatStreamSpan?.output).toBeDefined(); }); + if (options.supportsThinking) { + test("captures reasoning content for chatStream", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "cohere-chat-stream-thinking-operation", + ); + const span = findCohereSpan( + events, + operation?.span.id, + "cohere.chatStream", + ); + const output = span?.output as + | { + content?: Array<{ + text?: string; + thinking?: string; + type?: string; + }>; + } + | undefined; + const metrics = (span?.metrics ?? {}) as Record; + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.row.metadata).toMatchObject({ + model: "command-a-reasoning-08-2025", + provider: "cohere", + thinking: { + tokenBudget: 128, + type: "enabled", + }, + }); + expect(metrics).toMatchObject({ + completion_tokens: expect.any(Number), + prompt_tokens: expect.any(Number), + reasoning_tokens: expect.any(Number), + time_to_first_token: expect.any(Number), + }); + expect( + output?.content?.some( + (block) => + block.type === "thinking" && typeof block.thinking === "string", + ), + ).toBe(true); + expect( + output?.content?.some( + (block) => block.type === "text" && typeof block.text === "string", + ), + ).toBe(true); + }); + } + test("captures embed span", testConfig, () => { const operation = findLatestSpan(events, "cohere-embed-operation"); const span = findCohereSpan(events, operation?.span.id, "cohere.embed"); @@ -164,7 +244,9 @@ export function defineCohereInstrumentationAssertions(options: { test("matches span snapshot", testConfig, async () => { await expect( - formatJsonFileSnapshot(buildSpanSummary(events)), + formatJsonFileSnapshot( + buildSpanSummary(events, options.supportsThinking), + ), ).toMatchFileSnapshot(spanSnapshotPath); }); }); diff --git a/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.mjs b/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.mjs index 04a97bcf5..6b0f4d5bd 100644 --- a/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.mjs +++ b/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.mjs @@ -1,9 +1,16 @@ -import { CohereClient as CohereClientV7 } from "cohere-sdk-v7"; +import { + CohereClient as CohereClientV7, + CohereClientV2 as CohereClientV7V2, +} from "cohere-sdk-v7"; import { runMain } from "../../helpers/provider-runtime.mjs"; import { runAutoCohereInstrumentation } from "./scenario.impl.mjs"; runMain(async () => runAutoCohereInstrumentation(CohereClientV7, { apiVersion: "v7", + ThinkingCohereClient: + process.env.COHERE_SUPPORTS_THINKING === "1" + ? CohereClientV7V2 + : undefined, }), ); diff --git a/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.ts b/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.ts index c7b926f6e..3a01391a8 100644 --- a/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.ts +++ b/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.ts @@ -1,5 +1,8 @@ import { wrapCohere } from "braintrust"; -import { CohereClient as CohereClientV7 } from "cohere-sdk-v7"; +import { + CohereClient as CohereClientV7, + CohereClientV2 as CohereClientV7V2, +} from "cohere-sdk-v7"; import { runMain } from "../../helpers/scenario-runtime"; import { runWrappedCohereInstrumentation } from "./scenario.impl.mjs"; @@ -7,5 +10,9 @@ runMain(async () => runWrappedCohereInstrumentation(CohereClientV7, { apiVersion: "v7", decorateClient: wrapCohere, + ThinkingCohereClient: + process.env.COHERE_SUPPORTS_THINKING === "1" + ? CohereClientV7V2 + : undefined, }), ); diff --git a/e2e/scenarios/cohere-instrumentation/scenario.impl.mjs b/e2e/scenarios/cohere-instrumentation/scenario.impl.mjs index 410919de0..b1e74a8a6 100644 --- a/e2e/scenarios/cohere-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/cohere-instrumentation/scenario.impl.mjs @@ -15,7 +15,7 @@ import { SCENARIO_NAME, } from "./constants.mjs"; -export const COHERE_SCENARIO_TIMEOUT_MS = 120_000; +export const COHERE_SCENARIO_TIMEOUT_MS = 240_000; export const COHERE_SCENARIO_SPECS = [ { @@ -23,6 +23,7 @@ export const COHERE_SCENARIO_SPECS = [ autoEntry: "scenario.cohere-v7.mjs", dependencyName: "cohere-sdk-v7-14-0", snapshotName: "cohere-v7-14-0", + supportsThinking: false, wrapperEntry: "scenario.cohere-v7.ts", }, { @@ -84,6 +85,36 @@ function getChatRequest(apiVersion, { stream = false } = {}) { }; } +function getThinkingChatRequest() { + return { + model: "command-a-reasoning-08-2025", + messages: [ + { + role: "user", + content: "What is 2+2? Reply with the number only.", + }, + ], + maxTokens: 256, + temperature: 0, + thinking: { + type: "enabled", + tokenBudget: 128, + }, + }; +} + +function shouldRunThinkingScenario(apiVersion) { + if (process.env.COHERE_SUPPORTS_THINKING === "1") { + return true; + } + + if (process.env.COHERE_SUPPORTS_THINKING === "0") { + return false; + } + + return apiVersion === "v8"; +} + function getEmbedRequest(apiVersion) { if (apiVersion === "v8") { return { @@ -128,7 +159,7 @@ function getRerankRequest(apiVersion) { async function runCohereInstrumentationScenario( CohereClient, - { apiVersion, decorateClient } = {}, + { apiVersion, decorateClient, ThinkingCohereClient } = {}, ) { const apiKey = getApiKey(); if (!apiKey) { @@ -139,6 +170,16 @@ async function runCohereInstrumentationScenario( token: apiKey, }); const client = decorateClient ? decorateClient(baseClient) : baseClient; + const thinkingClientClass = ThinkingCohereClient ?? CohereClient; + const thinkingBaseClient = + thinkingClientClass === CohereClient + ? baseClient + : new thinkingClientClass({ + token: apiKey, + }); + const thinkingClient = decorateClient + ? decorateClient(thinkingBaseClient) + : thinkingBaseClient; await runTracedScenario({ callback: async () => { @@ -157,6 +198,19 @@ async function runCohereInstrumentationScenario( }, ); + if (shouldRunThinkingScenario(apiVersion)) { + await runOperation( + "cohere-chat-stream-thinking-operation", + "chat-stream-thinking", + async () => { + const stream = await thinkingClient.chatStream( + getThinkingChatRequest(), + ); + await collectAsync(stream); + }, + ); + } + await runOperation("cohere-embed-operation", "embed", async () => { await client.embed(getEmbedRequest(apiVersion)); }); diff --git a/e2e/scenarios/cohere-instrumentation/scenario.test.ts b/e2e/scenarios/cohere-instrumentation/scenario.test.ts index 652733408..70c6ad6ea 100644 --- a/e2e/scenarios/cohere-instrumentation/scenario.test.ts +++ b/e2e/scenarios/cohere-instrumentation/scenario.test.ts @@ -25,18 +25,24 @@ const cohereScenarios = await Promise.all( ); for (const scenario of cohereScenarios) { + const supportsThinking = scenario.supportsThinking ?? true; + describe(`cohere sdk ${scenario.version}`, () => { defineCohereInstrumentationAssertions({ name: "wrapped instrumentation", runScenario: async ({ runScenarioDir }) => { await runScenarioDir({ entry: scenario.wrapperEntry, + env: { + COHERE_SUPPORTS_THINKING: supportsThinking ? "1" : "0", + }, runContext: { variantKey: scenario.snapshotName }, scenarioDir, timeoutMs: COHERE_SCENARIO_TIMEOUT_MS, }); }, snapshotName: scenario.snapshotName, + supportsThinking, testFileUrl: import.meta.url, timeoutMs: COHERE_SCENARIO_TIMEOUT_MS, }); @@ -46,6 +52,9 @@ for (const scenario of cohereScenarios) { runScenario: async ({ runNodeScenarioDir }) => { await runNodeScenarioDir({ entry: scenario.autoEntry, + env: { + COHERE_SUPPORTS_THINKING: supportsThinking ? "1" : "0", + }, nodeArgs: ["--import", "braintrust/hook.mjs"], runContext: { variantKey: scenario.snapshotName }, scenarioDir, @@ -53,6 +62,7 @@ for (const scenario of cohereScenarios) { }); }, snapshotName: scenario.snapshotName, + supportsThinking, testFileUrl: import.meta.url, timeoutMs: COHERE_SCENARIO_TIMEOUT_MS, }); diff --git a/js/src/auto-instrumentations/configs/cohere.ts b/js/src/auto-instrumentations/configs/cohere.ts index c12d0a594..955038c84 100644 --- a/js/src/auto-instrumentations/configs/cohere.ts +++ b/js/src/auto-instrumentations/configs/cohere.ts @@ -15,6 +15,19 @@ export const cohereConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + { + channelName: cohereChannels.chat.channelName, + module: { + name: "cohere-ai", + versionRange: ">=7.20.0 <8.0.0", + filePath: "api/resources/v2/client/Client.js", + }, + functionQuery: { + className: "V2Client", + methodName: "chat", + kind: "Async", + }, + }, { channelName: cohereChannels.chat.channelName, module: { @@ -54,6 +67,19 @@ export const cohereConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + { + channelName: cohereChannels.chatStream.channelName, + module: { + name: "cohere-ai", + versionRange: ">=7.20.0 <8.0.0", + filePath: "api/resources/v2/client/Client.js", + }, + functionQuery: { + className: "V2Client", + methodName: "chatStream", + kind: "Async", + }, + }, { channelName: cohereChannels.chatStream.channelName, module: { @@ -93,6 +119,19 @@ export const cohereConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + { + channelName: cohereChannels.embed.channelName, + module: { + name: "cohere-ai", + versionRange: ">=7.20.0 <8.0.0", + filePath: "api/resources/v2/client/Client.js", + }, + functionQuery: { + className: "V2Client", + methodName: "embed", + kind: "Async", + }, + }, { channelName: cohereChannels.embed.channelName, module: { @@ -132,6 +171,19 @@ export const cohereConfigs: InstrumentationConfig[] = [ kind: "Async", }, }, + { + channelName: cohereChannels.rerank.channelName, + module: { + name: "cohere-ai", + versionRange: ">=7.20.0 <8.0.0", + filePath: "api/resources/v2/client/Client.js", + }, + functionQuery: { + className: "V2Client", + methodName: "rerank", + kind: "Async", + }, + }, { channelName: cohereChannels.rerank.channelName, module: { diff --git a/js/src/instrumentation/plugins/cohere-plugin.test.ts b/js/src/instrumentation/plugins/cohere-plugin.test.ts index b760e2ae6..c2ea4fbf1 100644 --- a/js/src/instrumentation/plugins/cohere-plugin.test.ts +++ b/js/src/instrumentation/plugins/cohere-plugin.test.ts @@ -13,6 +13,7 @@ describe("parseCohereMetricsFromUsage", () => { tokens: { inputTokens: 10, outputTokens: 4, + reasoning_tokens: 6, }, cachedTokens: 3, }, @@ -25,6 +26,7 @@ describe("parseCohereMetricsFromUsage", () => { ).toEqual({ prompt_tokens: 10, completion_tokens: 4, + reasoning_tokens: 6, tokens: 14, prompt_cached_tokens: 3, search_units: 1, @@ -193,4 +195,121 @@ describe("aggregateCohereChatStreamChunks", () => { ], }); }); + + it("aggregates v8 thinking blocks and reasoning token metrics", () => { + const aggregated = aggregateCohereChatStreamChunks([ + { + type: "message-start", + id: "resp_reasoning", + delta: { + message: { + role: "assistant", + }, + }, + }, + { + type: "content-start", + index: 0, + delta: { + message: { + content: { + type: "thinking", + thinking: "", + }, + }, + }, + }, + { + type: "content-delta", + index: 0, + delta: { + message: { + content: { + thinking: "Let me think. ", + }, + }, + }, + }, + { + type: "content-delta", + index: 0, + delta: { + message: { + content: { + thinking: "2 + 2 = 4.", + }, + }, + }, + }, + { + type: "content-start", + index: 1, + delta: { + message: { + content: { + type: "text", + text: "", + }, + }, + }, + }, + { + type: "content-delta", + index: 1, + delta: { + message: { + content: { + text: "4", + }, + }, + }, + }, + { + type: "tool-plan-delta", + delta: { + message: { + toolPlan: "Answer directly", + }, + }, + }, + { + type: "message-end", + delta: { + finishReason: "COMPLETE", + usage: { + tokens: { + inputTokens: 7, + outputTokens: 3, + reasoning_tokens: 11, + }, + }, + }, + }, + ]); + + expect(aggregated.metadata).toEqual({ + id: "resp_reasoning", + finish_reason: "COMPLETE", + }); + expect(aggregated.metrics).toEqual({ + prompt_tokens: 7, + completion_tokens: 3, + reasoning_tokens: 11, + tokens: 10, + }); + expect(aggregated.output).toEqual({ + role: "assistant", + toolPlan: "Answer directly", + content: [ + { + type: "thinking", + thinking: "Let me think. 2 + 2 = 4.", + }, + { + type: "text", + text: "4", + }, + ], + }); + }); }); diff --git a/js/src/instrumentation/plugins/cohere-plugin.ts b/js/src/instrumentation/plugins/cohere-plugin.ts index 5e337a200..afaea628c 100644 --- a/js/src/instrumentation/plugins/cohere-plugin.ts +++ b/js/src/instrumentation/plugins/cohere-plugin.ts @@ -124,6 +124,11 @@ const CHAT_REQUEST_METADATA_ALLOWLIST = new Set([ "strictTools", "strict_tools", "temperature", + "thinking", + "thinkingTokenBudget", + "thinkingType", + "thinking_token_budget", + "thinking_type", "toolChoice", "tool_choice", ]); @@ -428,6 +433,14 @@ function mergeUsageMetrics( "tokens", tokenContainer.totalTokens ?? tokenContainer.total_tokens, ); + setMetricIfNumber( + metrics, + "reasoning_tokens", + tokenContainer.reasoningTokens ?? + tokenContainer.reasoning_tokens ?? + tokenContainer.thinkingTokens ?? + tokenContainer.thinking_tokens, + ); } const billedUnits = @@ -551,6 +564,129 @@ function extractV8DeltaText(chunk: CohereChatStreamEvent): string | undefined { return undefined; } +type CohereContentBlockType = "text" | "thinking"; + +type AggregatedCohereContentBlock = { + text: string; + thinking: string; + type?: CohereContentBlockType; +}; + +type SerializedCohereContentBlock = + | { + text: string; + type: "text"; + } + | { + thinking: string; + type: "thinking"; + }; + +function getV8ContentIndex(chunk: CohereChatStreamEvent): number { + return typeof chunk.index === "number" ? chunk.index : 0; +} + +function toContentBlockType( + value: unknown, +): CohereContentBlockType | undefined { + return value === "text" || value === "thinking" ? value : undefined; +} + +function getOrCreateContentBlock( + contentBlocksByIndex: Record, + contentBlockOrder: number[], + index: number, +): AggregatedCohereContentBlock { + if (!contentBlockOrder.includes(index)) { + contentBlockOrder.push(index); + } + + if (!(index in contentBlocksByIndex)) { + contentBlocksByIndex[index] = { + text: "", + thinking: "", + }; + } + + return contentBlocksByIndex[index]; +} + +function appendV8ContentBlock( + contentBlocksByIndex: Record, + contentBlockOrder: number[], + index: number, + content: unknown, +): void { + if (typeof content === "string") { + const block = getOrCreateContentBlock( + contentBlocksByIndex, + contentBlockOrder, + index, + ); + block.type ??= "text"; + block.text += content; + return; + } + + if (!isObject(content)) { + return; + } + + const block = getOrCreateContentBlock( + contentBlocksByIndex, + contentBlockOrder, + index, + ); + const contentType = toContentBlockType(content.type); + if (contentType) { + block.type = contentType; + } + + if (typeof content.text === "string") { + block.type ??= "text"; + block.text += content.text; + } + + if (typeof content.thinking === "string") { + block.type ??= "thinking"; + block.thinking += content.thinking; + } +} + +function serializeAggregatedContentBlocks( + contentBlocksByIndex: Record, + contentBlockOrder: number[], +): SerializedCohereContentBlock[] { + return contentBlockOrder + .sort((left, right) => left - right) + .flatMap((index) => { + const block = contentBlocksByIndex[index]; + if (!block) { + return []; + } + + if (block.type === "thinking" && block.thinking.length > 0) { + return [{ type: "thinking", thinking: block.thinking }]; + } + + if (block.text.length > 0) { + return [{ type: "text", text: block.text }]; + } + + if (block.thinking.length > 0) { + return [{ type: "thinking", thinking: block.thinking }]; + } + + return []; + }); +} + +function hasThinkingContent( + contentBlocks: Array<{ type: "text" | "thinking" }>, +): boolean { + return contentBlocks.some((block) => block.type === "thinking"); +} + export function aggregateCohereChatStreamChunks( chunks: CohereChatStreamEvent[], ): { @@ -559,11 +695,14 @@ export function aggregateCohereChatStreamChunks( metadata: Record; } { const textDeltas: string[] = []; + const contentBlocksByIndex: Record = {}; + const contentBlockOrder: number[] = []; const toolCallsByIndex: Record = {}; const toolCallOrder: number[] = []; let terminalResponse: CohereChatResponse | undefined; let role: string | undefined; let finishReason: string | undefined; + let toolPlan = ""; let metadata: Record = {}; let metrics: Record = {}; @@ -642,6 +781,14 @@ export function aggregateCohereChatStreamChunks( } if (eventType === "content-delta") { + appendV8ContentBlock( + contentBlocksByIndex, + contentBlockOrder, + getV8ContentIndex(chunk), + isObject(chunk.delta) && isObject(chunk.delta.message) + ? chunk.delta.message.content + : undefined, + ); const text = extractV8DeltaText(chunk); if (text) { textDeltas.push(text); @@ -649,6 +796,34 @@ export function aggregateCohereChatStreamChunks( continue; } + if (eventType === "content-start") { + appendV8ContentBlock( + contentBlocksByIndex, + contentBlockOrder, + getV8ContentIndex(chunk), + isObject(chunk.delta) && isObject(chunk.delta.message) + ? chunk.delta.message.content + : undefined, + ); + continue; + } + + if (eventType === "tool-plan-delta") { + if (isObject(chunk.delta) && isObject(chunk.delta.message)) { + const deltaToolPlan = + typeof chunk.delta.message.toolPlan === "string" + ? chunk.delta.message.toolPlan + : typeof chunk.delta.message.tool_plan === "string" + ? chunk.delta.message.tool_plan + : undefined; + + if (deltaToolPlan) { + toolPlan += deltaToolPlan; + } + } + continue; + } + if (eventType === "tool-call-start") { const toolCalls = isObject(chunk.delta) && isObject(chunk.delta.message) @@ -721,14 +896,41 @@ export function aggregateCohereChatStreamChunks( .sort((left, right) => left - right) .map((index) => toolCallsByIndex[index]) .filter((toolCall): toolCall is CohereToolCall => isObject(toolCall)); + const aggregatedContentBlocks = serializeAggregatedContentBlocks( + contentBlocksByIndex, + contentBlockOrder, + ); let output: unknown = extractCohereChatOutput(terminalResponse); if (output === undefined) { const mergedText = textDeltas.join(""); - if (mergedToolCalls.length > 0 || role || mergedText.length > 0) { + const shouldUseStructuredContent = + hasThinkingContent(aggregatedContentBlocks) || toolPlan.length > 0; + + if (shouldUseStructuredContent) { + output = { + ...(role ? { role } : {}), + ...(aggregatedContentBlocks.length > 0 + ? { content: aggregatedContentBlocks } + : {}), + ...(toolPlan.length > 0 ? { toolPlan } : {}), + ...(mergedToolCalls.length > 0 ? { toolCalls: mergedToolCalls } : {}), + }; + } else if ( + mergedToolCalls.length > 0 || + role || + mergedText.length > 0 || + aggregatedContentBlocks.length > 0 + ) { + const textContent = + mergedText.length > 0 + ? mergedText + : aggregatedContentBlocks[0]?.type === "text" + ? aggregatedContentBlocks[0].text + : undefined; output = { ...(role ? { role } : {}), - ...(mergedText.length > 0 ? { content: mergedText } : {}), + ...(textContent ? { content: textContent } : {}), ...(mergedToolCalls.length > 0 ? { toolCalls: mergedToolCalls } : {}), }; } diff --git a/js/src/vendor-sdk-types/cohere.ts b/js/src/vendor-sdk-types/cohere.ts index fc6347c45..920c01f3b 100644 --- a/js/src/vendor-sdk-types/cohere.ts +++ b/js/src/vendor-sdk-types/cohere.ts @@ -3,6 +3,10 @@ export type CohereTokenUsage = { input_tokens?: number; outputTokens?: number; output_tokens?: number; + reasoningTokens?: number; + reasoning_tokens?: number; + thinkingTokens?: number; + thinking_tokens?: number; totalTokens?: number; total_tokens?: number; [key: string]: unknown; @@ -68,6 +72,8 @@ export type CohereChatResponse = { message?: { role?: string; content?: unknown; + toolPlan?: string; + tool_plan?: string; toolCalls?: CohereToolCall[]; tool_calls?: CohereToolCall[]; [key: string]: unknown; @@ -96,9 +102,13 @@ export type CohereChatStreamEvent = { usage?: CohereUsageLike; message?: { role?: string; + toolPlan?: string; + tool_plan?: string; content?: | string | { + type?: "text" | "thinking"; + thinking?: string; text?: string; [key: string]: unknown; } From 6e32e74539adfa0f7cb65d333b1da5a5087162aa Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 20 Apr 2026 11:20:54 +0200 Subject: [PATCH 2/2] changeset --- .changeset/honest-mails-show.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/honest-mails-show.md diff --git a/.changeset/honest-mails-show.md b/.changeset/honest-mails-show.md new file mode 100644 index 000000000..ca2cf8bbb --- /dev/null +++ b/.changeset/honest-mails-show.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +feat: Capture thinking with cohere