From 2548a48884d4ae2fb0ce50b75d71cff3d60efb33 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Tue, 21 Jul 2026 00:18:07 +0000 Subject: [PATCH] fix(mistral): drop eager response serialization, use allowlist metadata pass-through MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Align the Mistral integration with `.agents/skills/sdk-integrations/SKILL.md` following the same pattern as #594 (litellm) and #583 (anthropic): - **Response finalizers** no longer call `_normalized_mistral_dict(response)` to full-`model_dump` the response just to read a few fields. Read `choices` / `outputs` / `pages` / `usage` / `id` / `model` directly via the existing `_get_value` helper (dict-or-pydantic aware, handles Mistral's `Unset` sentinel). `bt_safe_deep_copy` in `logger.log_internal` already `model_dump`s pydantic values at log time. - **`_build_request_metadata`** stops walk-and-dumping every metadata value (`temperature`, `top_p`, `tools`, `response_format`, `stop`, …). These don't carry attachments; pass through and let bt_json handle it. - **Streaming aggregators** stop pre-dumping `usage` / `segments` into the intermediate dict; `_parse_usage_metrics` already accepts pydantic and the log-time deep-copy handles the output. - **Tool-span walkers** (`_completion_tool_calls`, `_conversation_tool_outputs`, `_log_completion_tool_spans`, `_log_conversation_tool_spans`) switched to `_get_value` so they work on raw pydantic responses. The multimodal walker is retained where attachment materialization matters: the `_start_span` input path, `_append_delta_content` / `_merge_tool_calls` streaming assembly, and per-field accumulators for conversation outputs. Metadata is still allowlist-per-surface — no changes there. Test updates (no new mocks, no cassette re-recording): - `test_wrappers_ignore_usage_when_response_normalization_fails` → `test_wrappers_read_usage_from_plain_python_responses`. The old test asserted metrics were dropped for plain-Python responses — that was a side-effect of `.model_dump()` failing on non-pydantic objects. The new behavior correctly reads usage via `getattr`; assert the positive. - `test_aggregate_completion_events_merges_tool_calls_and_content`: `aggregated["usage"]` is now the raw pydantic `UsageInfo` (dumps at log time, not in the aggregator); switched to `getattr(..., "total_tokens")`. Co-Authored-By: Claude Opus 4.7 --- .../integrations/mistral/test_mistral.py | 20 +- .../integrations/mistral/tracing.py | 172 +++++++----------- 2 files changed, 78 insertions(+), 114 deletions(-) diff --git a/py/src/braintrust/integrations/mistral/test_mistral.py b/py/src/braintrust/integrations/mistral/test_mistral.py index e8a9f41f..04fb0ddb 100644 --- a/py/src/braintrust/integrations/mistral/test_mistral.py +++ b/py/src/braintrust/integrations/mistral/test_mistral.py @@ -1339,14 +1339,14 @@ def test_normalize_mistral_multimodal_value_leaves_non_base64_input_audio_unchan assert sanitized == original -class _UnsupportedUsageResponse: +class _PlainResponse: def __init__(self, **attrs): for key, value in attrs.items(): setattr(self, key, value) @pytest.mark.parametrize( - ("wrapper", "kwargs", "response", "missing_metric_keys"), + ("wrapper", "kwargs", "response", "expected_metric_keys"), [ ( _chat_complete_wrapper, @@ -1354,7 +1354,7 @@ def __init__(self, **attrs): "model": CHAT_MODEL, "messages": [{"role": "user", "content": "hello"}], }, - _UnsupportedUsageResponse(usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}), + _PlainResponse(usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}), {"tokens", "prompt_tokens", "completion_tokens"}, ), ( @@ -1363,7 +1363,7 @@ def __init__(self, **attrs): "model": CHAT_MODEL, "inputs": [{"role": "user", "content": "hello"}], }, - _UnsupportedUsageResponse(usage={"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5}), + _PlainResponse(usage={"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5}), {"tokens", "prompt_tokens", "completion_tokens"}, ), ( @@ -1372,13 +1372,13 @@ def __init__(self, **attrs): "model": OCR_MODEL, "document": {"type": "document_url", "document_url": TEST_PDF_DATA_URL}, }, - _UnsupportedUsageResponse(usage_info={"pages_processed": 1, "doc_size_bytes": 70}), + _PlainResponse(usage_info={"pages_processed": 1, "doc_size_bytes": 70}), {"pages_processed", "doc_size_bytes"}, ), ], ) -def test_wrappers_ignore_usage_when_response_normalization_fails( - memory_logger, wrapper, kwargs, response, missing_metric_keys +def test_wrappers_read_usage_from_plain_python_responses( + memory_logger, wrapper, kwargs, response, expected_metric_keys ): assert not memory_logger.pop() @@ -1389,8 +1389,8 @@ def test_wrappers_ignore_usage_when_response_normalization_fails( spans = memory_logger.pop() assert len(spans) == 1 span = spans[0] - for key in missing_metric_keys: - assert key not in span["metrics"] + for key in expected_metric_keys: + assert key in span["metrics"] assert span["metrics"]["duration"] >= 0 @@ -1447,7 +1447,7 @@ def test_aggregate_completion_events_merges_tool_calls_and_content(): assert aggregated["id"] == "cmpl_123" assert aggregated["model"] == CHAT_MODEL - assert aggregated["usage"]["total_tokens"] == 14 + assert getattr(aggregated["usage"], "total_tokens", None) == 14 assert aggregated["choices"][0]["finish_reason"] == "tool_calls" tool_call = aggregated["choices"][0]["message"]["tool_calls"][0] assert tool_call["id"] == "call_1" diff --git a/py/src/braintrust/integrations/mistral/tracing.py b/py/src/braintrust/integrations/mistral/tracing.py index 7a1ed296..d4a6d082 100644 --- a/py/src/braintrust/integrations/mistral/tracing.py +++ b/py/src/braintrust/integrations/mistral/tracing.py @@ -261,11 +261,11 @@ def _build_request_metadata( value = kwargs.get(key) if value is None or _is_unset(value): continue - metadata[key] = _normalize_mistral_multimodal_value(value) + metadata[key] = value request_metadata = kwargs.get("metadata") if request_metadata is not None and not _is_unset(request_metadata): - metadata["request_metadata"] = _normalize_mistral_multimodal_value(request_metadata) + metadata["request_metadata"] = request_metadata if stream is not None: metadata["stream"] = stream @@ -452,47 +452,38 @@ def _merge_ocr_metrics(start_time: float, usage: Any) -> dict[str, Any]: return _merge_timing_and_usage_metrics(start_time, usage, _parse_ocr_usage_metrics) -def _response_data_to_metadata(data: dict[str, Any] | None) -> dict[str, Any]: - if data is None: +def _response_to_metadata(response: Any) -> dict[str, Any]: + if response is None: return {} metadata = {} for key in ("id", "model", "object", "created"): - value = data.get(key) + value = _get_value(response, key) if value is not None: metadata[key] = value return metadata -def _response_to_metadata(response: Any) -> dict[str, Any]: - return _response_data_to_metadata(_normalized_mistral_dict(response)) - - -def _conversation_outputs_data(data: dict[str, Any] | None) -> list[Any]: - if data is None: - return [] - - outputs = data.get("outputs") +def _conversation_outputs(response: Any) -> list[Any]: + outputs = _get_value(response, "outputs") return outputs if isinstance(outputs, list) else [] -def _conversation_response_data_to_metadata(data: dict[str, Any] | None) -> dict[str, Any]: - if data is None: +def _conversation_response_metadata(response: Any) -> dict[str, Any]: + if response is None: return {} metadata = {} - conversation_id = data.get("conversation_id") + conversation_id = _get_value(response, "conversation_id") if conversation_id is not None: metadata["conversation_id"] = conversation_id - object_type = data.get("object") + object_type = _get_value(response, "object") if object_type is not None: metadata["object"] = object_type - for output in _conversation_outputs_data(data): - if not isinstance(output, dict): - continue - model = output.get("model") + for output in _conversation_outputs(response): + model = _get_value(output, "model") if model is not None: metadata["model"] = model break @@ -514,19 +505,6 @@ def _embeddings_output(response: Any) -> dict[str, Any]: return output -def _ocr_output_data(data: dict[str, Any] | None) -> dict[str, Any]: - if data is None: - return {"pages": []} - - output = { - "pages": data.get("pages") or [], - } - document_annotation = data.get("document_annotation") - if document_annotation is not None: - output["document_annotation"] = document_annotation - return output - - def _transcription_response_to_metadata(response: Any) -> dict[str, Any]: metadata = _response_to_metadata(response) language = _get_value(response, "language") @@ -719,11 +697,11 @@ def _aggregate_transcription_events(items: list[Any]) -> dict[str, Any]: if model is not None: result["model"] = model if usage is not None: - result["usage"] = _normalize_mistral_multimodal_value(usage) + result["usage"] = usage if language is not None: result["language"] = language if isinstance(segments, list): - result["segments"] = _normalize_mistral_multimodal_value(segments) + result["segments"] = segments if finish_reason is not None: result["finish_reason"] = finish_reason return result @@ -754,7 +732,7 @@ def _aggregate_speech_events(items: list[Any]) -> dict[str, Any]: result = {"audio_data": "".join(audio_parts)} if usage is not None: - result["usage"] = _normalize_mistral_multimodal_value(usage) + result["usage"] = usage return result @@ -814,7 +792,7 @@ def _aggregate_completion_events(items: list[Any]) -> dict[str, Any]: if created is not None: result["created"] = created if usage is not None: - result["usage"] = _normalize_mistral_multimodal_value(usage) + result["usage"] = usage return result @@ -962,7 +940,7 @@ def _aggregate_conversation_events(items: list[Any]) -> dict[str, Any]: if conversation_id is not None: result["conversation_id"] = conversation_id if usage is not None: - result["usage"] = _normalize_mistral_multimodal_value(usage) + result["usage"] = usage return result @@ -975,27 +953,21 @@ def _maybe_parse_tool_arguments(arguments: Any) -> Any: return arguments -def _completion_tool_calls(response_data: dict[str, Any] | None) -> list[tuple[int, int, dict[str, Any]]]: - if not response_data: - return [] - - tool_calls = [] - choices = response_data.get("choices") +def _completion_tool_calls(response: Any) -> list[tuple[int, int, Any]]: + choices = _get_value(response, "choices") if not isinstance(choices, list): return [] + tool_calls = [] for choice_index, choice in enumerate(choices): - if not isinstance(choice, dict): - continue - message = choice.get("message") - if not isinstance(message, dict): + message = _get_value(choice, "message") + if message is None: continue - choice_tool_calls = message.get("tool_calls") + choice_tool_calls = _get_value(message, "tool_calls") if not isinstance(choice_tool_calls, list): continue for tool_index, tool_call in enumerate(choice_tool_calls): - if isinstance(tool_call, dict): - tool_calls.append((choice_index, tool_index, tool_call)) + tool_calls.append((choice_index, tool_index, tool_call)) return tool_calls @@ -1020,111 +992,105 @@ def _start_child_tool_span( pass -def _log_completion_tool_spans(response_data: dict[str, Any] | None, *, parent_span: Any) -> None: - tool_calls = _completion_tool_calls(response_data) +def _log_completion_tool_spans(response: Any, *, parent_span: Any) -> None: + tool_calls = _completion_tool_calls(response) if not tool_calls: return parent_export = parent_span.export() if parent_span is not None else None for choice_index, tool_index, tool_call in tool_calls: - 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" + function = _get_value(tool_call, "function") + function_name = _get_value(function, "name") if function is not None else None + tool_type = _get_value(tool_call, "type") or ("function" if function is not None else None) + arguments = _get_value(function, "arguments") if function is not None else None + call_index = _get_value(tool_call, "index") _start_child_tool_span( parent_export=parent_export, - name=name, - tool_input=_maybe_parse_tool_arguments(function.get("arguments")), + name=function_name or tool_type or "tool", + tool_input=_maybe_parse_tool_arguments(arguments), metadata={ - "tool_call_id": tool_call.get("id"), + "tool_call_id": _get_value(tool_call, "id"), "tool_type": tool_type, - "tool_index": tool_call.get("index", tool_index), + "tool_index": call_index if isinstance(call_index, int) else tool_index, "choice_index": choice_index, }, ) -def _conversation_tool_outputs(response_data: dict[str, Any] | None) -> list[tuple[int, dict[str, Any]]]: +def _conversation_tool_outputs(response: Any) -> list[tuple[int, 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") + for output_index, output in enumerate(_conversation_outputs(response)): + output_type = _get_value(output, "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) +def _log_conversation_tool_spans(response: Any, *, parent_span: Any) -> None: + tool_outputs = _conversation_tool_outputs(response) 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") + tool_type = _get_value(output, "type") + info = _get_value(output, "info") + tool_output = info if info is not None else _get_value(output, "output") _start_child_tool_span( parent_export=parent_export, - name=name, - tool_input=_maybe_parse_tool_arguments(output.get("arguments")), + name=_get_value(output, "name") or tool_type or "tool", + tool_input=_maybe_parse_tool_arguments(_get_value(output, "arguments")), output=tool_output, metadata={ - "tool_call_id": output.get("id") or output.get("tool_call_id"), + "tool_call_id": _get_value(output, "id", "tool_call_id"), "tool_type": tool_type, "tool_index": output_index, - "agent_id": output.get("agent_id"), - "model": output.get("model"), + "agent_id": _get_value(output, "agent_id"), + "model": _get_value(output, "model"), }, ) def _finalize_completion_response(span: Any, request_metadata: dict[str, Any], response: Any, start_time: float): - response_data = _normalized_mistral_dict(response) - response_metadata = _response_data_to_metadata(response_data) - usage = response_data.get("usage") if response_data else None - - _log_completion_tool_spans(response_data, parent_span=span) + _log_completion_tool_spans(response, parent_span=span) _log_and_end_span( span, - output=response_data.get("choices") if response_data else None, - metrics=_merge_metrics(start_time, usage), - metadata={**request_metadata, **response_metadata}, + output=_get_value(response, "choices"), + metrics=_merge_metrics(start_time, _get_value(response, "usage")), + metadata={**request_metadata, **_response_to_metadata(response)}, ) def _finalize_conversation_response(span: Any, request_metadata: dict[str, Any], response: Any, start_time: float): - 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_conversation_tool_spans(response, parent_span=span) _log_and_end_span( span, - output=_conversation_outputs_data(response_data), - metrics=_merge_metrics(start_time, usage), - metadata={**request_metadata, **response_metadata}, + output=_conversation_outputs(response), + metrics=_merge_metrics(start_time, _get_value(response, "usage")), + metadata={**request_metadata, **_conversation_response_metadata(response)}, ) def _finalize_embeddings_response(span: Any, request_metadata: dict[str, Any], response: Any, start_time: float): - response_metadata = _response_to_metadata(response) _log_and_end_span( span, output=_embeddings_output(response), - metrics=_merge_metrics(start_time, getattr(response, "usage", None)), - metadata={**request_metadata, **response_metadata}, + metrics=_merge_metrics(start_time, _get_value(response, "usage")), + metadata={**request_metadata, **_response_to_metadata(response)}, ) def _finalize_ocr_response(span: Any, request_metadata: dict[str, Any], response: Any, start_time: float): - response_data = _normalized_mistral_dict(response) - pages = response_data.get("pages") if response_data else None + pages = _get_value(response, "pages") response_metadata = {"page_count": len(pages)} if isinstance(pages, list) else {} - usage_info = response_data.get("usage_info") if response_data else None + output: dict[str, Any] = {"pages": pages if isinstance(pages, list) else []} + document_annotation = _get_value(response, "document_annotation") + if document_annotation is not None: + output["document_annotation"] = document_annotation _log_and_end_span( span, - output=_ocr_output_data(response_data), - metrics=_merge_ocr_metrics(start_time, usage_info), + output=output, + metrics=_merge_ocr_metrics(start_time, _get_value(response, "usage_info")), metadata={**request_metadata, **response_metadata}, ) @@ -1156,13 +1122,12 @@ def _finalize_completion_stream( first_token_time: float | None, ): response = _aggregate_completion_events(items) - response_metadata = _response_data_to_metadata(response) _log_completion_tool_spans(response, parent_span=span) _log_and_end_span( span, output=response.get("choices"), metrics=_merge_metrics(start_time, response.get("usage"), first_token_time), - metadata={**metadata, **response_metadata}, + metadata={**metadata, **_response_to_metadata(response)}, ) @@ -1174,13 +1139,12 @@ def _finalize_conversation_stream( first_token_time: float | None, ): 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"), metrics=_merge_metrics(start_time, response.get("usage"), first_token_time), - metadata={**metadata, **response_metadata}, + metadata={**metadata, **_conversation_response_metadata(response)}, )