From 36dd5d8cfc65e02595bc1da9392dbb65c98fa38d Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 4 May 2026 11:23:13 -0700 Subject: [PATCH] fix(mistral): add conversation tool spans Decompose Mistral Conversations tool.execution outputs into child TOOL spans for both sync and streaming responses. Add VCR-backed regression coverage for latest and 1.12.4. --- ...conversations_start_stream_tool_spans.yaml | 115 ++++++++++++++++++ ...l_beta_conversations_start_tool_spans.yaml | 86 +++++++++++++ ...conversations_start_stream_tool_spans.yaml | 115 ++++++++++++++++++ ...l_beta_conversations_start_tool_spans.yaml | 86 +++++++++++++ .../integrations/mistral/test_mistral.py | 52 ++++++++ .../integrations/mistral/tracing.py | 88 +++++++++++--- 6 files changed, 524 insertions(+), 18 deletions(-) create mode 100644 py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml create mode 100644 py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_tool_spans.yaml create mode 100644 py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml create mode 100644 py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_tool_spans.yaml diff --git a/py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml b/py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml new file mode 100644 index 00000000..e75965d9 --- /dev/null +++ b/py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml @@ -0,0 +1,115 @@ +interactions: +- request: + body: '{"inputs":"Use the code interpreter to calculate 21 * 2. Return only the + result.","stream":true,"tools":[{"type":"code_interpreter"}],"model":"mistral-small-latest"}' + headers: + Accept: + - text/event-stream + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '165' + Host: + - api.mistral.ai + content-type: + - application/json + user-agent: + - mistral-client-python/1.12.4 + method: POST + uri: https://api.mistral.ai/v1/conversations + response: + body: + string: 'event: conversation.response.started + + data: {"type":"conversation.response.started","created_at":"2026-05-04T18:17:38.536041Z","conversation_id":"conv_019df435c3a677148535d0889fa2490c"} + + + event: tool.execution.started + + data: {"type":"tool.execution.started","created_at":"2026-05-04T18:17:38.540129Z","output_index":0,"id":"tool_exec_019df435c7ec7497b1feba55b9e05e09","model":"mistral-small-latest","name":"code_interpreter","arguments":"","function":"code_interpreter"} + + + event: tool.execution.delta + + data: {"type":"tool.execution.delta","created_at":"2026-05-04T18:17:38.550567Z","output_index":0,"id":"tool_exec_019df435c7ec7497b1feba55b9e05e09","name":"code_interpreter","arguments":"{\"code\": + \"","function":"code_interpreter"} + + + event: tool.execution.delta + + data: {"type":"tool.execution.delta","created_at":"2026-05-04T18:17:38.569375Z","output_index":0,"id":"tool_exec_019df435c7ec7497b1feba55b9e05e09","name":"code_interpreter","arguments":"21 + * ","function":"code_interpreter"} + + + event: tool.execution.delta + + data: {"type":"tool.execution.delta","created_at":"2026-05-04T18:17:38.588794Z","output_index":0,"id":"tool_exec_019df435c7ec7497b1feba55b9e05e09","name":"code_interpreter","arguments":"2\"}","function":"code_interpreter"} + + + event: tool.execution.done + + data: {"type":"tool.execution.done","created_at":"2026-05-04T18:17:39.246174Z","output_index":0,"id":"tool_exec_019df435c7ec7497b1feba55b9e05e09","name":"code_interpreter","function":"code_interpreter","info":{"code":"21 + * 2","code_output":"42\n"}} + + + event: message.output.delta + + data: {"type":"message.output.delta","created_at":"2026-05-04T18:17:39.496747Z","output_index":1,"id":"msg_019df435cba875c0be02cb5196bb0c6b","content_index":0,"model":"mistral-small-latest","role":"assistant","content":"4"} + + + event: message.output.delta + + data: {"type":"message.output.delta","created_at":"2026-05-04T18:17:39.512024Z","output_index":1,"id":"msg_019df435cba875c0be02cb5196bb0c6b","content_index":0,"model":"mistral-small-latest","role":"assistant","content":"2"} + + + event: conversation.response.done + + data: {"type":"conversation.response.done","created_at":"2026-05-04T18:17:39.534463Z","usage":{"prompt_tokens":86,"completion_tokens":19,"total_tokens":110,"connector_tokens":5,"connectors":{"code_interpreter":1}}} + + + ' + headers: + CF-RAY: + - 9f6980d85e564d21-SJC + Connection: + - keep-alive + Content-Type: + - text/event-stream; charset=utf-8 + Date: + - Mon, 04 May 2026 18:17:38 GMT + Server: + - cloudflare + Strict-Transport-Security: + - max-age=15552000; includeSubDomains; preload + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + mistral-correlation-id: + - 019df435-c391-7b83-a2d1-5055d23b514f + set-cookie: + - __cf_bm=zMcNJeDI01Cs_2HPmgcbx8NK2vBd6q0te_lJkWy.Chg-1777918657.3398619-1.0.1.1-UHlbO0PlTlqBzpH16gQaciZIDn9Vo25APxvWZIsXdG5D5HBmLAro4wQfeyVtRwVUkk8hy_8aO5R90xX_XAIj0A.Ki1OeI.miQRj.WtJ71y1OK8CKJeoIPp2mVPaq7K8V; + HttpOnly; Secure; Path=/; Domain=mistral.ai; Expires=Mon, 04 May 2026 18:47:38 + GMT + - _cfuvid=v6_zR1PTQDx2W0GmJI_6mxxXKChwk4Yf3Pe6CfQjgsI-1777918657.3398619-1.0.1.1-MYti7JhXogZfcgmc7gCEjnj57SSCgtvPropKKHhzXYw; + HttpOnly; SameSite=None; Secure; Path=/; Domain=mistral.ai + x-envoy-upstream-service-time: + - '1101' + x-kong-proxy-latency: + - '14' + x-kong-request-id: + - 019df435-c391-7b83-a2d1-5055d23b514f + x-kong-upstream-latency: + - '1102' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_tool_spans.yaml b/py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_tool_spans.yaml new file mode 100644 index 00000000..bb35517f --- /dev/null +++ b/py/src/braintrust/integrations/mistral/cassettes/1.12.4/test_wrap_mistral_beta_conversations_start_tool_spans.yaml @@ -0,0 +1,86 @@ +interactions: +- request: + body: '{"inputs":"Use the code interpreter to calculate 21 * 2. Return only the + result.","stream":false,"tools":[{"type":"code_interpreter"}],"model":"mistral-small-latest"}' + headers: + Accept: + - application/json + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '166' + Host: + - api.mistral.ai + content-type: + - application/json + user-agent: + - mistral-client-python/1.12.4 + method: POST + uri: https://api.mistral.ai/v1/conversations + response: + body: + string: '{"object":"conversation.response","conversation_id":"conv_019df4349547763ab5a0acfe1591236d","outputs":[{"object":"entry","type":"tool.execution","created_at":"2026-05-04T18:16:20.369317Z","completed_at":"2026-05-04T18:16:21.125072Z","model":"mistral-small-latest","id":"tool_exec_019df43496917039a5066b144b0542b6","name":"code_interpreter","arguments":"{\"code\": + \"21 * 2\"}","function":"code_interpreter","info":{"code":"21 * 2","code_output":"42\n"}},{"object":"entry","type":"message.output","created_at":"2026-05-04T18:16:21.242155Z","completed_at":"2026-05-04T18:16:21.272804Z","model":"mistral-small-latest","id":"msg_019df43499fa767a8fb220270fd7ce5b","role":"assistant","content":"42"}],"usage":{"prompt_tokens":86,"completion_tokens":19,"total_tokens":110,"connector_tokens":5,"connectors":{"code_interpreter":1}},"guardrails":null}' + headers: + CF-RAY: + - 9f697ef48c7849fc-SJC + Connection: + - keep-alive + Content-Type: + - application/json + Date: + - Mon, 04 May 2026 18:16:21 GMT + Server: + - cloudflare + Strict-Transport-Security: + - max-age=15552000; includeSubDomains; preload + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + content-length: + - '841' + mistral-correlation-id: + - 019df434-952c-75b8-a13c-105b7be69de3 + set-cookie: + - __cf_bm=srYdq.iW.XvPJCfCeOzjZi.Zj_lMoQjIq_WJsNR9CuE-1777918579.925641-1.0.1.1-tzEi4sdxAdz2irm_6.4lxf1GK.l8x4oyTxQEgn7FYZiSIUkpToQU.ZlVanS4aqidM4ENwJ1kO1d0ZLh3xj1jp.Pa.c2kaZCka52mZkBySQBH1u6_AETzHR.A0EAWJThk; + HttpOnly; Secure; Path=/; Domain=mistral.ai; Expires=Mon, 04 May 2026 18:46:21 + GMT + - _cfuvid=Bs59mRiaoolnsjmFFfsJgMjpGWlBsQqzeA4B8Hpaztg-1777918579.925641-1.0.1.1-o3wi.3vNmTIzNBz5H38Uyb5lDmnUlSYDin8o1i1BWmM; + HttpOnly; SameSite=None; Secure; Path=/; Domain=mistral.ai + x-envoy-upstream-service-time: + - '1275' + x-kong-proxy-latency: + - '18' + x-kong-request-id: + - 019df434-952c-75b8-a13c-105b7be69de3 + x-kong-upstream-latency: + - '1277' + x-ratelimit-limit-code-interpreter-day: + - '10' + x-ratelimit-limit-code-interpreter-minute: + - '3' + x-ratelimit-limit-tokens-minute: + - '500000' + x-ratelimit-limit-tokens-month: + - '1000000000' + x-ratelimit-remaining-code-interpreter-day: + - '9' + x-ratelimit-remaining-code-interpreter-minute: + - '2' + x-ratelimit-remaining-tokens-minute: + - '499890' + x-ratelimit-remaining-tokens-month: + - '999999392' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml b/py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml new file mode 100644 index 00000000..414affb6 --- /dev/null +++ b/py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_stream_tool_spans.yaml @@ -0,0 +1,115 @@ +interactions: +- request: + body: '{"inputs":"Use the code interpreter to calculate 21 * 2. Return only the + result.","stream":true,"tools":[{"type":"code_interpreter"}],"model":"mistral-small-latest"}' + headers: + Accept: + - text/event-stream + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '165' + Host: + - api.mistral.ai + content-type: + - application/json + user-agent: + - mistral-client-python/2.4.4 + method: POST + uri: https://api.mistral.ai/v1/conversations + response: + body: + string: 'event: conversation.response.started + + data: {"type":"conversation.response.started","created_at":"2026-05-04T18:17:31.868571Z","conversation_id":"conv_019df435acbe72fe8159393921079b2b"} + + + event: tool.execution.started + + data: {"type":"tool.execution.started","created_at":"2026-05-04T18:17:31.872991Z","output_index":0,"id":"tool_exec_019df435ade0730c8aea545a915fb65e","model":"mistral-small-latest","name":"code_interpreter","arguments":"","function":"code_interpreter"} + + + event: tool.execution.delta + + data: {"type":"tool.execution.delta","created_at":"2026-05-04T18:17:31.877128Z","output_index":0,"id":"tool_exec_019df435ade0730c8aea545a915fb65e","name":"code_interpreter","arguments":"{\"code\": + \"","function":"code_interpreter"} + + + event: tool.execution.delta + + data: {"type":"tool.execution.delta","created_at":"2026-05-04T18:17:31.892107Z","output_index":0,"id":"tool_exec_019df435ade0730c8aea545a915fb65e","name":"code_interpreter","arguments":"21 + * ","function":"code_interpreter"} + + + event: tool.execution.delta + + data: {"type":"tool.execution.delta","created_at":"2026-05-04T18:17:31.910950Z","output_index":0,"id":"tool_exec_019df435ade0730c8aea545a915fb65e","name":"code_interpreter","arguments":"2\"}","function":"code_interpreter"} + + + event: tool.execution.done + + data: {"type":"tool.execution.done","created_at":"2026-05-04T18:17:32.570027Z","output_index":0,"id":"tool_exec_019df435ade0730c8aea545a915fb65e","name":"code_interpreter","function":"code_interpreter","info":{"code":"21 + * 2","code_output":"42\n"}} + + + event: message.output.delta + + data: {"type":"message.output.delta","created_at":"2026-05-04T18:17:32.707363Z","output_index":1,"id":"msg_019df435b12372eabb016e1e2f5d3d97","content_index":0,"model":"mistral-small-latest","role":"assistant","content":"4"} + + + event: message.output.delta + + data: {"type":"message.output.delta","created_at":"2026-05-04T18:17:32.721768Z","output_index":1,"id":"msg_019df435b12372eabb016e1e2f5d3d97","content_index":0,"model":"mistral-small-latest","role":"assistant","content":"2"} + + + event: conversation.response.done + + data: {"type":"conversation.response.done","created_at":"2026-05-04T18:17:32.748780Z","usage":{"prompt_tokens":86,"completion_tokens":19,"total_tokens":110,"connector_tokens":5,"connectors":{"code_interpreter":1}}} + + + ' + headers: + CF-RAY: + - 9f6980b3a8d3d59a-SJC + Connection: + - keep-alive + Content-Type: + - text/event-stream; charset=utf-8 + Date: + - Mon, 04 May 2026 18:17:31 GMT + Server: + - cloudflare + Strict-Transport-Security: + - max-age=15552000; includeSubDomains; preload + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + mistral-correlation-id: + - 019df435-aca5-76ef-af89-b68d546fb35c + set-cookie: + - __cf_bm=MiLFLA_gRlztjcWsGqmI3aQqgrwWwIBCZOWsIkRfSEs-1777918651.470252-1.0.1.1-Xhd7DGqnK0C5DI9xz_5WdP6KBerfk.TGvtYM.FCrA098iP3Sq5xf.gZqeP1ZLGo3ovGnxMrfUkuUZqjoD5a2PSSeLmJXETg4G65ng9Mz5lRZg1QvRjK.2Z.Uh8ubZp0Q; + HttpOnly; Secure; Path=/; Domain=mistral.ai; Expires=Mon, 04 May 2026 18:47:31 + GMT + - _cfuvid=UmqrIl_V_PrtzRN4zlVFwwLc7ynKa1oHPQqEjo1ytj0-1777918651.470252-1.0.1.1-Qc_29.Z00umjFznSjum_jqFURSIenwOKOGFMm.g7KnY; + HttpOnly; SameSite=None; Secure; Path=/; Domain=mistral.ai + x-envoy-upstream-service-time: + - '300' + x-kong-proxy-latency: + - '16' + x-kong-request-id: + - 019df435-aca5-76ef-af89-b68d546fb35c + x-kong-upstream-latency: + - '301' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_tool_spans.yaml b/py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_tool_spans.yaml new file mode 100644 index 00000000..59b3ae5d --- /dev/null +++ b/py/src/braintrust/integrations/mistral/cassettes/latest/test_wrap_mistral_beta_conversations_start_tool_spans.yaml @@ -0,0 +1,86 @@ +interactions: +- request: + body: '{"inputs":"Use the code interpreter to calculate 21 * 2. Return only the + result.","stream":false,"tools":[{"type":"code_interpreter"}],"model":"mistral-small-latest"}' + headers: + Accept: + - application/json + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '166' + Host: + - api.mistral.ai + content-type: + - application/json + user-agent: + - mistral-client-python/2.4.4 + method: POST + uri: https://api.mistral.ai/v1/conversations + response: + body: + string: '{"object":"conversation.response","conversation_id":"conv_019df434dbef71ba8344e8a71add58b7","outputs":[{"object":"entry","type":"tool.execution","created_at":"2026-05-04T18:16:38.455267Z","completed_at":"2026-05-04T18:16:39.337889Z","model":"mistral-small-latest","id":"tool_exec_019df434dd377608b47712bc02191110","name":"code_interpreter","arguments":"{\"code\": + \"21 * 2\"}","function":"code_interpreter","info":{"code":"21 * 2","code_output":"42\n"}},{"object":"entry","type":"message.output","created_at":"2026-05-04T18:16:39.462982Z","completed_at":"2026-05-04T18:16:39.476221Z","model":"mistral-small-latest","id":"msg_019df434e12674ec86dee2ea6a8207f0","role":"assistant","content":"42"}],"usage":{"prompt_tokens":86,"completion_tokens":19,"total_tokens":110,"connector_tokens":5,"connectors":{"code_interpreter":1}},"guardrails":null}' + headers: + CF-RAY: + - 9f697f658a91a364-SJC + Connection: + - keep-alive + Content-Type: + - application/json + Date: + - Mon, 04 May 2026 18:16:39 GMT + Server: + - cloudflare + Strict-Transport-Security: + - max-age=15552000; includeSubDomains; preload + Transfer-Encoding: + - chunked + X-Content-Type-Options: + - nosniff + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + content-length: + - '841' + mistral-correlation-id: + - 019df434-dbca-7275-8ff0-e0f7ee4f6859 + set-cookie: + - __cf_bm=UUT3oY.ENHmpOKbOiSt1qbFegNcd40otXOV5ewy5czM-1777918598.0101023-1.0.1.1-ATZXEMeDb3ITMcN105oJ7zsNqmEmjexVSlRhie8CMsWBqF.Kp4_Oc8u1DcX.wnzFBpj.I1gfHLrfdxuI_xGbtV7RY3bQop8aJyxyvb8MlTntbBAzDIud6t3qiqovRIPP; + HttpOnly; Secure; Path=/; Domain=mistral.ai; Expires=Mon, 04 May 2026 18:46:39 + GMT + - _cfuvid=xDWB4FcKcwbfb6NZ8Zo01zCm8URPF2nhydB4_Sbd1.Y-1777918598.0101023-1.0.1.1-Ah.tVR7Urd7WkxdpeVZOLAQodRr6V.wQu30T.LcYyhI; + HttpOnly; SameSite=None; Secure; Path=/; Domain=mistral.ai + x-envoy-upstream-service-time: + - '1396' + x-kong-proxy-latency: + - '28' + x-kong-request-id: + - 019df434-dbca-7275-8ff0-e0f7ee4f6859 + x-kong-upstream-latency: + - '1397' + x-ratelimit-limit-code-interpreter-day: + - '10' + x-ratelimit-limit-code-interpreter-minute: + - '3' + x-ratelimit-limit-tokens-minute: + - '500000' + x-ratelimit-limit-tokens-month: + - '1000000000' + x-ratelimit-remaining-code-interpreter-day: + - '8' + x-ratelimit-remaining-code-interpreter-minute: + - '1' + x-ratelimit-remaining-tokens-minute: + - '499780' + x-ratelimit-remaining-tokens-month: + - '999999282' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/mistral/test_mistral.py b/py/src/braintrust/integrations/mistral/test_mistral.py index 37b7161f..b38a47c3 100644 --- a/py/src/braintrust/integrations/mistral/test_mistral.py +++ b/py/src/braintrust/integrations/mistral/test_mistral.py @@ -116,6 +116,13 @@ def _assert_speech_complete_span(span, start, end): assert span["metrics"]["duration"] >= 0 +def _assert_conversation_tool_span(span): + assert span["span_attributes"]["type"] == SpanTypeAttribute.TOOL + assert span["input"] in ({"code": "21 * 2"}, {"code": "print(21 * 2)"}) + assert span["output"].get("stdout") == "42\n" or span["output"].get("code_output") == "42\n" + assert span["metadata"]["tool_type"] == "tool.execution" + + def _assert_conversation_span(span, expected_input, start, end, *, expected_content, stream=False): assert span["input"] == expected_input assert span["span_attributes"]["type"] == SpanTypeAttribute.TASK @@ -372,6 +379,51 @@ def test_wrap_mistral_beta_conversations_start_sync(memory_logger): ) +@pytest.mark.vcr +def test_wrap_mistral_beta_conversations_start_tool_spans(memory_logger): + assert not memory_logger.pop() + + client = wrap_mistral(_get_client()) + response = client.beta.conversations.start( + model=CHAT_MODEL, + inputs="Use the code interpreter to calculate 21 * 2. Return only the result.", + tools=[{"type": "code_interpreter"}], + ) + + assert response.outputs + assert any(output.type == "tool.execution" for output in response.outputs) + + spans = memory_logger.pop() + task_spans = find_spans_by_type(spans, SpanTypeAttribute.TASK) + tool_spans = find_spans_by_type(spans, SpanTypeAttribute.TOOL) + assert len(task_spans) == 1 + assert len(tool_spans) == 1 + _assert_conversation_tool_span(tool_spans[0]) + + +@pytest.mark.vcr +def test_wrap_mistral_beta_conversations_start_stream_tool_spans(memory_logger): + assert not memory_logger.pop() + + client = wrap_mistral(_get_client()) + with client.beta.conversations.start_stream( + model=CHAT_MODEL, + inputs="Use the code interpreter to calculate 21 * 2. Return only the result.", + tools=[{"type": "code_interpreter"}], + ) as stream: + events = list(stream) + + assert events + assert any(getattr(event, "event", None) == "tool.execution.done" for event in events) + + spans = memory_logger.pop() + task_spans = find_spans_by_type(spans, SpanTypeAttribute.TASK) + tool_spans = find_spans_by_type(spans, SpanTypeAttribute.TOOL) + assert len(task_spans) == 1 + assert len(tool_spans) == 1 + _assert_conversation_tool_span(tool_spans[0]) + + @pytest.mark.vcr @pytest.mark.asyncio async def test_wrap_mistral_beta_conversations_append_async(memory_logger): diff --git a/py/src/braintrust/integrations/mistral/tracing.py b/py/src/braintrust/integrations/mistral/tracing.py index 6928170a..30cfb20c 100644 --- a/py/src/braintrust/integrations/mistral/tracing.py +++ b/py/src/braintrust/integrations/mistral/tracing.py @@ -987,6 +987,27 @@ def _completion_tool_calls(response_data: dict[str, Any] | None) -> list[tuple[i return tool_calls +def _start_child_tool_span( + *, + parent_export: Any, + name: str, + tool_input: Any, + output: Any = None, + metadata: dict[str, Any] | None = None, +) -> None: + span_args = { + "name": f"tool: {name}", + "type": SpanTypeAttribute.TOOL, + "input": tool_input, + "output": output, + "metadata": clean_nones(metadata or {}) or None, + } + if parent_export is not None: + span_args["parent"] = parent_export + with start_span(**span_args): + pass + + def _log_completion_tool_spans(response_data: dict[str, Any] | None, *, parent_span: Any) -> None: tool_calls = _completion_tool_calls(response_data) if not tool_calls: @@ -997,24 +1018,53 @@ def _log_completion_tool_spans(response_data: dict[str, Any] | None, *, parent_s function = tool_call.get("function") if isinstance(tool_call.get("function"), dict) else {} tool_type = tool_call.get("type") or ("function" if function else None) name = function.get("name") or tool_type or "tool" - span_args = { - "name": f"tool: {name}", - "type": SpanTypeAttribute.TOOL, - "input": _maybe_parse_tool_arguments(function.get("arguments")), - "metadata": clean_nones( - { - "tool_call_id": tool_call.get("id"), - "tool_type": tool_type, - "tool_index": tool_call.get("index", tool_index), - "choice_index": choice_index, - } - ) - or None, - } - if parent_export is not None: - span_args["parent"] = parent_export - with start_span(**span_args): - pass + _start_child_tool_span( + parent_export=parent_export, + name=name, + tool_input=_maybe_parse_tool_arguments(function.get("arguments")), + metadata={ + "tool_call_id": tool_call.get("id"), + "tool_type": tool_type, + "tool_index": tool_call.get("index", tool_index), + "choice_index": choice_index, + }, + ) + + +def _conversation_tool_outputs(response_data: dict[str, Any] | None) -> list[tuple[int, dict[str, Any]]]: + tool_outputs = [] + for output_index, output in enumerate(_conversation_outputs_data(response_data)): + if not isinstance(output, dict): + continue + output_type = output.get("type") + if output_type in {"tool.execution", "tool_execution"}: + tool_outputs.append((output_index, output)) + return tool_outputs + + +def _log_conversation_tool_spans(response_data: dict[str, Any] | None, *, parent_span: Any) -> None: + tool_outputs = _conversation_tool_outputs(response_data) + if not tool_outputs: + return + + parent_export = parent_span.export() if parent_span is not None else None + for output_index, output in tool_outputs: + tool_type = output.get("type") + name = output.get("name") or tool_type or "tool" + tool_output = output.get("info") if "info" in output else output.get("output") + _start_child_tool_span( + parent_export=parent_export, + name=name, + tool_input=_maybe_parse_tool_arguments(output.get("arguments")), + output=tool_output, + metadata={ + "tool_call_id": output.get("id") or output.get("tool_call_id"), + "tool_type": tool_type, + "tool_index": output_index, + "agent_id": output.get("agent_id"), + "model": output.get("model"), + }, + ) def _finalize_completion_response(span: Any, request_metadata: dict[str, Any], response: Any, start_time: float): @@ -1035,6 +1085,7 @@ def _finalize_conversation_response(span: Any, request_metadata: dict[str, Any], response_data = _normalized_mistral_dict(response) response_metadata = _conversation_response_data_to_metadata(response_data) usage = response_data.get("usage") if response_data else None + _log_conversation_tool_spans(response_data, parent_span=span) _log_and_end_span( span, output=_conversation_outputs_data(response_data), @@ -1112,6 +1163,7 @@ def _finalize_conversation_stream( ): response = _aggregate_conversation_events(items) response_metadata = _conversation_response_data_to_metadata(response) + _log_conversation_tool_spans(response, parent_span=span) _log_and_end_span( span, output=response.get("outputs"),