feat(openai-agents) Add guardrail spans

sentry_sdk/consts.py

@@ -467,6 +467,12 @@ class SPANDATA:
     Example: "The weather in Paris is rainy and overcast, with temperatures around 57°F"
     """
 
+    GEN_AI_GUARDRAIL_TRIPWIRE_TRIGGERED = "gen_ai.guardrail.tripwire_triggered"
+    """
+    Whether the guardrail tripwire was triggered.
+    Example: true
+    """
+
     GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
     """
     The name of the operation being performed.
@@ -798,6 +804,8 @@ class OP:
     GEN_AI_EMBEDDINGS = "gen_ai.embeddings"
     GEN_AI_EXECUTE_TOOL = "gen_ai.execute_tool"
     GEN_AI_GENERATE_TEXT = "gen_ai.generate_text"
+    GEN_AI_GUARDRAIL_INPUT = "gen_ai.guardrail.input"
+    GEN_AI_GUARDRAIL_OUTPUT = "gen_ai.guardrail.output"
     GEN_AI_HANDOFF = "gen_ai.handoff"
     GEN_AI_PIPELINE = "gen_ai.pipeline"
     GEN_AI_INVOKE_AGENT = "gen_ai.invoke_agent"

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -1,6 +1,10 @@
 from functools import wraps
 
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.integrations.openai_agents.spans.guardrail import (
+    guardrail_span,
+    update_guardrail_span,
+)
 from ..spans import invoke_agent_span, update_invoke_agent_span, handoff_span
 
 from typing import TYPE_CHECKING
@@ -26,6 +30,13 @@ def _patch_agent_run():
     original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
     original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
 
+    original_run_single_input_guardrail = (
+        agents._run_impl.RunImpl.run_single_input_guardrail
+    )
+    original_run_single_output_guardrail = (
+        agents._run_impl.RunImpl.run_single_output_guardrail
+    )
+
     def _start_invoke_agent_span(context_wrapper, agent, kwargs):
         # type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> None
         """Start an agent invocation span"""
@@ -132,9 +143,47 @@ async def patched_execute_final_output(cls, *args, **kwargs):
 
         return result
 
+    @wraps(
+        original_run_single_input_guardrail.__func__
+        if hasattr(original_run_single_input_guardrail, "__func__")
+        else original_run_single_input_guardrail
+    )
+    async def patched_run_single_input_guardrail(cls, *args, **kwargs):
+        # type: (agents.Runner, *Any, **Any) -> Any
+        agent = args[0]
+        guardrail = args[1]
+
+        with guardrail_span(guardrail, "input", args, kwargs) as span:
+            result = await original_run_single_input_guardrail(*args, **kwargs)
+            update_guardrail_span(span, agent, guardrail, "input", result)
+
+        return result
+
+    @wraps(
+        original_run_single_output_guardrail.__func__
+        if hasattr(original_run_single_output_guardrail, "__func__")
+        else original_run_single_output_guardrail
+    )
+    async def patched_run_single_output_guardrail(cls, *args, **kwargs):
+        # type: (agents.Runner, *Any, **Any) -> Any
+        guardrail = args[0]
+        agent = args[1]
+
+        with guardrail_span(guardrail, "output", args, kwargs) as span:
+            result = await original_run_single_output_guardrail(*args, **kwargs)
+            update_guardrail_span(span, agent, guardrail, "output", result)
+
+        return result
+
     # Apply patches
     agents.run.AgentRunner._run_single_turn = classmethod(patched_run_single_turn)
     agents._run_impl.RunImpl.execute_handoffs = classmethod(patched_execute_handoffs)
     agents._run_impl.RunImpl.execute_final_output = classmethod(
         patched_execute_final_output
     )
+    agents._run_impl.RunImpl.run_single_input_guardrail = classmethod(
+        patched_run_single_input_guardrail
+    )
+    agents._run_impl.RunImpl.run_single_output_guardrail = classmethod(
+        patched_run_single_output_guardrail
+    )

sentry_sdk/integrations/openai_agents/spans/guardrail.py

@@ -0,0 +1,56 @@
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.scope import should_send_default_pii
+
+from ..consts import SPAN_ORIGIN
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import agents
+    from typing import Any
+
+
+def guardrail_span(guardrail, guardrail_type, args, kwargs):
+    # type: (agents.Guardrail, str, tuple[Any, ...], dict[str, Any]) -> sentry_sdk.tracing.Span
+    op = (
+        OP.GEN_AI_GUARDRAIL_OUTPUT
+        if guardrail_type == "output"
+        else OP.GEN_AI_GUARDRAIL_INPUT
+    )
+
+    span = sentry_sdk.start_span(
+        op=op,
+        name=f"guardrail {guardrail.name or ''}".strip(),
+        origin=SPAN_ORIGIN,
+    )
+
+    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "guardrail")
+    span.set_data(SPANDATA.GEN_AI_TOOL_TYPE, f"guardrail.{guardrail_type}")
+
+    if guardrail.name is not None:
+        span.set_data(SPANDATA.GEN_AI_TOOL_NAME, guardrail.name)
+
+    try:
+        input = args[2]
+    except IndexError:
+        input = None
+
+    if should_send_default_pii() and input is not None:
+        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, input)
+
+    return span
+
+
+def update_guardrail_span(span, agent, guardrail, guardrail_type, result):
+    # type: (sentry_sdk.tracing.Span, agents.Agent, agents.Guardrail, str, Any) -> None
+    if agent.name is not None:
+        span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent.name)
+
+    output = result.output.output_info.get("reason")
+    if should_send_default_pii() and output is not None:
+        span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, output)
+
+    tripwire_triggered = result.output.tripwire_triggered
+    if tripwire_triggered is not None:
+        span.set_data(SPANDATA.GEN_AI_GUARDRAIL_TRIPWIRE_TRIGGERED, tripwire_triggered)