diff --git a/sentry_sdk/integrations/mistral.py b/sentry_sdk/integrations/mistral.py index 8452c580c4..e6e4e417db 100644 --- a/sentry_sdk/integrations/mistral.py +++ b/sentry_sdk/integrations/mistral.py @@ -16,6 +16,7 @@ try: from mistralai.client.chat import Chat + from mistralai.client.models import ChatCompletionResponse except ImportError: raise DidNotEnable("mistralai not installed") @@ -69,7 +70,29 @@ def wrap_complete(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any": set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) - return f(self, *args, **kwargs) + + response = f(self, *args, **kwargs) + + if not isinstance(response, ChatCompletionResponse): + return response + + if response.usage.prompt_tokens is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, response.usage.prompt_tokens + ) + + if response.usage.completion_tokens is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + response.usage.completion_tokens, + ) + + if response.usage.total_tokens is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, response.usage.total_tokens + ) + + return response return wrap_complete @@ -112,6 +135,28 @@ async def wrap_complete_async(self: "Chat", *args: "Any", **kwargs: "Any") -> "A set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, False) - return await f(self, *args, **kwargs) + + response = await f(self, *args, **kwargs) + + if not isinstance(response, ChatCompletionResponse): + return response + + if response.usage.prompt_tokens is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, response.usage.prompt_tokens + ) + + if response.usage.completion_tokens is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, + response.usage.completion_tokens, + ) + + if response.usage.total_tokens is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, response.usage.total_tokens + ) + + return response return wrap_complete_async diff --git a/tests/integrations/mistral/test_mistral.py b/tests/integrations/mistral/test_mistral.py index 9719719f34..cc3c6e3c73 100644 --- a/tests/integrations/mistral/test_mistral.py +++ b/tests/integrations/mistral/test_mistral.py @@ -92,6 +92,10 @@ def test_nonstreaming_chat( span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "mistral-medium-latest" ) assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 else: items = capture_items("transaction") @@ -115,6 +119,10 @@ def test_nonstreaming_chat( assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "open-mistral" assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + @pytest.mark.asyncio @pytest.mark.parametrize("span_streaming", [True, False]) @@ -172,6 +180,10 @@ async def test_nonstreaming_chat_async( span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "mistral-medium-latest" ) assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 else: items = capture_items("transaction") @@ -194,3 +206,7 @@ async def test_nonstreaming_chat_async( assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "mistral-medium-latest" assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30