fix(cohere): add tool spans for chat tool calls (#380)

AbhiPrasad · web-flow · commit 6968fe508631 · 2026-05-01T18:21:02.000-04:00
Emit child TOOL spans for Cohere chat responses that include tool calls, covering both v1 and v2 response shapes. Rename the shared Cohere field accessor to document why tracing handles both dict and SDK model objects. resolves #377
diff --git a/py/src/braintrust/integrations/cohere/cassettes/latest/test_wrap_cohere_chat_v1_tool_call_spans.yaml b/py/src/braintrust/integrations/cohere/cassettes/latest/test_wrap_cohere_chat_v1_tool_call_spans.yaml
@@ -0,0 +1,83 @@
+interactions:
+- request:
+    body: '{"message":"Use the get_weather tool for Paris.","model":"command-a-03-2025","max_tokens":64,"tools":[{"name":"get_weather","description":"Get
+      the weather for a city.","parameter_definitions":{"city":{"description":"City
+      name","type":"str","required":true}}}],"force_single_step":true,"stream":false}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '300'
+      Host:
+      - api.cohere.com
+      User-Agent:
+      - cohere/6.1.0
+      X-Fern-Language:
+      - Python
+      X-Fern-Platform:
+      - darwin/25.2.0
+      X-Fern-Runtime:
+      - python/3.12.12
+      X-Fern-SDK-Name:
+      - cohere
+      X-Fern-SDK-Version:
+      - 6.1.0
+      content-type:
+      - application/json
+    method: POST
+    uri: https://api.cohere.com/v1/chat
+  response:
+    body:
+      string: '{"response_id":"490af82c-d74b-4774-96aa-a48abc77b39d","text":"I will
+        use one or more of the available tools to find the answer","generation_id":"465b7c4c-823c-4932-bd72-cc91b3aa154d","chat_history":[{"role":"USER","message":"Use
+        the get_weather tool for Paris."},{"role":"CHATBOT","message":"I will use
+        one or more of the available tools to find the answer","tool_calls":[{"name":"get_weather","parameters":{"city":"Paris"}}]}],"finish_reason":"COMPLETE","meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":42,"output_tokens":22},"tokens":{"input_tokens":1462,"output_tokens":33},"cached_tokens":0},"tool_calls":[{"name":"get_weather","parameters":{"city":"Paris"}}]}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000
+      Transfer-Encoding:
+      - chunked
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-length:
+      - '684'
+      content-type:
+      - application/json
+      date:
+      - Fri, 01 May 2026 17:35:39 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 GMT
+      num_chars:
+      - '6894'
+      num_tokens:
+      - '64'
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      vary:
+      - Origin,Accept-Encoding
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - 9a84e7af81536f5a0ae6409b072713da
+      x-endpoint-monthly-call-limit:
+      - '1000'
+      x-envoy-upstream-service-time:
+      - '988'
+      x-trial-endpoint-call-limit:
+      - '20'
+      x-trial-endpoint-call-remaining:
+      - '17'
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/py/src/braintrust/integrations/cohere/cassettes/latest/test_wrap_cohere_chat_v2_tool_call_spans.yaml b/py/src/braintrust/integrations/cohere/cassettes/latest/test_wrap_cohere_chat_v2_tool_call_spans.yaml
@@ -0,0 +1,81 @@
+interactions:
+- request:
+    body: '{"model":"command-a-03-2025","messages":[{"role":"user","content":"Use
+      the get_weather tool for Paris."}],"tools":[{"type":"function","function":{"name":"get_weather","description":"Get
+      the weather for a city.","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}}],"max_tokens":64,"tool_choice":"REQUIRED","stream":false}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '361'
+      Host:
+      - api.cohere.com
+      User-Agent:
+      - cohere/6.1.0
+      X-Fern-Language:
+      - Python
+      X-Fern-Platform:
+      - darwin/25.2.0
+      X-Fern-Runtime:
+      - python/3.12.12
+      X-Fern-SDK-Name:
+      - cohere
+      X-Fern-SDK-Version:
+      - 6.1.0
+      content-type:
+      - application/json
+    method: POST
+    uri: https://api.cohere.com/v2/chat
+  response:
+    body:
+      string: '{"id":"74529407-6112-4e14-9cb3-41ab95fe7725","message":{"role":"assistant","tool_plan":"I
+        will use one or more of the available tools to find the answer","tool_calls":[{"id":"get_weather_p1927913dgpp","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Paris\"}"}}]},"finish_reason":"TOOL_CALL","usage":{"billed_units":{"input_tokens":37,"output_tokens":22},"tokens":{"input_tokens":1455,"output_tokens":33},"cached_tokens":0}}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000
+      Transfer-Encoding:
+      - chunked
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-length:
+      - '451'
+      content-type:
+      - application/json
+      date:
+      - Fri, 01 May 2026 17:34:51 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 GMT
+      num_chars:
+      - '6866'
+      num_tokens:
+      - '59'
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      vary:
+      - Origin,Accept-Encoding
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - b3c262976868c73da50611172c4163a1
+      x-endpoint-monthly-call-limit:
+      - '1000'
+      x-envoy-upstream-service-time:
+      - '913'
+      x-trial-endpoint-call-limit:
+      - '20'
+      x-trial-endpoint-call-remaining:
+      - '18'
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/py/src/braintrust/integrations/cohere/test_cohere.py b/py/src/braintrust/integrations/cohere/test_cohere.py
@@ -30,7 +30,8 @@
     V2RerankPatcher,
 )
 from braintrust.integrations.test_utils import assert_metrics_are_valid, verify_autoinstrument_script
-from braintrust.test_helpers import init_test_logger
+from braintrust.span_types import SpanTypeAttribute
+from braintrust.test_helpers import find_spans_by_type, init_test_logger
 
 
 pytest.importorskip("cohere")
@@ -260,6 +261,91 @@ def test_wrap_cohere_chat_v2_sync(memory_logger):
     assert_metrics_are_valid(span["metrics"], start, end)
 
 
+@pytest.mark.vcr
+def test_wrap_cohere_chat_v2_tool_call_spans(memory_logger):
+    if os.environ.get("BRAINTRUST_TEST_PACKAGE_VERSION") != "latest":
+        pytest.skip("v2 tool-call cassette is recorded for the latest Cohere SDK")
+
+    assert not memory_logger.pop()
+    client = wrap_cohere(_v2_client(require_methods=("chat",)))
+
+    response = client.chat(
+        model=CHAT_MODEL,
+        messages=[{"role": "user", "content": "Use the get_weather tool for Paris."}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "description": "Get the weather for a city.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"city": {"type": "string"}},
+                        "required": ["city"],
+                    },
+                },
+            }
+        ],
+        tool_choice="REQUIRED",
+        max_tokens=64,
+    )
+
+    tool_calls = response.message.tool_calls
+    assert tool_calls
+    assert tool_calls[0].function.name == "get_weather"
+
+    spans = memory_logger.pop()
+    llm_spans = find_spans_by_type(spans, SpanTypeAttribute.LLM)
+    tool_spans = find_spans_by_type(spans, SpanTypeAttribute.TOOL)
+
+    assert len(llm_spans) == 1
+    assert len(tool_spans) == 1
+    tool_span = tool_spans[0]
+    assert tool_span["span_attributes"]["name"] == "tool: get_weather"
+    assert tool_span["span_parents"] == [llm_spans[0]["span_id"]]
+    assert tool_span["metadata"]["tool_call_id"] == tool_calls[0].id
+    assert tool_span["metadata"]["tool_type"] == "function"
+    assert "Paris" in str(tool_span["input"])
+
+
+@pytest.mark.vcr
+def test_wrap_cohere_chat_v1_tool_call_spans(memory_logger):
+    if os.environ.get("BRAINTRUST_TEST_PACKAGE_VERSION") != "latest":
+        pytest.skip("v1 tool-call cassette is recorded for the latest Cohere SDK")
+
+    assert not memory_logger.pop()
+    client = wrap_cohere(_v1_client())
+
+    response = client.chat(
+        model=CHAT_MODEL,
+        message="Use the get_weather tool for Paris.",
+        tools=[
+            {
+                "name": "get_weather",
+                "description": "Get the weather for a city.",
+                "parameter_definitions": {"city": {"description": "City name", "type": "str", "required": True}},
+            }
+        ],
+        force_single_step=True,
+        max_tokens=64,
+    )
+
+    tool_calls = response.tool_calls
+    assert tool_calls
+    assert tool_calls[0].name == "get_weather"
+
+    spans = memory_logger.pop()
+    llm_spans = find_spans_by_type(spans, SpanTypeAttribute.LLM)
+    tool_spans = find_spans_by_type(spans, SpanTypeAttribute.TOOL)
+
+    assert len(llm_spans) == 1
+    assert len(tool_spans) == 1
+    tool_span = tool_spans[0]
+    assert tool_span["span_attributes"]["name"] == "tool: get_weather"
+    assert tool_span["span_parents"] == [llm_spans[0]["span_id"]]
+    assert "Paris" in str(tool_span["input"])
+
+
 @pytest.mark.vcr
 def test_wrap_cohere_chat_v1_sync(memory_logger):
     assert not memory_logger.pop()
diff --git a/py/src/braintrust/integrations/cohere/tracing.py b/py/src/braintrust/integrations/cohere/tracing.py
