braintrustdata · Luca Forstner (lforst) · May 4, 2026 · May 4, 2026 · May 4, 2026 · May 4, 2026
diff --git a/.changeset/fine-cameras-allow.md b/.changeset/fine-cameras-allow.md
@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix(cohere): Wrap v2 subclient
diff --git a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-14-0-wrapped.span-events.json b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-14-0-wrapped.span-events.json
@@ -0,0 +1,152 @@
+[
+  {
+    "has_input": false,
+    "has_output": false,
+    "metadata": {
+      "scenario": "cohere-instrumentation"
+    },
+    "metric_keys": [],
+    "name": "cohere-instrumentation-root",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:1>",
+    "span_parents": [],
+    "type": "task"
+  },
+  {
+    "has_input": false,
+    "has_output": false,
+    "metadata": {
+      "operation": "chat"
+    },
+    "metric_keys": [],
+    "name": "cohere-v2-chat-operation",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:2>",
+    "span_parents": [
+      "<span:1>"
+    ],
+    "type": null
+  },
+  {
+    "has_input": true,
+    "has_output": true,
+    "metadata": {
+      "model": "command-a-03-2025",
+      "provider": "cohere"
+    },
+    "metric_keys": [
+      "completion_tokens",
+      "prompt_cached_tokens",
+      "prompt_tokens",
+      "time_to_first_token",
+      "tokens"
+    ],
+    "name": "cohere.chat",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:3>",
+    "span_parents": [
+      "<span:2>"
+    ],
+    "type": "llm"
+  },
+  {
+    "has_input": false,
+    "has_output": false,
+    "metadata": {
+      "operation": "chat-stream"
+    },
+    "metric_keys": [],
+    "name": "cohere-v2-chat-stream-operation",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:4>",
+    "span_parents": [
+      "<span:1>"
+    ],
+    "type": null
+  },
+  {
+    "has_input": true,
+    "has_output": false,
+    "metadata": {
+      "model": "command-a-03-2025",
+      "provider": "cohere"
+    },
+    "metric_keys": [],
+    "name": "cohere.chatStream",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:5>",
+    "span_parents": [
+      "<span:4>"
+    ],
+    "type": "llm"
+  },
+  {
+    "has_input": false,
+    "has_output": false,
+    "metadata": {
+      "operation": "embed"
+    },
+    "metric_keys": [],
+    "name": "cohere-v2-embed-operation",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:6>",
+    "span_parents": [
+      "<span:1>"
+    ],
+    "type": null
+  },
+  {
+    "has_input": true,
+    "has_output": true,
+    "metadata": {
+      "inputType": "search_document",
+      "model": "embed-english-v3.0",
+      "provider": "cohere"
+    },
+    "metric_keys": [
+      "prompt_tokens"
+    ],
+    "name": "cohere.embed",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:7>",
+    "span_parents": [
+      "<span:6>"
+    ],
+    "type": "llm"
+  },
+  {
+    "has_input": false,
+    "has_output": false,
+    "metadata": {
+      "operation": "rerank"
+    },
+    "metric_keys": [],
+    "name": "cohere-v2-rerank-operation",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:8>",
+    "span_parents": [
+      "<span:1>"
+    ],
+    "type": null
+  },
+  {
+    "has_input": true,
+    "has_output": true,
+    "metadata": {
+      "document_count": 3,
+      "model": "rerank-english-v3.0",
+      "provider": "cohere",
+      "topN": 2
+    },
+    "metric_keys": [
+      "search_units"
+    ],
+    "name": "cohere.rerank",
+    "root_span_id": "<span:1>",
+    "span_id": "<span:9>",
+    "span_parents": [
+      "<span:8>"
+    ],
+    "type": "llm"
+  }
+]
diff --git a/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-14-0.span-events.json b/e2e/scenarios/cohere-instrumentation/__snapshots__/cohere-v7-14-0.span-events.json
@@ -19,7 +19,7 @@
       "operation": "chat"
     },
     "metric_keys": [],
-    "name": "cohere-chat-operation",
+    "name": "cohere-v2-chat-operation",
     "root_span_id": "<span:1>",
     "span_id": "<span:2>",
     "span_parents": [
@@ -56,7 +56,7 @@
       "operation": "chat-stream"
     },
     "metric_keys": [],
-    "name": "cohere-chat-stream-operation",
+    "name": "cohere-v2-chat-stream-operation",
     "root_span_id": "<span:1>",
     "span_id": "<span:4>",
     "span_parents": [
@@ -93,7 +93,7 @@
       "operation": "embed"
     },
     "metric_keys": [],
-    "name": "cohere-embed-operation",
+    "name": "cohere-v2-embed-operation",
     "root_span_id": "<span:1>",
     "span_id": "<span:6>",
     "span_parents": [
@@ -127,7 +127,7 @@
       "operation": "rerank"
     },
     "metric_keys": [],
-    "name": "cohere-rerank-operation",
+    "name": "cohere-v2-rerank-operation",
     "root_span_id": "<span:1>",
     "span_id": "<span:8>",
     "span_parents": [

diff --git a/e2e/scenarios/cohere-instrumentation/assertions.ts b/e2e/scenarios/cohere-instrumentation/assertions.ts
@@ -42,21 +42,40 @@ function isCohereProviderLimitError(error: unknown): boolean {
   return message.includes("TooManyRequestsError") || message.includes("429");
 }
 
+function getOperationName(
+  baseName: string,
+  { useV2Namespace }: { useV2Namespace: boolean },
+) {
+  return useV2Namespace
+    ? `cohere-v2-${baseName}-operation`
+    : `cohere-${baseName}-operation`;
+}
+
 function buildSpanSummary(
   events: CapturedLogEvent[],
   supportsThinking: boolean,
+  useV2Namespace: boolean,
 ): Json {
-  const chatOperation = findLatestSpan(events, "cohere-chat-operation");
+  const chatOperation = findLatestSpan(
+    events,
+    getOperationName("chat", { useV2Namespace }),
+  );
   const chatStreamOperation = findLatestSpan(
     events,
-    "cohere-chat-stream-operation",
+    getOperationName("chat-stream", { useV2Namespace }),
   );
   const chatStreamThinkingOperation = findLatestSpan(
     events,
-    "cohere-chat-stream-thinking-operation",
+    getOperationName("chat-stream-thinking", { useV2Namespace }),
+  );
+  const embedOperation = findLatestSpan(
+    events,
+    getOperationName("embed", { useV2Namespace }),
+  );
+  const rerankOperation = findLatestSpan(
+    events,
+    getOperationName("rerank", { useV2Namespace }),
   );
-  const embedOperation = findLatestSpan(events, "cohere-embed-operation");
-  const rerankOperation = findLatestSpan(events, "cohere-rerank-operation");
 
   const summaryEvents = [
     findLatestSpan(events, ROOT_NAME),
@@ -104,6 +123,8 @@ export function defineCohereInstrumentationAssertions(options: {
   supportsThinking: boolean;
   testFileUrl: string;
   timeoutMs: number;
+  requireChatStreamOutput?: boolean;
+  useV2Namespace?: boolean;
 }): void {
   const spanSnapshotPath = resolveFileSnapshotPath(
     options.testFileUrl,
@@ -151,15 +172,21 @@ export function defineCohereInstrumentationAssertions(options: {
         context.skip();
       }
 
-      const chatOperation = findLatestSpan(events, "cohere-chat-operation");
+      const chatOperationName = getOperationName("chat", {
+        useV2Namespace: options.useV2Namespace ?? false,
+      });
+      const chatStreamOperationName = getOperationName("chat-stream", {
+        useV2Namespace: options.useV2Namespace ?? false,
+      });
+      const chatOperation = findLatestSpan(events, chatOperationName);
       const chatSpan = findCohereSpan(
         events,
         chatOperation?.span.id,
         "cohere.chat",
       );
       const chatStreamOperation = findLatestSpan(
         events,
-        "cohere-chat-stream-operation",
+        chatStreamOperationName,
       );
       const chatStreamSpan = findCohereSpan(
         events,
@@ -179,7 +206,9 @@ export function defineCohereInstrumentationAssertions(options: {
       expect(chatStreamSpan?.row.metadata).toMatchObject({
         provider: "cohere",
       });
-      expect(chatStreamSpan?.output).toBeDefined();
+      if (options.requireChatStreamOutput ?? true) {
+        expect(chatStreamSpan?.output).toBeDefined();
+      }
     });
 
     if (options.supportsThinking) {
@@ -194,7 +223,9 @@ export function defineCohereInstrumentationAssertions(options: {
           const root = findLatestSpan(events, ROOT_NAME);
           const operation = findLatestSpan(
             events,
-            "cohere-chat-stream-thinking-operation",
+            getOperationName("chat-stream-thinking", {
+              useV2Namespace: options.useV2Namespace ?? false,
+            }),
           );
           const span = findCohereSpan(
             events,
@@ -250,7 +281,12 @@ export function defineCohereInstrumentationAssertions(options: {
         context.skip();
       }
 
-      const operation = findLatestSpan(events, "cohere-embed-operation");
+      const operation = findLatestSpan(
+        events,
+        getOperationName("embed", {
+          useV2Namespace: options.useV2Namespace ?? false,
+        }),
+      );
       const span = findCohereSpan(events, operation?.span.id, "cohere.embed");
       const output = span?.output as { embedding_length?: number } | undefined;
 
@@ -268,7 +304,12 @@ export function defineCohereInstrumentationAssertions(options: {
         context.skip();
       }
 
-      const operation = findLatestSpan(events, "cohere-rerank-operation");
+      const operation = findLatestSpan(
+        events,
+        getOperationName("rerank", {
+          useV2Namespace: options.useV2Namespace ?? false,
+        }),
+      );
       const span = findCohereSpan(events, operation?.span.id, "cohere.rerank");
 
       expect(operation).toBeDefined();
@@ -290,7 +331,11 @@ export function defineCohereInstrumentationAssertions(options: {
 
       await expect(
         formatJsonFileSnapshot(
-          buildSpanSummary(events, options.supportsThinking),
+          buildSpanSummary(
+            events,
+            options.supportsThinking,
+            options.useV2Namespace ?? false,
+          ),
         ),
       ).toMatchFileSnapshot(spanSnapshotPath);
     });

diff --git a/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.mjs b/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.mjs
@@ -1,16 +1,13 @@
-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, {
+runMain(async () => {
+  const cohere = await import(
+    process.env.COHERE_PACKAGE_NAME ?? "cohere-sdk-v7"
+  );
+
+  await runAutoCohereInstrumentation(cohere.CohereClient, {
     apiVersion: "v7",
-    ThinkingCohereClient:
-      process.env.COHERE_SUPPORTS_THINKING === "1"
-        ? CohereClientV7V2
-        : undefined,
-  }),
-);
+    useV2Namespace: true,
+  });
+});
diff --git a/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.ts b/e2e/scenarios/cohere-instrumentation/scenario.cohere-v7.ts
@@ -1,18 +1,15 @@
 import { wrapCohere } from "braintrust";
-import {
-  CohereClient as CohereClientV7,
-  CohereClientV2 as CohereClientV7V2,
-} from "cohere-sdk-v7";
 import { runMain } from "../../helpers/scenario-runtime";
 import { runWrappedCohereInstrumentation } from "./scenario.impl.mjs";
 
-runMain(async () =>
-  runWrappedCohereInstrumentation(CohereClientV7, {
+runMain(async () => {
+  const cohere = await import(
+    process.env.COHERE_PACKAGE_NAME ?? "cohere-sdk-v7"
+  );
+
+  await runWrappedCohereInstrumentation(cohere.CohereClient, {
     apiVersion: "v7",
     decorateClient: wrapCohere,
-    ThinkingCohereClient:
-      process.env.COHERE_SUPPORTS_THINKING === "1"
-        ? CohereClientV7V2
-        : undefined,
-  }),
-);
+    useV2Namespace: true,
+  });
+});