feat(typesafe): Record token usage (#7719)

Record `gen_ai.usage.input_tokens` and `gen_ai.usage.output_tokens` attributes.

Author: alexander-alderman-webb

sentry_sdk/integrations/typesafe.py

@@ -221,6 +221,16 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A
 
             set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model)
 
+            if response.usage.input_tokens is not None:
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, response.usage.input_tokens
+                )
+
+            if response.usage.output_tokens is not None:
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens
+                )
+
             return response
 
     return wrap_system_one
@@ -297,6 +307,16 @@ async def wrap_system_one_async(
 
             set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model)
 
+            if response.usage.input_tokens is not None:
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, response.usage.input_tokens
+                )
+
+            if response.usage.output_tokens is not None:
+                set_on_span(
+                    SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, response.usage.output_tokens
+                )
+
             return response
 
     return wrap_system_one_async

tests/integrations/typesafe/test_typesafe.py

@@ -122,6 +122,9 @@ def test_system_one(
 
         assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest"
 
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 12
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3
+
         assert json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [
             {
                 "type": "evaluation",
@@ -205,6 +208,9 @@ def test_system_one(
 
         assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest"
 
+        assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 12
+        assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3
+
         assert json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [
             {
                 "type": "evaluation",
@@ -318,6 +324,9 @@ async def test_system_one_async(
 
         assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest"
 
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 12
+        assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3
+
         assert json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [
             {
                 "type": "evaluation",
@@ -404,6 +413,9 @@ async def test_system_one_async(
 
         assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "jev-latest"
 
+        assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 12
+        assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 3
+
         assert json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == [
             {
                 "type": "evaluation",