Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ typing = [
"openai>=3.7.0",
"tiktoken>=0.14.0",
"aiohttp>=3.14.3",
"openai-agents>=0.22.0",
]
test = [
"dataclasses ; python_full_version < '3.7'",
Expand Down
44 changes: 22 additions & 22 deletions sentry_sdk/integrations/openai_agents/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@
# https://github.com/openai/openai-agents-python/commit/3ce7c24d349b77bb750062b7e0e856d9ff48a5d5#diff-7470b3a5c5cbe2fcbb2703dc24f326f45a5819d853be2b1f395d122d278cd911
from agents.run_internal import run_loop, turn_preparation, turn_resolution
except ImportError:
run_loop = None
turn_preparation = None
turn_resolution = None
run_loop = None # type: ignore[assignment]
turn_preparation = None # type: ignore[assignment]
turn_resolution = None # type: ignore[assignment]

from typing import TYPE_CHECKING

Expand All @@ -50,13 +50,13 @@ def _patch_runner(use_run_hooks: "bool") -> None:
# Create the root span for one full agent run (including eventual handoffs)
# Note agents.run.DEFAULT_AGENT_RUNNER.run_sync is a wrapper around
# agents.run.DEFAULT_AGENT_RUNNER.run. It does not need to be wrapped separately.
agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper( # type: ignore[method-assign]
agents.run.DEFAULT_AGENT_RUNNER.run,
use_run_hooks=use_run_hooks,
)

# Patch streaming runner
agents.run.DEFAULT_AGENT_RUNNER.run_streamed = _create_run_streamed_wrapper(
agents.run.DEFAULT_AGENT_RUNNER.run_streamed = _create_run_streamed_wrapper( # type: ignore[method-assign]
agents.run.DEFAULT_AGENT_RUNNER.run_streamed,
use_run_hooks=use_run_hooks,
)
Expand Down Expand Up @@ -175,9 +175,9 @@ async def new_wrapped_final_output(
return

if not use_run_hooks:
original_get_all_tools = AgentRunner._get_all_tools
original_get_all_tools = AgentRunner._get_all_tools # type: ignore[attr-defined]

@wraps(AgentRunner._get_all_tools.__func__)
@wraps(AgentRunner._get_all_tools.__func__) # type: ignore[attr-defined]
async def old_wrapped_get_all_tools(
cls: "agents.Runner",
agent: "agents.Agent",
Expand All @@ -187,44 +187,44 @@ async def old_wrapped_get_all_tools(
original_get_all_tools, agent, context_wrapper
)

agents.run.AgentRunner._get_all_tools = classmethod(
old_wrapped_get_all_tools
agents.run.AgentRunner._get_all_tools = classmethod( # type: ignore[attr-defined]
old_wrapped_get_all_tools # type: ignore[arg-type]
)

original_get_model = AgentRunner._get_model
original_get_model = AgentRunner._get_model # type: ignore[attr-defined]

@wraps(AgentRunner._get_model.__func__)
@wraps(AgentRunner._get_model.__func__) # type: ignore[attr-defined]
def old_wrapped_get_model(
cls: "agents.Runner", agent: "agents.Agent", run_config: "agents.RunConfig"
) -> "agents.Model":
return _get_model(original_get_model, agent, run_config)

agents.run.AgentRunner._get_model = classmethod(old_wrapped_get_model)
agents.run.AgentRunner._get_model = classmethod(old_wrapped_get_model) # type: ignore[arg-type,attr-defined]

original_run_single_turn = AgentRunner._run_single_turn
original_run_single_turn = AgentRunner._run_single_turn # type: ignore[attr-defined]

@wraps(AgentRunner._run_single_turn.__func__)
@wraps(AgentRunner._run_single_turn.__func__) # type: ignore[attr-defined]
async def old_wrapped_run_single_turn(
cls: "agents.Runner", *args: "Any", **kwargs: "Any"
) -> "SingleStepResult":
return await _run_single_turn(original_run_single_turn, *args, **kwargs)

agents.run.AgentRunner._run_single_turn = classmethod(
old_wrapped_run_single_turn
agents.run.AgentRunner._run_single_turn = classmethod( # type: ignore[attr-defined]
old_wrapped_run_single_turn # type: ignore[arg-type]
)

original_run_single_turn_streamed = AgentRunner._run_single_turn_streamed
original_run_single_turn_streamed = AgentRunner._run_single_turn_streamed # type: ignore[attr-defined]

@wraps(AgentRunner._run_single_turn_streamed.__func__)
@wraps(AgentRunner._run_single_turn_streamed.__func__) # type: ignore[attr-defined]
async def old_wrapped_run_single_turn_streamed(
cls: "agents.Runner", *args: "Any", **kwargs: "Any"
) -> "SingleStepResult":
return await _run_single_turn_streamed(
original_run_single_turn_streamed, *args, **kwargs
)

agents.run.AgentRunner._run_single_turn_streamed = classmethod(
old_wrapped_run_single_turn_streamed
agents.run.AgentRunner._run_single_turn_streamed = classmethod( # type: ignore[attr-defined]
old_wrapped_run_single_turn_streamed # type: ignore[arg-type]
)

original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
Expand All @@ -236,7 +236,7 @@ async def old_wrapped_execute_handoffs(
return await _execute_handoffs(original_execute_handoffs, *args, **kwargs)

agents._run_impl.RunImpl.execute_handoffs = classmethod(
old_wrapped_execute_handoffs
old_wrapped_execute_handoffs # type: ignore[arg-type]
)

original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
Expand All @@ -250,5 +250,5 @@ async def old_wrapped_final_output(
)

agents._run_impl.RunImpl.execute_final_output = classmethod(
old_wrapped_final_output
old_wrapped_final_output # type: ignore[arg-type]
)
28 changes: 16 additions & 12 deletions sentry_sdk/integrations/openai_agents/patches/agent_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
)

if TYPE_CHECKING:
from typing import Any, Awaitable, Callable, Optional, Union
from typing import Any, Awaitable, Callable, Coroutine, Optional, Union

from agents.run_internal.run_steps import SingleStepResult

Expand Down Expand Up @@ -46,8 +46,8 @@ def _close_streaming_workflow_span(agent: "Optional[agents.Agent]") -> None:


def _maybe_start_agent_span(
context_wrapper: "agents.RunContextWrapper",
agent: "agents.Agent",
context_wrapper: "Optional[agents.RunContextWrapper]",
agent: "Optional[agents.Agent]",
should_run_agent_start_hooks: bool,
span_kwargs: "dict[str, Any]",
is_streaming: bool = False,
Expand All @@ -74,10 +74,10 @@ def _maybe_start_agent_span(
delattr(context_wrapper, "_sentry_agent_span")

# Store the agent on the context wrapper so we can access it later
context_wrapper._sentry_current_agent = agent
context_wrapper._sentry_current_agent = agent # type: ignore[attr-defined]
span = invoke_agent_span(context_wrapper, agent, span_kwargs)
context_wrapper._sentry_agent_span = span
agent._sentry_agent_span = span
context_wrapper._sentry_agent_span = span # type: ignore[attr-defined]
agent._sentry_agent_span = span # type: ignore[attr-defined]

if not is_streaming:
return span
Expand All @@ -102,12 +102,14 @@ async def _run_single_turn(
"""
# openai-agents >= 0.14 passes `bindings: AgentBindings` instead of `agent`.
bindings = kwargs.get("bindings")
agent = (
agent: "Optional[agents.Agent]" = (
getattr(bindings, "public_agent", None)
if bindings is not None
else kwargs.get("agent")
)
context_wrapper = kwargs.get("context_wrapper")
context_wrapper: "Optional[agents.RunContextWrapper]" = kwargs.get(
"context_wrapper"
)
should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks", False)

span = _maybe_start_agent_span(
Expand Down Expand Up @@ -224,7 +226,7 @@ async def _run_single_turn_streamed(


async def _execute_handoffs(
original_execute_handoffs: "Callable[..., SingleStepResult]",
original_execute_handoffs: "Callable[..., Coroutine[Any, Any, SingleStepResult]]",
*args: "Any",
**kwargs: "Any",
) -> "SingleStepResult":
Expand All @@ -235,10 +237,12 @@ async def _execute_handoffs(
- ends the workflow span if the response is streamed and an exception is raised in `execute_handoffs()`.
"""

context_wrapper = kwargs.get("context_wrapper")
context_wrapper: "Optional[agents.RunContextWrapper]" = kwargs.get(
"context_wrapper"
)
run_handoffs = kwargs.get("run_handoffs")
# openai-agents >= 0.14 renamed `agent` to `public_agent`.
agent = kwargs.get("public_agent", kwargs.get("agent"))
agent: "Optional[agents.Agent]" = kwargs.get("public_agent", kwargs.get("agent"))

# Create Sentry handoff span for the first handoff (agents library only processes the first one)
if run_handoffs:
Expand Down Expand Up @@ -282,7 +286,7 @@ async def _execute_handoffs(


async def _execute_final_output(
original_execute_final_output: "Callable[..., SingleStepResult]",
original_execute_final_output: "Callable[..., Coroutine[Any, Any, SingleStepResult]]",
*args: "Any",
**kwargs: "Any",
) -> "SingleStepResult":
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ def sentry_attach_error_to_current_span(
# Call the original function
return original_attach_error(error, *args, **kwargs)

error_tracing_module.attach_error_to_current_span = (
sentry_attach_error_to_current_span
setattr(
error_tracing_module,
"attach_error_to_current_span",
sentry_attach_error_to_current_span,
)
8 changes: 4 additions & 4 deletions sentry_sdk/integrations/openai_agents/patches/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def _get_model(

# Capture the request model name for spans (agent.model can be None when using defaults)
request_model_name = model.model if hasattr(model, "model") else str(model)
agent._sentry_request_model = request_model_name
agent._sentry_request_model = request_model_name # type: ignore[attr-defined]

# Wrap _fetch_response if it exists (for OpenAI models) to capture response model
if hasattr(model, "_fetch_response"):
Expand All @@ -83,7 +83,7 @@ def _get_model(
async def wrapped_fetch_response(*args: "Any", **kwargs: "Any") -> "Any":
response = await original_fetch_response(*args, **kwargs)
if hasattr(response, "model") and response.model:
agent._sentry_response_model = str(response.model)
agent._sentry_response_model = str(response.model) # type: ignore[attr-defined]
return response

model._fetch_response = wrapped_fetch_response
Expand Down Expand Up @@ -114,7 +114,7 @@ async def wrapped_get_response(*args: "Any", **kwargs: "Any") -> "Any":

return result

model.get_response = wrapped_get_response
model.get_response = wrapped_get_response # type: ignore[method-assign]

# Also wrap stream_response for streaming support
if hasattr(model, "stream_response"):
Expand Down Expand Up @@ -177,6 +177,6 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any":
span, streaming_response, response_model, agent
)

model.stream_response = wrapped_stream_response
model.stream_response = wrapped_stream_response # type: ignore[method-assign]

return model
8 changes: 4 additions & 4 deletions sentry_sdk/integrations/openai_agents/patches/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
TContext = TypeVar("TContext")


class _SentryRunHooks(RunHooks[TContext]): # type: ignore[misc]
class _SentryRunHooks(RunHooks[TContext]):
"""
Responsible for creating and managing Execute Tool spans. These spans are
stored on the ToolContext reference that is shared between `on_tool_start()`
Expand Down Expand Up @@ -120,9 +120,9 @@ async def on_tool_end(
await sentry_hooks.on_tool_end(context, agent, tool, result)
await original_on_tool_end(context, agent, tool, result)

hooks._sentry_is_patched = True
hooks.on_tool_start = on_tool_start
hooks.on_tool_end = on_tool_end
hooks._sentry_is_patched = True # type: ignore[attr-defined]
hooks.on_tool_start = on_tool_start # type: ignore[method-assign]
hooks.on_tool_end = on_tool_end # type: ignore[method-assign]


def _create_run_wrapper(
Expand Down
4 changes: 2 additions & 2 deletions sentry_sdk/integrations/openai_agents/patches/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ async def _get_all_tools(
# Get the original tools
tools = await original_get_all_tools(agent, context_wrapper)

wrapped_tools = []
wrapped_tools: "list[agents.Tool]" = []
for tool in tools:
# Wrap only the function tools (for now)
if not isinstance(tool, FunctionTool):
Expand All @@ -42,7 +42,7 @@ async def _get_all_tools(
original_on_invoke = tool.on_invoke_tool

def create_wrapped_invoke(
current_tool: "agents.Tool", current_on_invoke: "Callable[..., Any]"
current_tool: "FunctionTool", current_on_invoke: "Callable[..., Any]"
) -> "Callable[..., Any]":
@wraps(current_on_invoke)
async def sentry_wrapped_on_invoke_tool(
Expand Down
11 changes: 5 additions & 6 deletions sentry_sdk/integrations/openai_agents/spans/ai_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,23 +25,23 @@
try:
from agents import ApplyPatchTool, ShellTool
except ImportError:
ShellTool = None
ApplyPatchTool = None
ShellTool = None # type: ignore[assignment,misc]
ApplyPatchTool = None # type: ignore[assignment,misc]

try:
from agents import ToolSearchTool
except ImportError:
ToolSearchTool = None
ToolSearchTool = None # type: ignore[assignment,misc]

try:
from agents import CustomTool
except ImportError:
CustomTool = None
CustomTool = None # type: ignore[assignment,misc]

try:
from agents import ProgrammaticToolCallingTool
except ImportError:
ProgrammaticToolCallingTool = None
ProgrammaticToolCallingTool = None # type: ignore[assignment,misc]


from ..consts import SPAN_ORIGIN
Expand Down Expand Up @@ -201,7 +201,6 @@ def ai_client_span(
agent: "Agent", get_response_kwargs: "dict[str, Any]"
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
# TODO-anton: implement other types of operations. Now "chat" is hardcoded.
# Get model name from agent.model or fall back to request model (for when agent.model is None/default)
model_name = None
if agent.model:
model_name = agent.model.model if hasattr(agent.model, "model") else agent.model
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@


def execute_tool_span(
tool: "agents.Tool", *args: "Any", **kwargs: "Any"
tool: "agents.FunctionTool", *args: "Any", **kwargs: "Any"
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)

Expand Down
10 changes: 7 additions & 3 deletions sentry_sdk/integrations/openai_agents/spans/handoff.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,20 @@
from ..consts import SPAN_ORIGIN

if TYPE_CHECKING:
from typing import Optional

import agents


def handoff_span(
context: "agents.RunContextWrapper", from_agent: "agents.Agent", to_agent_name: str
context: "Optional[agents.RunContextWrapper]",
from_agent: "Optional[agents.Agent]",
to_agent_name: str,
) -> None:
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
if span_streaming:
with sentry_sdk.traces.start_span(
name=f"handoff from {from_agent.name} to {to_agent_name}",
name=f"handoff from {from_agent.name} to {to_agent_name}", # type: ignore[union-attr]
attributes={
"sentry.op": OP.GEN_AI_HANDOFF,
"sentry.origin": SPAN_ORIGIN,
Expand All @@ -30,7 +34,7 @@ def handoff_span(
else:
with sentry_sdk.start_span(
op=OP.GEN_AI_HANDOFF,
name=f"handoff from {from_agent.name} to {to_agent_name}",
name=f"handoff from {from_agent.name} to {to_agent_name}", # type: ignore[union-attr]
Comment thread
alexander-alderman-webb marked this conversation as resolved.
origin=SPAN_ORIGIN,
) as span:
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "handoff")
Expand Down
8 changes: 4 additions & 4 deletions sentry_sdk/integrations/openai_agents/spans/invoke_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
from ..utils import _set_agent_data, _set_usage_data

if TYPE_CHECKING:
from typing import Any, Union
from typing import Any, Optional, Union

import agents

Expand Down Expand Up @@ -110,12 +110,12 @@ def invoke_agent_span(

def update_invoke_agent_span(
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]",
context: "agents.RunContextWrapper",
agent: "agents.Agent",
context: "Optional[agents.RunContextWrapper]",
agent: "Optional[agents.Agent]",
output: "Any" = None,
) -> None:
# Add aggregated usage data from context_wrapper
if hasattr(context, "usage"):
if context and hasattr(context, "usage"):
_set_usage_data(span, context.usage)

client = sentry_sdk.get_client()
Expand Down
Loading
Loading