openai-agents 0.2.8__py3-none-any.whl → 0.6.8__py3-none-any.whl

agents/tool_guardrails.py

@@ -0,0 +1,279 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, overload
+
+from typing_extensions import TypedDict, TypeVar
+
+from .exceptions import UserError
+from .tool_context import ToolContext
+from .util._types import MaybeAwaitable
+
+if TYPE_CHECKING:
+    from .agent import Agent
+
+
+@dataclass
+class ToolInputGuardrailResult:
+    """The result of a tool input guardrail run."""
+
+    guardrail: ToolInputGuardrail[Any]
+    """The guardrail that was run."""
+
+    output: ToolGuardrailFunctionOutput
+    """The output of the guardrail function."""
+
+
+@dataclass
+class ToolOutputGuardrailResult:
+    """The result of a tool output guardrail run."""
+
+    guardrail: ToolOutputGuardrail[Any]
+    """The guardrail that was run."""
+
+    output: ToolGuardrailFunctionOutput
+    """The output of the guardrail function."""
+
+
+class RejectContentBehavior(TypedDict):
+    """Rejects the tool call/output but continues execution with a message to the model."""
+
+    type: Literal["reject_content"]
+    message: str
+
+
+class RaiseExceptionBehavior(TypedDict):
+    """Raises an exception to halt execution."""
+
+    type: Literal["raise_exception"]
+
+
+class AllowBehavior(TypedDict):
+    """Allows normal tool execution to continue."""
+
+    type: Literal["allow"]
+
+
+@dataclass
+class ToolGuardrailFunctionOutput:
+    """The output of a tool guardrail function."""
+
+    output_info: Any
+    """
+    Optional data about checks performed. For example, the guardrail could include
+    information about the checks it performed and granular results.
+    """
+
+    behavior: RejectContentBehavior | RaiseExceptionBehavior | AllowBehavior = field(
+        default_factory=lambda: AllowBehavior(type="allow")
+    )
+    """
+    Defines how the system should respond when this guardrail result is processed.
+    - allow: Allow normal tool execution to continue without interference (default)
+    - reject_content: Reject the tool call/output but continue execution with a message to the model
+    - raise_exception: Halt execution by raising a ToolGuardrailTripwireTriggered exception
+    """
+
+    @classmethod
+    def allow(cls, output_info: Any = None) -> ToolGuardrailFunctionOutput:
+        """Create a guardrail output that allows the tool execution to continue normally.
+
+        Args:
+            output_info: Optional data about checks performed.
+
+        Returns:
+            ToolGuardrailFunctionOutput configured to allow normal execution.
+        """
+        return cls(output_info=output_info, behavior=AllowBehavior(type="allow"))
+
+    @classmethod
+    def reject_content(cls, message: str, output_info: Any = None) -> ToolGuardrailFunctionOutput:
+        """Create a guardrail output that rejects the tool call/output but continues execution.
+
+        Args:
+            message: Message to send to the model instead of the tool result.
+            output_info: Optional data about checks performed.
+
+        Returns:
+            ToolGuardrailFunctionOutput configured to reject the content.
+        """
+        return cls(
+            output_info=output_info,
+            behavior=RejectContentBehavior(type="reject_content", message=message),
+        )
+
+    @classmethod
+    def raise_exception(cls, output_info: Any = None) -> ToolGuardrailFunctionOutput:
+        """Create a guardrail output that raises an exception to halt execution.
+
+        Args:
+            output_info: Optional data about checks performed.
+
+        Returns:
+            ToolGuardrailFunctionOutput configured to raise an exception.
+        """
+        return cls(output_info=output_info, behavior=RaiseExceptionBehavior(type="raise_exception"))
+
+
+@dataclass
+class ToolInputGuardrailData:
+    """Input data passed to a tool input guardrail function."""
+
+    context: ToolContext[Any]
+    """
+    The tool context containing information about the current tool execution.
+    """
+
+    agent: Agent[Any]
+    """
+    The agent that is executing the tool.
+    """
+
+
+@dataclass
+class ToolOutputGuardrailData(ToolInputGuardrailData):
+    """Input data passed to a tool output guardrail function.
+
+    Extends input data with the tool's output.
+    """
+
+    output: Any
+    """
+    The output produced by the tool function.
+    """
+
+
+TContext_co = TypeVar("TContext_co", bound=Any, covariant=True)
+
+
+@dataclass
+class ToolInputGuardrail(Generic[TContext_co]):
+    """A guardrail that runs before a function tool is invoked."""
+
+    guardrail_function: Callable[
+        [ToolInputGuardrailData], MaybeAwaitable[ToolGuardrailFunctionOutput]
+    ]
+    """
+    The function that implements the guardrail logic.
+    """
+
+    name: str | None = None
+    """
+    Optional name for the guardrail. If not provided, uses the function name.
+    """
+
+    def get_name(self) -> str:
+        return self.name or self.guardrail_function.__name__
+
+    async def run(self, data: ToolInputGuardrailData) -> ToolGuardrailFunctionOutput:
+        if not callable(self.guardrail_function):
+            raise UserError(f"Guardrail function must be callable, got {self.guardrail_function}")
+
+        result = self.guardrail_function(data)
+        if inspect.isawaitable(result):
+            return await result
+        return result
+
+
+@dataclass
+class ToolOutputGuardrail(Generic[TContext_co]):
+    """A guardrail that runs after a function tool is invoked."""
+
+    guardrail_function: Callable[
+        [ToolOutputGuardrailData], MaybeAwaitable[ToolGuardrailFunctionOutput]
+    ]
+    """
+    The function that implements the guardrail logic.
+    """
+
+    name: str | None = None
+    """
+    Optional name for the guardrail. If not provided, uses the function name.
+    """
+
+    def get_name(self) -> str:
+        return self.name or self.guardrail_function.__name__
+
+    async def run(self, data: ToolOutputGuardrailData) -> ToolGuardrailFunctionOutput:
+        if not callable(self.guardrail_function):
+            raise UserError(f"Guardrail function must be callable, got {self.guardrail_function}")
+
+        result = self.guardrail_function(data)
+        if inspect.isawaitable(result):
+            return await result
+        return result
+
+
+# Decorators
+_ToolInputFuncSync = Callable[[ToolInputGuardrailData], ToolGuardrailFunctionOutput]
+_ToolInputFuncAsync = Callable[[ToolInputGuardrailData], Awaitable[ToolGuardrailFunctionOutput]]
+
+
+@overload
+def tool_input_guardrail(func: _ToolInputFuncSync): ...
+
+
+@overload
+def tool_input_guardrail(func: _ToolInputFuncAsync): ...
+
+
+@overload
+def tool_input_guardrail(
+    *, name: str | None = None
+) -> Callable[[_ToolInputFuncSync | _ToolInputFuncAsync], ToolInputGuardrail[Any]]: ...
+
+
+def tool_input_guardrail(
+    func: _ToolInputFuncSync | _ToolInputFuncAsync | None = None,
+    *,
+    name: str | None = None,
+) -> (
+    ToolInputGuardrail[Any]
+    | Callable[[_ToolInputFuncSync | _ToolInputFuncAsync], ToolInputGuardrail[Any]]
+):
+    """Decorator to create a ToolInputGuardrail from a function."""
+
+    def decorator(f: _ToolInputFuncSync | _ToolInputFuncAsync) -> ToolInputGuardrail[Any]:
+        return ToolInputGuardrail(guardrail_function=f, name=name or f.__name__)
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+_ToolOutputFuncSync = Callable[[ToolOutputGuardrailData], ToolGuardrailFunctionOutput]
+_ToolOutputFuncAsync = Callable[[ToolOutputGuardrailData], Awaitable[ToolGuardrailFunctionOutput]]
+
+
+@overload
+def tool_output_guardrail(func: _ToolOutputFuncSync): ...
+
+
+@overload
+def tool_output_guardrail(func: _ToolOutputFuncAsync): ...
+
+
+@overload
+def tool_output_guardrail(
+    *, name: str | None = None
+) -> Callable[[_ToolOutputFuncSync | _ToolOutputFuncAsync], ToolOutputGuardrail[Any]]: ...
+
+
+def tool_output_guardrail(
+    func: _ToolOutputFuncSync | _ToolOutputFuncAsync | None = None,
+    *,
+    name: str | None = None,
+) -> (
+    ToolOutputGuardrail[Any]
+    | Callable[[_ToolOutputFuncSync | _ToolOutputFuncAsync], ToolOutputGuardrail[Any]]
+):
+    """Decorator to create a ToolOutputGuardrail from a function."""
+
+    def decorator(f: _ToolOutputFuncSync | _ToolOutputFuncAsync) -> ToolOutputGuardrail[Any]:
+        return ToolOutputGuardrail(guardrail_function=f, name=name or f.__name__)
+
+    if func is not None:
+        return decorator(func)
+    return decorator

agents/tracing/__init__.py

@@ -1,5 +1,6 @@
 import atexit
 
+from .config import TracingConfig
 from .create import (
     agent_span,
     custom_span,
@@ -53,6 +54,7 @@ __all__ = [
     "set_trace_processors",
     "set_trace_provider",
     "set_tracing_disabled",
+    "TracingConfig",
     "trace",
     "Trace",
     "SpanError",

agents/tracing/config.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from typing_extensions import TypedDict
+
+
+class TracingConfig(TypedDict, total=False):
+    """Configuration for tracing export."""
+
+    api_key: str

agents/tracing/create.py

@@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
 from ..logger import logger
+from .config import TracingConfig
 from .setup import get_trace_provider
 from .span_data import (
     AgentSpanData,
@@ -30,6 +31,7 @@ def trace(
     trace_id: str | None = None,
     group_id: str | None = None,
     metadata: dict[str, Any] | None = None,
+    tracing: TracingConfig | None = None,
     disabled: bool = False,
 ) -> Trace:
     """
@@ -50,6 +52,7 @@ def trace(
         group_id: Optional grouping identifier to link multiple traces from the same conversation
             or process. For instance, you might use a chat thread ID.
         metadata: Optional dictionary of additional metadata to attach to the trace.
+        tracing: Optional tracing configuration for exporting this trace.
         disabled: If True, we will return a Trace but the Trace will not be recorded.
 
     Returns:
@@ -66,6 +69,7 @@ def trace(
         trace_id=trace_id,
         group_id=group_id,
         metadata=metadata,
+        tracing=tracing,
         disabled=disabled,
     )
 

agents/tracing/processor_interface.py

@@ -7,52 +7,125 @@ if TYPE_CHECKING:
 
 
 class TracingProcessor(abc.ABC):
-    """Interface for processing spans."""
+    """Interface for processing and monitoring traces and spans in the OpenAI Agents system.
+
+    This abstract class defines the interface that all tracing processors must implement.
+    Processors receive notifications when traces and spans start and end, allowing them
+    to collect, process, and export tracing data.
+
+    Example:
+        ```python
+        class CustomProcessor(TracingProcessor):
+            def __init__(self):
+                self.active_traces = {}
+                self.active_spans = {}
+
+            def on_trace_start(self, trace):
+                self.active_traces[trace.trace_id] = trace
+
+            def on_trace_end(self, trace):
+                # Process completed trace
+                del self.active_traces[trace.trace_id]
+
+            def on_span_start(self, span):
+                self.active_spans[span.span_id] = span
+
+            def on_span_end(self, span):
+                # Process completed span
+                del self.active_spans[span.span_id]
+
+            def shutdown(self):
+                # Clean up resources
+                self.active_traces.clear()
+                self.active_spans.clear()
+
+            def force_flush(self):
+                # Force processing of any queued items
+                pass
+        ```
+
+    Notes:
+        - All methods should be thread-safe
+        - Methods should not block for long periods
+        - Handle errors gracefully to prevent disrupting agent execution
+    """
 
     @abc.abstractmethod
     def on_trace_start(self, trace: "Trace") -> None:
-        """Called when a trace is started.
+        """Called when a new trace begins execution.
 
         Args:
-            trace: The trace that started.
+            trace: The trace that started. Contains workflow name and metadata.
+
+        Notes:
+            - Called synchronously on trace start
+            - Should return quickly to avoid blocking execution
+            - Any errors should be caught and handled internally
         """
         pass
 
     @abc.abstractmethod
     def on_trace_end(self, trace: "Trace") -> None:
-        """Called when a trace is finished.
+        """Called when a trace completes execution.
 
         Args:
-            trace: The trace that finished.
+            trace: The completed trace containing all spans and results.
+
+        Notes:
+            - Called synchronously when trace finishes
+            - Good time to export/process the complete trace
+            - Should handle cleanup of any trace-specific resources
         """
         pass
 
     @abc.abstractmethod
     def on_span_start(self, span: "Span[Any]") -> None:
-        """Called when a span is started.
+        """Called when a new span begins execution.
 
         Args:
-            span: The span that started.
+            span: The span that started. Contains operation details and context.
+
+        Notes:
+            - Called synchronously on span start
+            - Should return quickly to avoid blocking execution
+            - Spans are automatically nested under current trace/span
         """
         pass
 
     @abc.abstractmethod
     def on_span_end(self, span: "Span[Any]") -> None:
-        """Called when a span is finished. Should not block or raise exceptions.
+        """Called when a span completes execution.
 
         Args:
-            span: The span that finished.
+            span: The completed span containing execution results.
+
+        Notes:
+            - Called synchronously when span finishes
+            - Should not block or raise exceptions
+            - Good time to export/process the individual span
         """
         pass
 
     @abc.abstractmethod
     def shutdown(self) -> None:
-        """Called when the application stops."""
+        """Called when the application stops to clean up resources.
+
+        Should perform any necessary cleanup like:
+        - Flushing queued traces/spans
+        - Closing connections
+        - Releasing resources
+        """
         pass
 
     @abc.abstractmethod
     def force_flush(self) -> None:
-        """Forces an immediate flush of all queued spans/traces."""
+        """Forces immediate processing of any queued traces/spans.
+
+        Notes:
+            - Should process all queued items before returning
+            - Useful before shutdown or when immediate processing is needed
+            - May block while processing completes
+        """
         pass
 
 

agents/tracing/processors.py

@@ -70,8 +70,8 @@ class BackendSpanExporter(TracingExporter):
                 client.
         """
         # Clear the cached property if it exists
-        if 'api_key' in self.__dict__:
-            del self.__dict__['api_key']
+        if "api_key" in self.__dict__:
+            del self.__dict__["api_key"]
 
         # Update the private attribute
         self._api_key = api_key
@@ -92,62 +92,73 @@ class BackendSpanExporter(TracingExporter):
         if not items:
             return
 
-        if not self.api_key:
-            logger.warning("OPENAI_API_KEY is not set, skipping trace export")
-            return
-
-        data = [item.export() for item in items if item.export()]
-        payload = {"data": data}
-
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json",
-            "OpenAI-Beta": "traces=v1",
-        }
-
-        if self.organization:
-            headers["OpenAI-Organization"] = self.organization
-
-        if self.project:
-            headers["OpenAI-Project"] = self.project
-
-        # Exponential backoff loop
-        attempt = 0
-        delay = self.base_delay
-        while True:
-            attempt += 1
-            try:
-                response = self._client.post(url=self.endpoint, headers=headers, json=payload)
-
-                # If the response is successful, break out of the loop
-                if response.status_code < 300:
-                    logger.debug(f"Exported {len(items)} items")
-                    return
+        grouped_items: dict[str | None, list[Trace | Span[Any]]] = {}
+        for item in items:
+            key = item.tracing_api_key
+            grouped_items.setdefault(key, []).append(item)
+
+        for item_key, grouped in grouped_items.items():
+            api_key = item_key or self.api_key
+            if not api_key:
+                logger.warning("OPENAI_API_KEY is not set, skipping trace export")
+                continue
+
+            data = [item.export() for item in grouped if item.export()]
+            payload = {"data": data}
+
+            headers = {
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+                "OpenAI-Beta": "traces=v1",
+            }
+
+            if self.organization:
+                headers["OpenAI-Organization"] = self.organization
+
+            if self.project:
+                headers["OpenAI-Project"] = self.project
+
+            # Exponential backoff loop
+            attempt = 0
+            delay = self.base_delay
+            while True:
+                attempt += 1
+                try:
+                    response = self._client.post(url=self.endpoint, headers=headers, json=payload)
+
+                    # If the response is successful, break out of the loop
+                    if response.status_code < 300:
+                        logger.debug(f"Exported {len(grouped)} items")
+                        break
+
+                    # If the response is a client error (4xx), we won't retry
+                    if 400 <= response.status_code < 500:
+                        logger.error(
+                            "[non-fatal] Tracing client error %s: %s",
+                            response.status_code,
+                            response.text,
+                        )
+                        break
+
+                    # For 5xx or other unexpected codes, treat it as transient and retry
+                    logger.warning(
+                        f"[non-fatal] Tracing: server error {response.status_code}, retrying."
+                    )
+                except httpx.RequestError as exc:
+                    # Network or other I/O error, we'll retry
+                    logger.warning(f"[non-fatal] Tracing: request failed: {exc}")
 
-                # If the response is a client error (4xx), we won't retry
-                if 400 <= response.status_code < 500:
+                # If we reach here, we need to retry or give up
+                if attempt >= self.max_retries:
                     logger.error(
-                        f"[non-fatal] Tracing client error {response.status_code}: {response.text}"
+                        "[non-fatal] Tracing: max retries reached, giving up on this batch."
                     )
-                    return
-
-                # For 5xx or other unexpected codes, treat it as transient and retry
-                logger.warning(
-                    f"[non-fatal] Tracing: server error {response.status_code}, retrying."
-                )
-            except httpx.RequestError as exc:
-                # Network or other I/O error, we'll retry
-                logger.warning(f"[non-fatal] Tracing: request failed: {exc}")
-
-            # If we reach here, we need to retry or give up
-            if attempt >= self.max_retries:
-                logger.error("[non-fatal] Tracing: max retries reached, giving up on this batch.")
-                return
+                    break
 
-            # Exponential backoff + jitter
-            sleep_time = delay + random.uniform(0, 0.1 * delay)  # 10% jitter
-            time.sleep(sleep_time)
-            delay = min(delay * 2, self.max_delay)
+                # Exponential backoff + jitter
+                sleep_time = delay + random.uniform(0, 0.1 * delay)  # 10% jitter
+                time.sleep(sleep_time)
+                delay = min(delay * 2, self.max_delay)
 
     def close(self):
         """Close the underlying HTTP client."""