diff --git a/sentry_sdk/integrations/openai_agents/__init__.py b/sentry_sdk/integrations/openai_agents/__init__.py index 94c5ecc7a5..dc2399947b 100644 --- a/sentry_sdk/integrations/openai_agents/__init__.py +++ b/sentry_sdk/integrations/openai_agents/__init__.py @@ -112,7 +112,7 @@ async def new_wrapped_run_single_turn( *args: "Any", **kwargs: "Any" ) -> "SingleStepResult": return await _run_single_turn( - run_loop.run_single_turn, *args, **kwargs + run_loop.run_single_turn, use_run_hooks, *args, **kwargs ) agents.run.run_single_turn = new_wrapped_run_single_turn @@ -125,6 +125,7 @@ async def new_wrapped_run_single_turn_streamed( ) -> "SingleStepResult": return await _run_single_turn_streamed( original_run_single_turn_streamed, + use_run_hooks, *args, **kwargs, ) @@ -151,7 +152,7 @@ async def new_wrapped_execute_handoffs( *args: "Any", **kwargs: "Any" ) -> "SingleStepResult": return await _execute_handoffs( - original_execute_handoffs, *args, **kwargs + original_execute_handoffs, use_run_hooks, *args, **kwargs ) agents.run_internal.turn_resolution.execute_handoffs = ( @@ -207,7 +208,9 @@ def old_wrapped_get_model( 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) + return await _run_single_turn( + original_run_single_turn, use_run_hooks, *args, **kwargs + ) agents.run.AgentRunner._run_single_turn = classmethod( old_wrapped_run_single_turn @@ -220,7 +223,7 @@ 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 + original_run_single_turn_streamed, use_run_hooks, *args, **kwargs ) agents.run.AgentRunner._run_single_turn_streamed = classmethod( @@ -233,7 +236,9 @@ async def old_wrapped_run_single_turn_streamed( async def old_wrapped_execute_handoffs( cls: "agents.Runner", *args: "Any", **kwargs: "Any" ) -> "SingleStepResult": - return await _execute_handoffs(original_execute_handoffs, *args, **kwargs) + return await _execute_handoffs( + original_execute_handoffs, use_run_hooks, *args, **kwargs + ) agents._run_impl.RunImpl.execute_handoffs = classmethod( old_wrapped_execute_handoffs diff --git a/sentry_sdk/integrations/openai_agents/patches/agent_run.py b/sentry_sdk/integrations/openai_agents/patches/agent_run.py index 71883b2eef..966577e0c9 100644 --- a/sentry_sdk/integrations/openai_agents/patches/agent_run.py +++ b/sentry_sdk/integrations/openai_agents/patches/agent_run.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from typing import Any, Awaitable, Callable, Optional, Union + from agents import TResponseInputItem from agents.run_internal.run_steps import SingleStepResult from sentry_sdk.tracing import Span @@ -49,7 +50,7 @@ def _maybe_start_agent_span( context_wrapper: "agents.RunContextWrapper", agent: "agents.Agent", should_run_agent_start_hooks: bool, - span_kwargs: "dict[str, Any]", + turn_input: "Optional[list[TResponseInputItem]]", is_streaming: bool = False, ) -> "Optional[Union[Span, StreamedSpan]]": """ @@ -68,14 +69,14 @@ def _maybe_start_agent_span( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent + span=span, usage=context_wrapper.usage, agent=agent ) span.__exit__(None, None, None) 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 - span = invoke_agent_span(context_wrapper, agent, span_kwargs) + span = invoke_agent_span(agent, turn_input) context_wrapper._sentry_agent_span = span agent._sentry_agent_span = span @@ -92,6 +93,7 @@ def _maybe_start_agent_span( async def _run_single_turn( original_run_single_turn: "Callable[..., Awaitable[SingleStepResult]]", + use_run_hooks: "bool", *args: "Any", **kwargs: "Any", ) -> "SingleStepResult": @@ -107,32 +109,48 @@ async def _run_single_turn( if bindings is not None else kwargs.get("agent") ) - context_wrapper = kwargs.get("context_wrapper") - should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks", False) - span = _maybe_start_agent_span( - context_wrapper, agent, should_run_agent_start_hooks, kwargs - ) + context_wrapper: "agents.RunContextWrapper[Any]" = kwargs.get("context_wrapper") + if not use_run_hooks: + should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks", False) - if ( - span is None - or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) - or (not isinstance(span, StreamedSpan) and span.timestamp is not None) - ): - return await original_run_single_turn(*args, **kwargs) + span = _maybe_start_agent_span( + context_wrapper, + agent, + should_run_agent_start_hooks, + kwargs.get("input"), + ) + + if ( + span is None + or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) + or (not isinstance(span, StreamedSpan) and span.timestamp is not None) + ): + return await original_run_single_turn(*args, **kwargs) try: result = await original_run_single_turn(*args, **kwargs) except Exception: exc_info = sys.exc_info() with capture_internal_exceptions(): - span = getattr(context_wrapper, "_sentry_agent_span", None) - if span: - update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent - ) - span.__exit__(*exc_info) - delattr(context_wrapper, "_sentry_agent_span") + if use_run_hooks: + run_hooks = kwargs.get("hooks") + if run_hooks is not None: + span = getattr(run_hooks, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + del run_hooks._sentry_invoke_agent_span + span.__exit__(*exc_info) + else: + span = getattr(context_wrapper, "_sentry_agent_span", None) + if span: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + span.__exit__(*exc_info) + delattr(context_wrapper, "_sentry_agent_span") reraise(*exc_info) return result @@ -140,6 +158,7 @@ async def _run_single_turn( async def _run_single_turn_streamed( original_run_single_turn_streamed: "Callable[..., Awaitable[SingleStepResult]]", + use_run_hooks: "bool", *args: "Any", **kwargs: "Any", ) -> "SingleStepResult": @@ -181,42 +200,59 @@ async def _run_single_turn_streamed( args[1] if len(args) > 1 else kwargs.get("bindings", kwargs.get("agent")) ) agent = getattr(agent_or_bindings, "public_agent", agent_or_bindings) - context_wrapper = args[3] if len(args) > 3 else kwargs.get("context_wrapper") - should_run_agent_start_hooks = bool( - args[5] if len(args) > 5 else kwargs.get("should_run_agent_start_hooks", False) + + context_wrapper: "agents.RunContextWrapper[Any]" = ( + args[3] if len(args) > 3 else kwargs.get("context_wrapper") ) + if not use_run_hooks: + should_run_agent_start_hooks = bool( + args[5] + if len(args) > 5 + else kwargs.get("should_run_agent_start_hooks", False) + ) - span_kwargs: "dict[str, Any]" = {} - if streamed_result and hasattr(streamed_result, "input"): - span_kwargs["original_input"] = streamed_result.input + span_kwargs: "dict[str, Any]" = {} + if streamed_result and hasattr(streamed_result, "input"): + span_kwargs["original_input"] = streamed_result.input - span = _maybe_start_agent_span( - context_wrapper, - agent, - should_run_agent_start_hooks, - span_kwargs, - is_streaming=True, - ) + span = _maybe_start_agent_span( + context_wrapper, + agent, + should_run_agent_start_hooks, + getattr(streamed_result, "input", None), + is_streaming=True, + ) - if ( - span is None - or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) - or (not isinstance(span, StreamedSpan) and span.timestamp is not None) - ): - return await original_run_single_turn_streamed(*args, **kwargs) + if ( + span is None + or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) + or (not isinstance(span, StreamedSpan) and span.timestamp is not None) + ): + return await original_run_single_turn_streamed(*args, **kwargs) try: result = await original_run_single_turn_streamed(*args, **kwargs) except Exception: exc_info = sys.exc_info() with capture_internal_exceptions(): - span = getattr(context_wrapper, "_sentry_agent_span", None) - if span: - update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent - ) - span.__exit__(*exc_info) - delattr(context_wrapper, "_sentry_agent_span") + if use_run_hooks: + run_hooks = args[2] if len(args) > 2 else kwargs.get("hooks") + if run_hooks is not None: + span = getattr(run_hooks, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + del run_hooks._sentry_invoke_agent_span + span.__exit__(*exc_info) + else: + span = getattr(context_wrapper, "_sentry_agent_span", None) + if span: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + span.__exit__(*exc_info) + delattr(context_wrapper, "_sentry_agent_span") _close_streaming_workflow_span(agent) reraise(*exc_info) @@ -225,6 +261,7 @@ async def _run_single_turn_streamed( async def _execute_handoffs( original_execute_handoffs: "Callable[..., SingleStepResult]", + use_run_hooks: "bool", *args: "Any", **kwargs: "Any", ) -> "SingleStepResult": @@ -241,12 +278,17 @@ async def _execute_handoffs( 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: + if not use_run_hooks and run_handoffs: first_handoff = run_handoffs[0] handoff_agent_name = first_handoff.handoff.agent_name handoff_span(context_wrapper, agent, handoff_agent_name) - if not agent or not context_wrapper or not _has_active_agent_span(context_wrapper): + if ( + use_run_hooks + or not agent + or not context_wrapper + or not _has_active_agent_span(context_wrapper) + ): # Call original method with all parameters try: return await original_execute_handoffs(*args, **kwargs) @@ -266,7 +308,7 @@ async def _execute_handoffs( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent + span=span, usage=context_wrapper.usage, agent=agent ) span.__exit__(*exc_info) delattr(context_wrapper, "_sentry_agent_span") @@ -274,7 +316,7 @@ async def _execute_handoffs( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: - update_invoke_agent_span(span=span, context=context_wrapper, agent=agent) + update_invoke_agent_span(span=span, usage=context_wrapper.usage, agent=agent) span.__exit__(None, None, None) delattr(context_wrapper, "_sentry_agent_span") @@ -315,7 +357,10 @@ async def _execute_final_output( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent, output=final_output + span=span, + usage=context_wrapper.usage, + agent=agent, + output=final_output, ) span.__exit__(*exc_info) delattr(context_wrapper, "_sentry_agent_span") @@ -324,7 +369,7 @@ async def _execute_final_output( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent, output=final_output + span=span, usage=context_wrapper.usage, agent=agent, output=final_output ) span.__exit__(None, None, None) delattr(context_wrapper, "_sentry_agent_span") diff --git a/sentry_sdk/integrations/openai_agents/patches/runner.py b/sentry_sdk/integrations/openai_agents/patches/runner.py index f091c5dfb9..2befe08e03 100644 --- a/sentry_sdk/integrations/openai_agents/patches/runner.py +++ b/sentry_sdk/integrations/openai_agents/patches/runner.py @@ -11,13 +11,15 @@ from ..spans import ( agent_workflow_span, execute_tool_span, + handoff_span, + invoke_agent_span, update_execute_tool_span, update_invoke_agent_span, ) from ..utils import _capture_exception try: - from agents import FunctionTool, RunHooks + from agents import FunctionTool, RunContextWrapper, RunHooks from agents.exceptions import AgentsException except ImportError: raise DidNotEnable("OpenAI Agents not installed") @@ -27,7 +29,7 @@ if TYPE_CHECKING: from typing import Any, AsyncIterator, Callable - from agents import Agent, Tool, ToolContext + from agents import Agent, AgentHookContext, Tool, ToolContext TContext = TypeVar("TContext") @@ -35,11 +37,66 @@ class _SentryRunHooks(RunHooks[TContext]): # type: ignore[misc] """ - Responsible for creating and managing Execute Tool spans. These spans are - stored on the ToolContext reference that is shared between `on_tool_start()` - and `on_tool_end()` + Responsible for creating Execute Tool and Invoke Agent spans. Both + spans can also be closed by hooks, and Execute Tool spans are exclusively + closed in the hook. + + Execute Tool spans are stored on the ToolContext reference that is shared + between `on_tool_start()` and `on_tool_end()`. + + Invoke Agent spans are stored as an instance variable on the hooks + instance. This works because agent invocations in a run are always + sequential. + + The `on_agent_end()` hook only fires once in a given run. + In a handoff, `on_handoff()` and `on_agent_start()` with the new + agent are invoked. The `on_handoff()` hook therefore finishes the + span corresponding to the `from_agent`. """ + async def on_agent_start( + self, + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + ) -> "None": + self._sentry_invoke_agent_span = invoke_agent_span(agent, context.turn_input) + + async def on_agent_end( + self, + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + output: "Any", + ) -> "None": + span = getattr(self, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, + usage=context.usage, + agent=agent, + output=output, + ) + del self._sentry_invoke_agent_span + span.__exit__(None, None, None) + + async def on_handoff( + self, + context: "RunContextWrapper[TContext]", + from_agent: "Agent[TContext]", + to_agent: "Agent[TContext]", + ) -> "None": + span = getattr(self, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, + usage=context.usage, + agent=from_agent, + output=None, + ) + del self._sentry_invoke_agent_span + span.__exit__(None, None, None) + + handoff_span(context, from_agent, to_agent.name) + async def on_tool_start( self, context: "ToolContext[TContext]", @@ -88,11 +145,44 @@ def _patch_run_hooks(hooks: "RunHooks[TContext]") -> None: if is_already_patched: return + original_on_agent_start = hooks.on_agent_start + original_on_agent_end = hooks.on_agent_end + original_on_handoff = hooks.on_handoff + original_on_tool_start = hooks.on_tool_start original_on_tool_end = hooks.on_tool_end sentry_hooks = _SentryRunHooks() # type: ignore[var-annotated] + @wraps(original_on_agent_start) + async def on_agent_start( + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + ) -> "None": + with capture_internal_exceptions(): + await sentry_hooks.on_agent_start(context, agent) + await original_on_agent_start(context, agent) + + @wraps(original_on_agent_end) + async def on_agent_end( + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + output: "Any", + ) -> "None": + with capture_internal_exceptions(): + await sentry_hooks.on_agent_end(context, agent, output) + await original_on_agent_end(context, agent, output) + + @wraps(original_on_handoff) + async def on_handoff( + context: "RunContextWrapper[TContext]", + from_agent: "Agent[TContext]", + to_agent: "Agent[TContext]", + ) -> "None": + with capture_internal_exceptions(): + await sentry_hooks.on_handoff(context, from_agent, to_agent) + await original_on_handoff(context, from_agent, to_agent) + @wraps(original_on_tool_start) async def on_tool_start( context: "ToolContext[TContext]", agent: "Agent[TContext]", tool: "Tool" @@ -113,6 +203,11 @@ async def on_tool_end( await original_on_tool_end(context, agent, tool, result) hooks._sentry_is_patched = True + + hooks.on_agent_start = on_agent_start + hooks.on_agent_end = on_agent_end + hooks.on_handoff = on_handoff + hooks.on_tool_start = on_tool_start hooks.on_tool_end = on_tool_end @@ -136,7 +231,8 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": if hooks is not None: _patch_run_hooks(hooks=hooks) else: - kwargs["hooks"] = _SentryRunHooks() + hooks = _SentryRunHooks() + kwargs["hooks"] = hooks # Isolate each workflow so that when agents are run in asyncio tasks they # don't touch each other's scopes @@ -175,7 +271,20 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": _capture_exception(exc) context_wrapper = getattr(exc.run_data, "context_wrapper", None) - if context_wrapper is not None: + if context_wrapper is not None and use_run_hooks: + invoke_agent_span = getattr( + hooks, "_sentry_invoke_agent_span", None + ) + + if invoke_agent_span is not None: + update_invoke_agent_span( + span=invoke_agent_span, + usage=context_wrapper.usage, + agent=agent, + ) + del hooks._sentry_invoke_agent_span + invoke_agent_span.__exit__(*exc_info) + elif context_wrapper is not None: invoke_agent_span = getattr( context_wrapper, "_sentry_agent_span", None ) @@ -192,7 +301,7 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": ): update_invoke_agent_span( span=invoke_agent_span, - context=context_wrapper, + usage=context_wrapper.usage, agent=agent, ) @@ -208,20 +317,21 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": _capture_exception(exc) reraise(*exc_info) - invoke_agent_span = getattr( - run_result.context_wrapper, "_sentry_agent_span", None - ) - if not invoke_agent_span: - return run_result - - update_invoke_agent_span( - span=invoke_agent_span, - context=run_result.context_wrapper, - agent=agent, - ) - - invoke_agent_span.__exit__(None, None, None) - delattr(run_result.context_wrapper, "_sentry_agent_span") + if not use_run_hooks: + invoke_agent_span = getattr( + run_result.context_wrapper, "_sentry_agent_span", None + ) + if not invoke_agent_span: + return run_result + + update_invoke_agent_span( + span=invoke_agent_span, + usage=run_result.context_wrapper.usage, + agent=agent, + ) + + invoke_agent_span.__exit__(None, None, None) + delattr(run_result.context_wrapper, "_sentry_agent_span") return run_result return wrapper diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py index c21145ac4a..06a8fe0c88 100644 --- a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py +++ b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py @@ -20,13 +20,15 @@ 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 + from agents import TResponseInputItem + from agents.usage import Usage def invoke_agent_span( - context: "agents.RunContextWrapper", agent: "agents.Agent", kwargs: "dict[str, Any]" + agent: "agents.Agent", turn_input: "Optional[list[TResponseInputItem]]" ) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: @@ -64,12 +66,11 @@ def invoke_agent_span( } ) - original_input = kwargs.get("original_input") - if original_input is not None: + if turn_input is not None: message = ( - original_input - if isinstance(original_input, str) - else safe_serialize(original_input) + turn_input + if isinstance(turn_input, str) + else safe_serialize(turn_input) ) messages.append( { @@ -102,13 +103,12 @@ def invoke_agent_span( def update_invoke_agent_span( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - context: "agents.RunContextWrapper", + usage: "Optional[Usage]", agent: "agents.Agent", output: "Any" = None, ) -> None: - # Add aggregated usage data from context_wrapper - if hasattr(context, "usage"): - _set_usage_data(span, context.usage) + if usage is not None: + _set_usage_data(span, usage) if should_send_default_pii(): set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False) diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py index 757808f192..cfc45c4703 100644 --- a/tests/integrations/openai_agents/test_openai_agents.py +++ b/tests/integrations/openai_agents/test_openai_agents.py @@ -1938,6 +1938,14 @@ async def test_handoff_span( sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] + + agent_spans = [ + span + for span in spans + if span["attributes"].get("sentry.op") == OP.GEN_AI_INVOKE_AGENT + ] + assert len(agent_spans) == 2 + handoff_span = next( span for span in spans @@ -1971,6 +1979,12 @@ async def test_handoff_span( (transaction,) = events spans = transaction["spans"] + + agent_spans = [ + span for span in spans if span.get("op") == OP.GEN_AI_INVOKE_AGENT + ] + assert len(agent_spans) == 2 + handoff_span = next( span for span in spans if span.get("op") == OP.GEN_AI_HANDOFF )