fix: mypy issues

Author: constantinius

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -9,7 +9,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Callable
+    from typing import Any, Callable, Optional
 
 try:
     import pydantic_ai
@@ -25,8 +25,8 @@ def __init__(self, agent, original_ctx_manager, is_streaming=True):
         self.agent = agent
         self.original_ctx_manager = original_ctx_manager
         self.is_streaming = is_streaming
-        self._isolation_scope = None
-        self._workflow_span = None
+        self._isolation_scope = None  # type: Any
+        self._workflow_span = None  # type: Optional[sentry_sdk.tracing.Span]
 
     async def __aenter__(self):
         # type: () -> Any
@@ -217,6 +217,6 @@ def _patch_agent_run():
 
     # Wrap and apply patches for streaming methods
     Agent.run_stream = _create_streaming_wrapper(original_run_stream)  # type: ignore
-    Agent.run_stream_events = _create_streaming_events_wrapper(
+    Agent.run_stream_events = _create_streaming_events_wrapper(  # type: ignore[method-assign]
         original_run_stream_events
-    )  # type: ignore
+    )

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -83,12 +83,12 @@ async def wrapped_user_prompt_run(self, ctx):
         span = invoke_agent_span(user_prompt, agent, model, model_settings)
         # Store span in context for later use
         if hasattr(ctx, "state"):
-            ctx.state._sentry_invoke_span = span  # type: ignore
+            ctx.state._sentry_invoke_span = span
 
         result = await original_user_prompt_run(self, ctx)
         return result
 
-    UserPromptNode.run = wrapped_user_prompt_run  # type: ignore
+    UserPromptNode.run = wrapped_user_prompt_run  # type: ignore[method-assign]
 
     # Patch ModelRequestNode to create ai_client spans
     original_model_request_run = ModelRequestNode.run

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -9,7 +9,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, List, Dict
     from pydantic_ai.usage import RequestUsage
 
 try:
@@ -213,7 +213,7 @@ def _set_input_messages(span, messages):
                         elif "ToolReturn" in part.__class__.__name__:
                             role = "tool"
 
-                    content = []
+                    content = []  # type: List[Dict[str, Any] | str]
                     tool_calls = None
                     tool_call_id = None
 
@@ -247,7 +247,7 @@ def _set_input_messages(span, messages):
 
                     # Add message if we have content or tool calls
                     if content or tool_calls:
-                        message = {"role": role}
+                        message = {"role": role}  # type: Dict[str, Any]
                         if content:
                             message["content"] = content
                         if tool_calls: