diff --git a/py/src/braintrust/integrations/mistral/test_mistral.py b/py/src/braintrust/integrations/mistral/test_mistral.py index 098ff3176..f3089ae7c 100644 --- a/py/src/braintrust/integrations/mistral/test_mistral.py +++ b/py/src/braintrust/integrations/mistral/test_mistral.py @@ -289,7 +289,14 @@ def test_wrap_mistral_chat_complete_tool_spans(memory_logger): assert response.choices[0].message.tool_calls[0].function.name == "get_weather" 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 + span = llm_spans[0] + service_tier = getattr(response.usage, "service_tier", None) + if service_tier is not None: + assert span["metadata"]["usage_service_tier"] == service_tier + else: + assert "usage_service_tier" not in span["metadata"] 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 c32305822..b8b2842ae 100644 --- a/py/src/braintrust/integrations/mistral/tracing.py +++ b/py/src/braintrust/integrations/mistral/tracing.py @@ -45,6 +45,7 @@ def start_span(*args, **kwargs): } _CHAT_METADATA_KEYS = ( "model", + "service_tier", "temperature", "top_p", "max_tokens", @@ -65,6 +66,7 @@ def start_span(*args, **kwargs): ) _AGENTS_METADATA_KEYS = ( "agent_id", + "service_tier", "max_tokens", "stop", "random_seed", @@ -472,6 +474,11 @@ def _response_to_metadata(response: Any) -> dict[str, Any]: value = _get_value(response, key) if value is not None: metadata[key] = value + + service_tier = _get_value(_get_value(response, "usage"), "service_tier") + if service_tier is not None: + metadata["usage_service_tier"] = service_tier + return metadata