linting

Author: antonpirker

pyproject.toml

@@ -183,6 +183,10 @@ ignore_missing_imports = true
 module = "grpc.*"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "agents.*"
+ignore_missing_imports = true
+
 #
 # Tool: Flake8
 #

sentry_sdk/integrations/openai_agents/__init__.py

@@ -15,19 +15,19 @@
 
 def _patch_runner():
     # type: () -> None
-    agents.Runner.run = _create_run_wrapper(agents.Runner.run)  # type: ignore[method-assign]
-    agents.Runner.run_sync = _create_run_wrapper(agents.Runner.run_sync)  # type: ignore[method-assign]
-    agents.Runner.run_streamed = _create_run_wrapper(agents.Runner.run_streamed)  # type: ignore[method-assign]
+    agents.Runner.run = _create_run_wrapper(agents.Runner.run)
+    agents.Runner.run_sync = _create_run_wrapper(agents.Runner.run_sync)
+    agents.Runner.run_streamed = _create_run_wrapper(agents.Runner.run_streamed)
 
 
 def _patch_model():
     # type: () -> None
-    agents.Runner._get_model = _create_get_model_wrapper(agents.Runner._get_model)  # type: ignore[method-assign]
+    agents.Runner._get_model = _create_get_model_wrapper(agents.Runner._get_model)
 
 
 def _patch_tools():
     # type: () -> None
-    agents.Runner._get_all_tools = _create_get_all_tools_wrapper(  # type: ignore[method-assign]
+    agents.Runner._get_all_tools = _create_get_all_tools_wrapper(
         agents.Runner._get_all_tools
     )
 

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -4,6 +4,11 @@
 
 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 agents
@@ -12,6 +17,7 @@
 
 
 def _create_get_model_wrapper(original_get_model):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
     """
     Wraps the agents.Runner._get_model method to wrap the get_response method of the model to create a AI client span.
     """
@@ -26,6 +32,7 @@ def wrapped_get_model(cls, agent, run_config):
 
         @wraps(original_get_response)
         async def wrapped_get_response(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
             with ai_client_span(agent, kwargs) as span:
                 result = await original_get_response(*args, **kwargs)
                 update_ai_client_span(span, agent, kwargs, result)

sentry_sdk/integrations/openai_agents/patches/runner.py

@@ -10,8 +10,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
-    from typing import Callable
+    from typing import Any, Callable
 
 try:
     import agents
@@ -20,7 +19,7 @@
 
 
 def _create_run_wrapper(original_func):
-    # type: (Callable) -> Callable
+    # type: (Callable[..., Any]) -> Callable[..., Any]
     """
     Wraps the agents.Runner.run* methods to create a root span for the agent workflow runs.
     """

sentry_sdk/integrations/openai_agents/patches/tools.py

@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Callable
 
 try:
     import agents
@@ -16,6 +16,7 @@
 
 
 def _create_get_all_tools_wrapper(original_get_all_tools):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
     """
     Wraps the agents.Runner._get_all_tools method of the Runner class to wrap all function tools with Sentry instrumentation.
     """
@@ -39,6 +40,7 @@ async def wrapped_get_all_tools(cls, agent, context_wrapper):
             original_on_invoke = tool.on_invoke_tool
 
             def create_wrapped_invoke(current_tool, current_on_invoke):
+                # type: (agents.Tool, Callable[..., Any]) -> Callable[..., Any]
                 @wraps(current_on_invoke)
                 async def sentry_wrapped_on_invoke_tool(*args, **kwargs):
                     # type: (*Any, **Any) -> Any

sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

@@ -5,12 +5,11 @@
 
 if TYPE_CHECKING:
     from agents import Agent
-    from sentry_sdk import Span
     from typing import Any
 
 
 def agent_workflow_span(*args, **kwargs):
-    # type: (*Any, **Any) -> Span
+    # type: (*Any, **Any) -> sentry_sdk.tracing.Span
     agent = args[0]
     # Create a transaction or a span if an transaction is already active
     span = _get_start_span_function()(
@@ -22,5 +21,5 @@ def agent_workflow_span(*args, **kwargs):
 
 
 def update_agent_workflow_span(span, agent, result):
-    # type: (Span, Agent, Any) -> None
+    # type: (sentry_sdk.tracing.Span, Agent, Any) -> None
     pass

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -11,7 +11,7 @@
 
 
 def invoke_agent_span(context, agent):
-    # type: (agents.RunContextWrapper, agents.Agent) -> None
+    # type: (agents.RunContextWrapper, agents.Agent) -> sentry_sdk.tracing.Span
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_INVOKE_AGENT,
         name=f"invoke_agent {agent.name}",

sentry_sdk/integrations/openai_agents/utils.py

@@ -44,6 +44,7 @@ def _create_hook_wrapper(original_hook, sentry_hook):
     # type: (Callable[..., Any], Callable[..., Any]) -> Callable[..., Any]
     @wraps(original_hook)
     async def async_wrapper(*args, **kwargs):
+        # type: (*Any, **Any) -> Any
         await sentry_hook(*args, **kwargs)
         return await original_hook(*args, **kwargs)
 
@@ -153,7 +154,7 @@ def _set_input_data(span, get_response_kwargs):
         "user": [],
         "assistant": [],
         "tool": [],
-    }
+    }  # type: (dict[str, list[Any]])
     system_instructions = get_response_kwargs.get("system_instructions")
     if system_instructions:
         messages_by_role["system"].append({"type": "text", "text": system_instructions})
@@ -185,7 +186,7 @@ def _set_output_data(span, result):
     output_messages = {
         "response": [],
         "tool": [],
-    }
+    }  # type: (dict[str, list[Any]])
 
     for output in result.output:
         if output.type == "function_call":