updated tests

Author: antonpirker

sentry_sdk/ai/monitoring.py

@@ -128,5 +128,8 @@ def record_token_usage(
             output_tokens_reasoning,
         )
 
+    if total_tokens is None and input_tokens is not None and output_tokens is not None:
+        total_tokens = input_tokens + output_tokens
+
     if total_tokens is not None:
         span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)

sentry_sdk/integrations/openai.py

@@ -366,22 +366,26 @@ def _new_embeddings_create_common(f, *args, **kwargs):
 
         response = yield f, args, kwargs
 
-        prompt_tokens = 0
+        input_tokens = 0
         total_tokens = 0
         if hasattr(response, "usage"):
             if hasattr(response.usage, "prompt_tokens") and isinstance(
                 response.usage.prompt_tokens, int
             ):
-                prompt_tokens = response.usage.prompt_tokens
+                input_tokens = response.usage.prompt_tokens
             if hasattr(response.usage, "total_tokens") and isinstance(
                 response.usage.total_tokens, int
             ):
                 total_tokens = response.usage.total_tokens
 
-        if prompt_tokens == 0:
-            prompt_tokens = integration.count_tokens(kwargs["input"] or "")
+        if input_tokens == 0:
+            input_tokens = integration.count_tokens(kwargs["input"] or "")
 
-        record_token_usage(span, prompt_tokens, None, total_tokens or prompt_tokens)
+        record_token_usage(
+            span,
+            input_tokens=input_tokens,
+            total_tokens=total_tokens or input_tokens,
+        )
 
         return response
 

tests/integrations/openai/test_openai.py

@@ -90,9 +90,9 @@ def test_nonstreaming_chat_completion(
         assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
         assert SPANDATA.AI_RESPONSES not in span["data"]
 
-    assert span["measurements"]["ai_completion_tokens_used"]["value"] == 10
-    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
-    assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
+    assert span["data"]["gen_ai.usage.output_tokens"] == 10
+    assert span["data"]["gen_ai.usage.input_tokens"] == 20
+    assert span["data"]["gen_ai.usage.total_tokens"] == 30
 
 
 @pytest.mark.asyncio
@@ -132,9 +132,9 @@ async def test_nonstreaming_chat_completion_async(
         assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
         assert SPANDATA.AI_RESPONSES not in span["data"]
 
-    assert span["measurements"]["ai_completion_tokens_used"]["value"] == 10
-    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
-    assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
+    assert span["data"]["gen_ai.usage.output_tokens"] == 10
+    assert span["data"]["gen_ai.usage.input_tokens"] == 20
+    assert span["data"]["gen_ai.usage.total_tokens"] == 30
 
 
 def tiktoken_encoding_if_installed():
@@ -228,9 +228,9 @@ def test_streaming_chat_completion(
     try:
         import tiktoken  # type: ignore # noqa # pylint: disable=unused-import
 
-        assert span["measurements"]["ai_completion_tokens_used"]["value"] == 2
-        assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 1
-        assert span["measurements"]["ai_total_tokens_used"]["value"] == 3
+        assert span["data"]["gen_ai.usage.output_tokens"] == 2
+        assert span["data"]["gen_ai.usage.input_tokens"] == 1
+        assert span["data"]["gen_ai.usage.total_tokens"] == 3
     except ImportError:
         pass  # if tiktoken is not installed, we can't guarantee token usage will be calculated properly
 
@@ -324,9 +324,9 @@ async def test_streaming_chat_completion_async(
     try:
         import tiktoken  # type: ignore # noqa # pylint: disable=unused-import
 
-        assert span["measurements"]["ai_completion_tokens_used"]["value"] == 2
-        assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 1
-        assert span["measurements"]["ai_total_tokens_used"]["value"] == 3
+        assert span["data"]["gen_ai.usage.output_tokens"] == 2
+        assert span["data"]["gen_ai.usage.input_tokens"] == 1
+        assert span["data"]["gen_ai.usage.total_tokens"] == 3
     except ImportError:
         pass  # if tiktoken is not installed, we can't guarantee token usage will be calculated properly
 
@@ -409,8 +409,8 @@ def test_embeddings_create(
     else:
         assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
 
-    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
-    assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
+    assert span["data"]["gen_ai.usage.input_tokens"] == 20
+    assert span["data"]["gen_ai.usage.total_tokens"] == 30
 
 
 @pytest.mark.asyncio
@@ -457,8 +457,8 @@ async def test_embeddings_create_async(
     else:
         assert SPANDATA.AI_INPUT_MESSAGES not in span["data"]
 
-    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
-    assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
+    assert span["data"]["gen_ai.usage.input_tokens"] == 20
+    assert span["data"]["gen_ai.usage.total_tokens"] == 30
 
 
 @pytest.mark.parametrize(