From 91982a2d8d864fb708cf7004549b1c7fa55d2272 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 12 Aug 2026 15:39:54 -0400 Subject: [PATCH] fix(mistral): trace cached prompt tokens Parse nested Mistral token detail fields into standardized prompt and completion metrics. Extend VCR-backed sync and streaming coverage for cached-token usage.\n\nFixes #678. --- .../integrations/mistral/test_mistral.py | 9 ++++++++- py/src/braintrust/integrations/mistral/tracing.py | 15 +++++++++++++-- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/py/src/braintrust/integrations/mistral/test_mistral.py b/py/src/braintrust/integrations/mistral/test_mistral.py index 04fb0ddb..298fffc4 100644 --- a/py/src/braintrust/integrations/mistral/test_mistral.py +++ b/py/src/braintrust/integrations/mistral/test_mistral.py @@ -242,6 +242,7 @@ def test_wrap_mistral_chat_complete_sync(memory_logger): end = time.time() assert "4" in str(response.choices[0].message.content) + assert response.usage.prompt_tokens_details["cached_tokens"] == 0 spans = memory_logger.pop() assert len(spans) == 1 @@ -251,6 +252,7 @@ def test_wrap_mistral_chat_complete_sync(memory_logger): assert span["metadata"]["provider"] == "mistral" assert span["metadata"]["model"] == CHAT_MODEL assert "4" in str(span["output"]) + assert span["metrics"]["prompt_cached_tokens"] == 0 assert_metrics_are_valid(span["metrics"], start, end) @@ -558,8 +560,13 @@ def test_wrap_mistral_agents_stream_tool_spans(memory_logger): chunks = list(stream) assert chunks + usage = next(chunk.data.usage for chunk in reversed(chunks) if getattr(chunk.data, "usage", None) is not None) + assert usage.prompt_tokens_details["cached_tokens"] == 80 + spans = memory_logger.pop() - assert len(find_spans_by_type(spans, SpanTypeAttribute.LLM)) == 1 + llm_spans = find_spans_by_type(spans, SpanTypeAttribute.LLM) + assert len(llm_spans) == 1 + assert llm_spans[0]["metrics"]["prompt_cached_tokens"] == 80 assert len(find_spans_by_type(spans, SpanTypeAttribute.TOOL)) == 1 diff --git a/py/src/braintrust/integrations/mistral/tracing.py b/py/src/braintrust/integrations/mistral/tracing.py index d4a6d082..c3230582 100644 --- a/py/src/braintrust/integrations/mistral/tracing.py +++ b/py/src/braintrust/integrations/mistral/tracing.py @@ -39,6 +39,10 @@ def start_span(*args, **kwargs): _TOKEN_NAME_MAP = { "total_tokens": "tokens", } +_TOKEN_DETAIL_PREFIX_MAP = { + "prompt_tokens_details": "prompt", + "completion_tokens_details": "completion", +} _CHAT_METADATA_KEYS = ( "model", "temperature", @@ -419,9 +423,16 @@ def _parse_usage_metrics(usage: Any) -> dict[str, float]: metrics = {} for key, value in usage_data.items(): - if not _is_supported_metric_value(value): + if _is_supported_metric_value(value): + metrics[_TOKEN_NAME_MAP.get(key, _camel_to_snake(key))] = float(value) + continue + + prefix = _TOKEN_DETAIL_PREFIX_MAP.get(key) + if prefix is None or not isinstance(value, dict): continue - metrics[_TOKEN_NAME_MAP.get(key, _camel_to_snake(key))] = float(value) + for nested_key, nested_value in value.items(): + if _is_supported_metric_value(nested_value): + metrics[f"{prefix}_{_camel_to_snake(nested_key)}"] = float(nested_value) if "tokens" not in metrics and "prompt_tokens" in metrics and "completion_tokens" in metrics: metrics["tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"]