fix(mistral): trace cached prompt tokens

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
 
 

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"]