feat(typesafe): Set token usage attributes

Author: alexander-alderman-webb

sentry_sdk/integrations/typesafe.py

@@ -212,6 +212,16 @@ def wrap_system_one(self: "TypeSafeClient", *args: "Any", **kwargs: "Any") -> "A
             if not isinstance(response, SystemOneResponse):
                 return response
 
+            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
@@ -280,6 +290,16 @@ async def wrap_system_one_async(
             if not isinstance(response, SystemOneResponse):
                 return response
 
+            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

@@ -121,6 +121,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",
             "state": {
@@ -208,6 +211,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",
             "state": {
@@ -324,6 +330,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",
             "state": {
@@ -414,6 +423,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",
             "state": {