fix(integrations): cleanups and working in feedback

constantinius · constantinius · commit 7f335f9a4efb · 2025-10-15T16:59:06.000+02:00
diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py
@@ -9,7 +9,6 @@
 
 try:
     import pydantic_ai
-
 except ImportError:
     raise DidNotEnable("pydantic-ai not installed")
 
diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py
@@ -1,21 +1,28 @@
 from functools import wraps
 
 import sentry_sdk
-from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.tracing_utils import set_span_errored
+from sentry_sdk.utils import event_from_exception
 
-from ..spans import agent_workflow_span, invoke_agent_span, update_invoke_agent_span
-from ..utils import _capture_exception
+from ..spans import agent_workflow_span
 
 from typing import TYPE_CHECKING
+from pydantic_ai.agent import Agent
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Optional
 
-try:
-    import pydantic_ai
-    from pydantic_ai.agent import Agent
-except ImportError:
-    raise DidNotEnable("pydantic-ai not installed")
+
+def _capture_exception(exc):
+    # type: (Any) -> None
+    set_span_errored()
+
+    event, hint = event_from_exception(
+        exc,
+        client_options=sentry_sdk.get_client().options,
+        mechanism={"type": "pydantic_ai", "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)
 
 
 class _StreamingContextManagerWrapper:
diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py
@@ -1,25 +1,18 @@
 from functools import wraps
 
-from sentry_sdk.integrations import DidNotEnable
-
 from ..spans import (
     invoke_agent_span,
     update_invoke_agent_span,
     ai_client_span,
     update_ai_client_span,
 )
+from pydantic_ai._agent_graph import UserPromptNode, ModelRequestNode, CallToolsNode
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Callable
 
-try:
-    import pydantic_ai
-    from pydantic_ai._agent_graph import UserPromptNode, ModelRequestNode, CallToolsNode
-except ImportError:
-    raise DidNotEnable("pydantic-ai not installed")
-
 
 def _patch_graph_nodes():
     # type: () -> None
diff --git a/sentry_sdk/integrations/pydantic_ai/patches/model_request.py b/sentry_sdk/integrations/pydantic_ai/patches/model_request.py
@@ -1,18 +1,13 @@
 from functools import wraps
+from typing import TYPE_CHECKING
 
-from sentry_sdk.integrations import DidNotEnable
+from pydantic_ai import models
 
 from ..spans import ai_client_span, update_ai_client_span
 
-from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Callable
-
-try:
-    import pydantic_ai
-except ImportError:
-    raise DidNotEnable("pydantic-ai not installed")
+    from typing import Any
 
 
 def _patch_model_request():
@@ -23,7 +18,6 @@ def _patch_model_request():
     In pydantic-ai, model requests are handled through the Model interface.
     We need to patch the request method on models to create spans.
     """
-    from pydantic_ai import models
 
     # Patch the base Model class's request method
     if hasattr(models, "Model"):
diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py
@@ -1,6 +1,9 @@
 from functools import wraps
 
-from sentry_sdk.integrations import DidNotEnable
+from pydantic_ai.toolsets.abstract import AbstractToolset
+from pydantic_ai.toolsets.function import FunctionToolset
+
+import sentry_sdk
 
 from ..spans import execute_tool_span, update_execute_tool_span
 
@@ -9,11 +12,6 @@
 if TYPE_CHECKING:
     from typing import Any, Callable
 
-try:
-    import pydantic_ai
-except ImportError:
-    raise DidNotEnable("pydantic-ai not installed")
-
 
 def _patch_tool_execution():
     # type: () -> None
@@ -26,8 +24,6 @@ def _patch_tool_execution():
     Note: pydantic-ai has built-in OpenTelemetry instrumentation for tools.
     Our patching adds Sentry-specific span data on top of that.
     """
-    from pydantic_ai.toolsets.abstract import AbstractToolset
-    from pydantic_ai.toolsets.function import FunctionToolset
 
     def create_wrapped_call_tool(original_call_tool):
         # type: (Callable[..., Any]) -> Callable[..., Any]
@@ -37,7 +33,6 @@ def create_wrapped_call_tool(original_call_tool):
         async def wrapped_call_tool(self, name, args_dict, ctx, tool):
             # type: (Any, str, Any, Any, Any) -> Any
             # Always create span if we're in a Sentry transaction context
-            import sentry_sdk
 
             current_span = sentry_sdk.get_current_span()
             should_create_span = current_span is not None
diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py
@@ -1,14 +1,14 @@
 import sentry_sdk
+from sentry_sdk.ai.utils import set_data_normalized
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.utils import safe_serialize
 
 from ..consts import SPAN_ORIGIN
 from ..utils import (
-    _get_model_name,
     _set_agent_data,
     _set_model_data,
-    _set_usage_data,
-    _set_input_messages,
-    _set_output_data,
+    _should_send_prompts,
+    _get_model_name,
 )
 
 from typing import TYPE_CHECKING
@@ -17,6 +17,158 @@
     from typing import Any
 
 
+def _set_usage_data(span, usage):
+    # type: (sentry_sdk.tracing.Span, RequestUsage) -> None
+    """Set token usage data on a span."""
+    if usage is None:
+        return
+
+    if hasattr(usage, "input_tokens") and usage.input_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
+
+    if hasattr(usage, "output_tokens") and usage.output_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
+
+    if hasattr(usage, "total_tokens") and usage.total_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
+
+
+def _set_input_messages(span, messages):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set input messages data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not messages:
+        return
+
+    try:
+        formatted_messages = []
+        system_prompt = None
+
+        # Extract system prompt from any ModelRequest with instructions
+        for msg in messages:
+            if hasattr(msg, "instructions") and msg.instructions:
+                system_prompt = msg.instructions
+                break
+
+        # Add system prompt as first message if present
+        if system_prompt:
+            formatted_messages.append(
+                {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
+            )
+
+        for msg in messages:
+            if hasattr(msg, "parts"):
+                for part in msg.parts:
+                    role = "user"
+                    if hasattr(part, "__class__"):
+                        if "System" in part.__class__.__name__:
+                            role = "system"
+                        elif (
+                            "Assistant" in part.__class__.__name__
+                            or "Text" in part.__class__.__name__
+                            or "ToolCall" in part.__class__.__name__
+                        ):
+                            role = "assistant"
+                        elif "ToolReturn" in part.__class__.__name__:
+                            role = "tool"
+
+                    content = []  # type: List[Dict[str, Any] | str]
+                    tool_calls = None
+                    tool_call_id = None
+
+                    # Handle ToolCallPart (assistant requesting tool use)
+                    if "ToolCall" in part.__class__.__name__:
+                        tool_call_data = {}
+                        if hasattr(part, "tool_name"):
+                            tool_call_data["name"] = part.tool_name
+                        if hasattr(part, "args"):
+                            tool_call_data["arguments"] = safe_serialize(part.args)
+                        if tool_call_data:
+                            tool_calls = [tool_call_data]
+                    # Handle ToolReturnPart (tool result)
+                    elif "ToolReturn" in part.__class__.__name__:
+                        if hasattr(part, "tool_name"):
+                            tool_call_id = part.tool_name
+                        if hasattr(part, "content"):
+                            content.append({"type": "text", "text": str(part.content)})
+                    # Handle regular content
+                    elif hasattr(part, "content"):
+                        if isinstance(part.content, str):
+                            content.append({"type": "text", "text": part.content})
+                        elif isinstance(part.content, list):
+                            for item in part.content:
+                                if isinstance(item, str):
+                                    content.append({"type": "text", "text": item})
+                                else:
+                                    content.append(safe_serialize(item))
+                        else:
+                            content.append({"type": "text", "text": str(part.content)})
+
+                    # Add message if we have content or tool calls
+                    if content or tool_calls:
+                        message = {"role": role}  # type: Dict[str, Any]
+                        if content:
+                            message["content"] = content
+                        if tool_calls:
+                            message["tool_calls"] = tool_calls
+                        if tool_call_id:
+                            message["tool_call_id"] = tool_call_id
+                        formatted_messages.append(message)
+
+        if formatted_messages:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, formatted_messages, unpack=False
+            )
+    except Exception:
+        # If we fail to format messages, just skip it
+        pass
+
+
+def _set_output_data(span, response):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set output data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not response:
+        return
+
+    set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name)
+    try:
+        # Extract text from ModelResponse
+        if hasattr(response, "parts"):
+            texts = []
+            tool_calls = []
+
+            for part in response.parts:
+                if hasattr(part, "__class__"):
+                    if "Text" in part.__class__.__name__ and hasattr(part, "content"):
+                        texts.append(part.content)
+                    elif "ToolCall" in part.__class__.__name__:
+                        tool_call_data = {
+                            "type": "function",
+                        }
+                        if hasattr(part, "tool_name"):
+                            tool_call_data["name"] = part.tool_name
+                        if hasattr(part, "args"):
+                            tool_call_data["arguments"] = safe_serialize(part.args)
+                        tool_calls.append(tool_call_data)
+
+            if texts:
+                set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts)
+
+            if tool_calls:
+                span.set_data(
+                    SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)
+                )
+
+    except Exception:
+        # If we fail to format output, just skip it
+        pass
+
+
 def ai_client_span(messages, agent, model, model_settings):
     # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span
     """Create a span for an AI client call (model request).
@@ -72,11 +224,11 @@ def ai_client_span(messages, agent, model, model_settings):
                     # Add description from function_schema if available
                     if hasattr(tool, "function_schema"):
                         schema = tool.function_schema
-                        if hasattr(schema, "description") and schema.description:
+                        if getattr(schema, "description", None):
                             tool_info["description"] = schema.description
 
                         # Add parameters from json_schema
-                        if hasattr(schema, "json_schema") and schema.json_schema:
+                        if getattr(schema, "json_schema", None):
                             tool_info["parameters"] = schema.json_schema
 
                     tools.append(tool_info)
diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py
@@ -4,7 +4,11 @@
 from sentry_sdk.utils import safe_serialize
 
 from ..consts import SPAN_ORIGIN
-from ..utils import _set_agent_data, _set_model_data, _should_send_prompts
+from ..utils import (
+    _set_agent_data,
+    _set_model_data,
+    _should_send_prompts,
+)
 
 from typing import TYPE_CHECKING
 
@@ -46,11 +50,11 @@ def invoke_agent_span(user_prompt, agent, model, model_settings):
                     # Add description from function_schema if available
                     if hasattr(tool, "function_schema"):
                         schema = tool.function_schema
-                        if hasattr(schema, "description") and schema.description:
+                        if getattr(schema, "description", None):
                             tool_info["description"] = schema.description
 
                         # Add parameters from json_schema
-                        if hasattr(schema, "json_schema") and schema.json_schema:
+                        if getattr(schema, "json_schema", None):
                             tool_info["parameters"] = schema.json_schema
 
                     tools.append(tool_info)
@@ -72,14 +76,14 @@ def invoke_agent_span(user_prompt, agent, model, model_settings):
 
         if agent:
             # Check for system_prompt
-            if hasattr(agent, "_system_prompts") and agent._system_prompts:
-                for prompt in agent._system_prompts:
-                    if isinstance(prompt, str):
-                        system_texts.append(prompt)
+            system_prompts = getattr(agent, "_system_prompts", None) or []
+            for prompt in system_prompts:
+                if isinstance(prompt, str):
+                    system_texts.append(prompt)
 
             # Check for instructions (stored in _instructions)
-            if hasattr(agent, "_instructions") and agent._instructions:
-                instructions = agent._instructions
+            instructions = getattr(agent, "_instructions", None)
+            if instructions:
                 if isinstance(instructions, str):
                     system_texts.append(instructions)
                 elif isinstance(instructions, (list, tuple)):
@@ -134,9 +138,8 @@ def update_invoke_agent_span(span, output):
     # type: (sentry_sdk.tracing.Span, Any) -> None
     """Update and close the invoke agent span."""
     if span and _should_send_prompts() and output:
-        output_text = str(output) if not isinstance(output, str) else output
         set_data_normalized(
-            span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_text, unpack=False
+            span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False
         )
 
     if span:
diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py
