feat(mistral): Record token usage

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

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