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

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

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

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()