More tests

Author: antonpirker

sentry_sdk/tracing_utils.py

@@ -1094,18 +1094,31 @@ def _get_output_attributes(template, send_pii, result):
 
     if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
         attributes.update(_get_usage_attributes(result))
-        if hasattr(result, "usage"):
-            attributes.update(_get_usage_attributes(result.usage))
-        elif hasattr(result, "metadata"):
-            if hasattr(result.metadata, "usage"):
-                attributes.update(_get_usage_attributes(result.metadata.usage))
-        elif isinstance(result, dict) and "usage" in result:
-            attributes.update(_get_usage_attributes(result["usage"]))
 
-        elif hasattr(result, "model") and isinstance(result.model, str):
-            attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = result.model
-        elif hasattr(result, "model_name") and isinstance(result.model_name, str):
-            attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = result.model_name
+        if isinstance(result, dict):
+            if "usage" in result:
+                attributes.update(_get_usage_attributes(result["usage"]))
+            if (
+                "metadata" in result
+                and result["metadata"]
+                and "usage" in result["metadata"]
+            ):
+                attributes.update(_get_usage_attributes(result["metadata"]["usage"]))
+            if "model" in result:
+                attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = result["model"]
+            if "model_name" in result:
+                attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = result["model_name"]
+        else:
+            if hasattr(result, "usage"):
+                attributes.update(_get_usage_attributes(result.usage))
+            if hasattr(result, "metadata") and hasattr(result.metadata, "usage"):
+                attributes.update(_get_usage_attributes(result.metadata.usage))
+            if isinstance(result, dict) and "usage" in result:
+                attributes.update(_get_usage_attributes(result["usage"]))
+            if hasattr(result, "model") and isinstance(result.model, str):
+                attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = result.model
+            if hasattr(result, "model_name") and isinstance(result.model_name, str):
+                attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = result.model_name
 
     if template == SPANTEMPLATE.AI_TOOL:
         if send_pii:

tests/tracing/test_decorator.py

@@ -127,23 +127,43 @@ def test_span_templates(sentry_init, capture_events):
 
     @sentry_sdk.trace(template=SPANTEMPLATE.AI_TOOL)
     def my_tool(arg1, arg2):
-        return "my_tool_result"
+        mock_usage = mock.Mock()
+        mock_usage.input_tokens = 11
+        mock_usage.output_tokens = 22
+        mock_usage.total_tokens = 33
 
-    @sentry_sdk.trace(template=SPANTEMPLATE.AI_CHAT)
-    def my_chat():
         return {
-            "content": "my_chat_result",
-            "usage": {
-                "prompt_tokens": 10,
-                "completion_tokens": 20,
-                "total_tokens": 30,
-            },
+            "output": "my_tool_result",
+            "usage": mock_usage,
+        }
+
+    @sentry_sdk.trace(template=SPANTEMPLATE.AI_CHAT)
+    def my_chat(model=None, **kwargs):
+        mock_result = mock.Mock()
+        mock_result.content = "my_chat_result"
+        mock_result.usage = {
+            "prompt_tokens": 10,
+            "completion_tokens": 20,
+            "total_tokens": 30,
         }
+        mock_result.model = f"{model}-v123"
+
+        return mock_result
 
     @sentry_sdk.trace(template=SPANTEMPLATE.AI_AGENT)
     def my_agent():
         my_tool(1, 2)
-        my_chat()
+        my_chat(
+            model="my-gpt-4o-mini",
+            prompt="What is the weather in Tokyo?",
+            system_prompt="You are a helpful assistant that can answer questions about the weather.",
+            max_tokens=100,
+            temperature=0.5,
+            top_p=0.9,
+            top_k=40,
+            frequency_penalty=1.0,
+            presence_penalty=2.0,
+        )
 
     with sentry_sdk.start_transaction(name="test-transaction"):
         my_agent()
@@ -171,14 +191,26 @@ def my_agent():
     assert tool_span["data"] == {
         "gen_ai.tool.name": "test_decorator.test_span_templates.<locals>.my_tool",
         "gen_ai.operation.name": "execute_tool",
+        "gen_ai.usage.input_tokens": 11,
+        "gen_ai.usage.output_tokens": 22,
+        "gen_ai.usage.total_tokens": 33,
         "thread.id": mock.ANY,
         "thread.name": mock.ANY,
     }
 
     assert chat_span["op"] == "gen_ai.chat"
-    assert chat_span["description"] == "chat"
+    assert chat_span["description"] == "chat my-gpt-4o-mini"
     assert chat_span["data"] == {
         "gen_ai.operation.name": "chat",
+        "gen_ai.request.frequency_penalty": 1.0,
+        "gen_ai.request.max_tokens": 100,
+        "gen_ai.request.messages": "[{'role': 'user', 'content': 'What is the weather in Tokyo?'}, {'role': 'system', 'content': 'You are a helpful assistant that can answer questions about the weather.'}]",
+        "gen_ai.request.model": "my-gpt-4o-mini",
+        "gen_ai.request.presence_penalty": 2.0,
+        "gen_ai.request.temperature": 0.5,
+        "gen_ai.request.top_k": 40,
+        "gen_ai.request.top_p": 0.9,
+        "gen_ai.response.model": "my-gpt-4o-mini-v123",
         "gen_ai.usage.input_tokens": 10,
         "gen_ai.usage.output_tokens": 20,
         "gen_ai.usage.total_tokens": 30,