linting

Author: antonpirker

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)
-    agents.Runner.run_sync = _create_run_wrapper(agents.Runner.run_sync)
-    agents.Runner.run_streamed = _create_run_wrapper(agents.Runner.run_streamed)
+    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]
 
 
 def _patch_model():
     # type: () -> None
-    agents.Runner._get_model = _create_get_model_wrapper(agents.Runner._get_model)
+    agents.Runner._get_model = _create_get_model_wrapper(agents.Runner._get_model)  # type: ignore[method-assign]
 
 
 def _patch_tools():
     # type: () -> None
-    agents.Runner._get_all_tools = _create_get_all_tools_wrapper(
+    agents.Runner._get_all_tools = _create_get_all_tools_wrapper(  # type: ignore[method-assign]
         agents.Runner._get_all_tools
     )
 

sentry_sdk/integrations/openai_agents/run_hooks.py

@@ -16,7 +16,7 @@
 
 
 class SentryRunHooks(agents.RunHooks):
-    async def on_agent_start(self, context, agent) -> None:
+    async def on_agent_start(self, context, agent):
         # type: (agents.RunContextWrapper, agents.Agent) -> None
         invoke_agent_span(context, agent)
 

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -13,11 +13,10 @@
 if TYPE_CHECKING:
     from agents import Agent
     from typing import Any
-    from sentry_sdk import Span
 
 
 def ai_client_span(agent, get_response_kwargs):
-    # type: (Agent, dict[str, Any]) -> Span
+    # type: (Agent, dict[str, Any]) -> sentry_sdk.tracing.Span
     # TODO-anton: implement other types of operations. Now "chat" is hardcoded.
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_CHAT,
@@ -31,7 +30,7 @@ def ai_client_span(agent, get_response_kwargs):
 
 
 def update_ai_client_span(span, agent, get_response_kwargs, result):
-    # type: (Span, Agent, dict[str, Any], Any) -> None
+    # type: (sentry_sdk.tracing.Span, Agent, dict[str, Any], Any) -> None
     _set_agent_data(span, agent)
     _set_usage_data(span, result.usage)
     _set_input_data(span, get_response_kwargs)

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -8,12 +8,11 @@
 
 if TYPE_CHECKING:
     import agents
-    from sentry_sdk import Span
     from typing import Any
 
 
 def execute_tool_span(tool, *args, **kwargs):
-    # type: (agents.Tool, *Any, **Any) -> Span
+    # type: (agents.Tool, *Any, **Any) -> sentry_sdk.tracing.Span
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_EXECUTE_TOOL,
         name=f"execute_tool {tool.name}",
@@ -36,7 +35,7 @@ def execute_tool_span(tool, *args, **kwargs):
 
 
 def update_execute_tool_span(span, agent, tool, result):
-    # type: (Span, agents.Agent, agents.Tool, Any) -> None
+    # type: (sentry_sdk.tracing.Span, agents.Agent, agents.Tool, Any) -> None
     _set_agent_data(span, agent)
 
     if should_send_default_pii():

sentry_sdk/integrations/openai_agents/utils.py

@@ -32,7 +32,7 @@ def _capture_exception(exc):
 
 
 def _get_start_span_function():
-    # type: () -> Callable
+    # type: () -> Callable[..., Any]
     current_span = sentry_sdk.get_current_span()
     is_transaction = (
         current_span is not None and current_span.containing_transaction == current_span
@@ -41,7 +41,7 @@ def _get_start_span_function():
 
 
 def _create_hook_wrapper(original_hook, sentry_hook):
-    # type: (Callable, Callable) -> Callable
+    # type: (Callable[..., Any], Callable[..., Any]) -> Callable[..., Any]
     @wraps(original_hook)
     async def async_wrapper(*args, **kwargs):
         await sentry_hook(*args, **kwargs)
@@ -86,7 +86,7 @@ def _wrap_hooks(hooks):
 
 
 def _set_agent_data(span, agent):
-    # type: (sentry_sdk.Span, agents.Agent) -> None
+    # type: (sentry_sdk.tracing.Span, agents.Agent) -> None
     span.set_data(
         SPANDATA.GEN_AI_SYSTEM, "openai"
     )  # See footnote for  https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-system for explanation why.
@@ -129,7 +129,7 @@ def _set_agent_data(span, agent):
 
 
 def _set_usage_data(span, usage):
-    # type: (sentry_sdk.Span, Usage) -> None
+    # type: (sentry_sdk.tracing.Span, Usage) -> None
     span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
     span.set_data(
         SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED,
@@ -144,7 +144,7 @@ def _set_usage_data(span, usage):
 
 
 def _set_input_data(span, get_response_kwargs):
-    # type: (sentry_sdk.Span, dict[str, Any]) -> None
+    # type: (sentry_sdk.tracing.Span, dict[str, Any]) -> None
     if not should_send_default_pii():
         return
 
@@ -178,7 +178,7 @@ def _set_input_data(span, get_response_kwargs):
 
 
 def _set_output_data(span, result):
-    # type: (sentry_sdk.Span, Any) -> None
+    # type: (sentry_sdk.tracing.Span, Any) -> None
     if not should_send_default_pii():
         return
 

tests/integrations/openai_agents/__init__.py

@@ -1,3 +1,3 @@
 import pytest
 
-pytest.importorskip("openai-agents")
+pytest.importorskip("agents")