openai-agents 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

agents/tracing/create.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from .logger import logger
+from .setup import GLOBAL_TRACE_PROVIDER
+from .span_data import (
+    AgentSpanData,
+    CustomSpanData,
+    FunctionSpanData,
+    GenerationSpanData,
+    GuardrailSpanData,
+    HandoffSpanData,
+    ResponseSpanData,
+)
+from .spans import Span
+from .traces import Trace
+
+if TYPE_CHECKING:
+    from openai.types.responses import Response
+
+
+def trace(
+    workflow_name: str,
+    trace_id: str | None = None,
+    group_id: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    disabled: bool = False,
+) -> Trace:
+    """
+    Create a new trace. The trace will not be started automatically; you should either use
+    it as a context manager (`with trace(...):`) or call `trace.start()` + `trace.finish()`
+    manually.
+
+    In addition to the workflow name and optional grouping identifier, you can provide
+    an arbitrary metadata dictionary to attach additional user-defined information to
+    the trace.
+
+    Args:
+        workflow_name: The name of the logical app or workflow. For example, you might provide
+            "code_bot" for a coding agent, or "customer_support_agent" for a customer support agent.
+        trace_id: The ID of the trace. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_trace_id()` to generate a trace ID, to guarantee that IDs are
+            correctly formatted.
+        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.
+        disabled: If True, we will return a Trace but the Trace will not be recorded. This will
+            not be checked if there's an existing trace and `even_if_trace_running` is True.
+
+    Returns:
+        The newly created trace object.
+    """
+    current_trace = GLOBAL_TRACE_PROVIDER.get_current_trace()
+    if current_trace:
+        logger.warning(
+            "Trace already exists. Creating a new trace, but this is probably a mistake."
+        )
+
+    return GLOBAL_TRACE_PROVIDER.create_trace(
+        name=workflow_name,
+        trace_id=trace_id,
+        group_id=group_id,
+        metadata=metadata,
+        disabled=disabled,
+    )
+
+
+def get_current_trace() -> Trace | None:
+    """Returns the currently active trace, if present."""
+    return GLOBAL_TRACE_PROVIDER.get_current_trace()
+
+
+def get_current_span() -> Span[Any] | None:
+    """Returns the currently active span, if present."""
+    return GLOBAL_TRACE_PROVIDER.get_current_span()
+
+
+def agent_span(
+    name: str,
+    handoffs: list[str] | None = None,
+    tools: list[str] | None = None,
+    output_type: str | None = None,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[AgentSpanData]:
+    """Create a new agent span. The span will not be started automatically, you should either do
+    `with agent_span() ...` or call `span.start()` + `span.finish()` manually.
+
+    Args:
+        name: The name of the agent.
+        handoffs: Optional list of agent names to which this agent could hand off control.
+        tools: Optional list of tool names available to this agent.
+        output_type: Optional name of the output type produced by the agent.
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+
+    Returns:
+        The newly created agent span.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=AgentSpanData(name=name, handoffs=handoffs, tools=tools, output_type=output_type),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )
+
+
+def function_span(
+    name: str,
+    input: str | None = None,
+    output: str | None = None,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[FunctionSpanData]:
+    """Create a new function span. The span will not be started automatically, you should either do
+    `with function_span() ...` or call `span.start()` + `span.finish()` manually.
+
+    Args:
+        name: The name of the function.
+        input: The input to the function.
+        output: The output of the function.
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+
+    Returns:
+        The newly created function span.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=FunctionSpanData(name=name, input=input, output=output),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )
+
+
+def generation_span(
+    input: Sequence[Mapping[str, Any]] | None = None,
+    output: Sequence[Mapping[str, Any]] | None = None,
+    model: str | None = None,
+    model_config: Mapping[str, Any] | None = None,
+    usage: dict[str, Any] | None = None,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[GenerationSpanData]:
+    """Create a new generation span. The span will not be started automatically, you should either
+    do `with generation_span() ...` or call `span.start()` + `span.finish()` manually.
+
+    This span captures the details of a model generation, including the
+    input message sequence, any generated outputs, the model name and
+    configuration, and usage data. If you only need to capture a model
+    response identifier, use `response_span()` instead.
+
+    Args:
+        input: The sequence of input messages sent to the model.
+        output: The sequence of output messages received from the model.
+        model: The model identifier used for the generation.
+        model_config: The model configuration (hyperparameters) used.
+        usage: A dictionary of usage information (input tokens, output tokens, etc.).
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+
+    Returns:
+        The newly created generation span.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=GenerationSpanData(
+            input=input, output=output, model=model, model_config=model_config, usage=usage
+        ),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )
+
+
+def response_span(
+    response: Response | None = None,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[ResponseSpanData]:
+    """Create a new response span. The span will not be started automatically, you should either do
+    `with response_span() ...` or call `span.start()` + `span.finish()` manually.
+
+    Args:
+        response: The OpenAI Response object.
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=ResponseSpanData(response=response),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )
+
+
+def handoff_span(
+    from_agent: str | None = None,
+    to_agent: str | None = None,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[HandoffSpanData]:
+    """Create a new handoff span. The span will not be started automatically, you should either do
+    `with handoff_span() ...` or call `span.start()` + `span.finish()` manually.
+
+    Args:
+        from_agent: The name of the agent that is handing off.
+        to_agent: The name of the agent that is receiving the handoff.
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+
+    Returns:
+        The newly created handoff span.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=HandoffSpanData(from_agent=from_agent, to_agent=to_agent),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )
+
+
+def custom_span(
+    name: str,
+    data: dict[str, Any] | None = None,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[CustomSpanData]:
+    """Create a new custom span, to which you can add your own metadata. The span will not be
+    started automatically, you should either do `with custom_span() ...` or call
+    `span.start()` + `span.finish()` manually.
+
+    Args:
+        name: The name of the custom span.
+        data: Arbitrary structured data to associate with the span.
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+
+    Returns:
+        The newly created custom span.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=CustomSpanData(name=name, data=data or {}),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )
+
+
+def guardrail_span(
+    name: str,
+    triggered: bool = False,
+    span_id: str | None = None,
+    parent: Trace | Span[Any] | None = None,
+    disabled: bool = False,
+) -> Span[GuardrailSpanData]:
+    """Create a new guardrail span. The span will not be started automatically, you should either
+    do `with guardrail_span() ...` or call `span.start()` + `span.finish()` manually.
+
+    Args:
+        name: The name of the guardrail.
+        triggered: Whether the guardrail was triggered.
+        span_id: The ID of the span. Optional. If not provided, we will generate an ID. We
+            recommend using `util.gen_span_id()` to generate a span ID, to guarantee that IDs are
+            correctly formatted.
+        parent: The parent span or trace. If not provided, we will automatically use the current
+            trace/span as the parent.
+        disabled: If True, we will return a Span but the Span will not be recorded.
+    """
+    return GLOBAL_TRACE_PROVIDER.create_span(
+        span_data=GuardrailSpanData(name=name, triggered=triggered),
+        span_id=span_id,
+        parent=parent,
+        disabled=disabled,
+    )

agents/tracing/logger.py

@@ -0,0 +1,3 @@
+import logging
+
+logger = logging.getLogger("openai.agents.tracing")

agents/tracing/processor_interface.py

@@ -0,0 +1,69 @@
+import abc
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .spans import Span
+    from .traces import Trace
+
+
+class TracingProcessor(abc.ABC):
+    """Interface for processing spans."""
+
+    @abc.abstractmethod
+    def on_trace_start(self, trace: "Trace") -> None:
+        """Called when a trace is started.
+
+        Args:
+            trace: The trace that started.
+        """
+        pass
+
+    @abc.abstractmethod
+    def on_trace_end(self, trace: "Trace") -> None:
+        """Called when a trace is finished.
+
+        Args:
+            trace: The trace that started.
+        """
+        pass
+
+    @abc.abstractmethod
+    def on_span_start(self, span: "Span[Any]") -> None:
+        """Called when a span is started.
+
+        Args:
+            span: The span that started.
+        """
+        pass
+
+    @abc.abstractmethod
+    def on_span_end(self, span: "Span[Any]") -> None:
+        """Called when a span is finished. Should not block or raise exceptions.
+
+        Args:
+            span: The span that finished.
+        """
+        pass
+
+    @abc.abstractmethod
+    def shutdown(self) -> None:
+        """Called when the application stops."""
+        pass
+
+    @abc.abstractmethod
+    def force_flush(self) -> None:
+        """Forces an immediate flush of all queued spans/traces."""
+        pass
+
+
+class TracingExporter(abc.ABC):
+    """Exports traces and spans. For example, could log them or send them to a backend."""
+
+    @abc.abstractmethod
+    def export(self, items: list["Trace | Span[Any]"]) -> None:
+        """Exports a list of traces and spans.
+
+        Args:
+            items: The items to export.
+        """
+        pass

agents/tracing/processors.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import os
+import queue
+import random
+import threading
+import time
+from typing import Any
+
+import httpx
+
+from .logger import logger
+from .processor_interface import TracingExporter, TracingProcessor
+from .spans import Span
+from .traces import Trace
+
+
+class ConsoleSpanExporter(TracingExporter):
+    """Prints the traces and spans to the console."""
+
+    def export(self, items: list[Trace | Span[Any]]) -> None:
+        for item in items:
+            if isinstance(item, Trace):
+                print(f"[Exporter] Export trace_id={item.trace_id}, name={item.name}, ")
+            else:
+                print(f"[Exporter] Export span: {item.export()}")
+
+
+class BackendSpanExporter(TracingExporter):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        organization: str | None = None,
+        project: str | None = None,
+        endpoint: str = "https://api.openai.com/v1/traces/ingest",
+        max_retries: int = 3,
+        base_delay: float = 1.0,
+        max_delay: float = 30.0,
+    ):
+        """
+        Args:
+            api_key: The API key for the "Authorization" header. Defaults to
+                `os.environ["OPENAI_TRACE_API_KEY"]` if not provided.
+            organization: The OpenAI organization to use. Defaults to
+                `os.environ["OPENAI_ORG_ID"]` if not provided.
+            project: The OpenAI project to use. Defaults to
+                `os.environ["OPENAI_PROJECT_ID"]` if not provided.
+            endpoint: The HTTP endpoint to which traces/spans are posted.
+            max_retries: Maximum number of retries upon failures.
+            base_delay: Base delay (in seconds) for the first backoff.
+            max_delay: Maximum delay (in seconds) for backoff growth.
+        """
+        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
+        self.organization = organization or os.environ.get("OPENAI_ORG_ID")
+        self.project = project or os.environ.get("OPENAI_PROJECT_ID")
+        self.endpoint = endpoint
+        self.max_retries = max_retries
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+
+        # Keep a client open for connection pooling across multiple export calls
+        self._client = httpx.Client(timeout=httpx.Timeout(timeout=60, connect=5.0))
+
+    def set_api_key(self, api_key: str):
+        """Set the OpenAI API key for the exporter.
+
+        Args:
+            api_key: The OpenAI API key to use. This is the same key used by the OpenAI Python
+                client.
+        """
+        self.api_key = api_key
+
+    def export(self, items: list[Trace | Span[Any]]) -> None:
+        if not items:
+            return
+
+        if not self.api_key:
+            logger.warning("OPENAI_API_KEY is not set, skipping trace export")
+            return
+
+        traces: list[dict[str, Any]] = []
+        spans: list[dict[str, Any]] = []
+
+        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",
+        }
+
+        # 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(traces)} traces, {len(spans)} spans")
+                    return
+
+                # If the response is a client error (4xx), we wont retry
+                if 400 <= response.status_code < 500:
+                    logger.error(f"Tracing client error {response.status_code}: {response.text}")
+                    return
+
+                # For 5xx or other unexpected codes, treat it as transient and retry
+                logger.warning(f"Server error {response.status_code}, retrying.")
+            except httpx.RequestError as exc:
+                # Network or other I/O error, we'll retry
+                logger.warning(f"Request failed: {exc}")
+
+            # If we reach here, we need to retry or give up
+            if attempt >= self.max_retries:
+                logger.error("Max retries reached, giving up on this batch.")
+                return
+
+            # 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."""
+        self._client.close()
+
+
+class BatchTraceProcessor(TracingProcessor):
+    """Some implementation notes:
+    1. Using Queue, which is thread-safe.
+    2. Using a background thread to export spans, to minimize any performance issues.
+    3. Spans are stored in memory until they are exported.
+    """
+
+    def __init__(
+        self,
+        exporter: TracingExporter,
+        max_queue_size: int = 8192,
+        max_batch_size: int = 128,
+        schedule_delay: float = 5.0,
+        export_trigger_ratio: float = 0.7,
+    ):
+        """
+        Args:
+            exporter: The exporter to use.
+            max_queue_size: The maximum number of spans to store in the queue. After this, we will
+                start dropping spans.
+            max_batch_size: The maximum number of spans to export in a single batch.
+            schedule_delay: The delay between checks for new spans to export.
+            export_trigger_ratio: The ratio of the queue size at which we will trigger an export.
+        """
+        self._exporter = exporter
+        self._queue: queue.Queue[Trace | Span[Any]] = queue.Queue(maxsize=max_queue_size)
+        self._max_queue_size = max_queue_size
+        self._max_batch_size = max_batch_size
+        self._schedule_delay = schedule_delay
+        self._shutdown_event = threading.Event()
+
+        # The queue size threshold at which we export immediately.
+        self._export_trigger_size = int(max_queue_size * export_trigger_ratio)
+
+        # Track when we next *must* perform a scheduled export
+        self._next_export_time = time.time() + self._schedule_delay
+
+        self._shutdown_event = threading.Event()
+        self._worker_thread = threading.Thread(target=self._run, daemon=True)
+        self._worker_thread.start()
+
+    def on_trace_start(self, trace: Trace) -> None:
+        try:
+            self._queue.put_nowait(trace)
+        except queue.Full:
+            logger.warning("Queue is full, dropping trace.")
+
+    def on_trace_end(self, trace: Trace) -> None:
+        # We send traces via on_trace_start, so we don't need to do anything here.
+        pass
+
+    def on_span_start(self, span: Span[Any]) -> None:
+        # We send spans via on_span_end, so we don't need to do anything here.
+        pass
+
+    def on_span_end(self, span: Span[Any]) -> None:
+        try:
+            self._queue.put_nowait(span)
+        except queue.Full:
+            logger.warning("Queue is full, dropping span.")
+
+    def shutdown(self, timeout: float | None = None):
+        """
+        Called when the application stops. We signal our thread to stop, then join it.
+        """
+        self._shutdown_event.set()
+        self._worker_thread.join(timeout=timeout)
+
+    def force_flush(self):
+        """
+        Forces an immediate flush of all queued spans.
+        """
+        self._export_batches(force=True)
+
+    def _run(self):
+        while not self._shutdown_event.is_set():
+            current_time = time.time()
+            queue_size = self._queue.qsize()
+
+            # If it's time for a scheduled flush or queue is above the trigger threshold
+            if current_time >= self._next_export_time or queue_size >= self._export_trigger_size:
+                self._export_batches(force=False)
+                # Reset the next scheduled flush time
+                self._next_export_time = time.time() + self._schedule_delay
+            else:
+                # Sleep a short interval so we don't busy-wait.
+                time.sleep(0.2)
+
+        # Final drain after shutdown
+        self._export_batches(force=True)
+
+    def _export_batches(self, force: bool = False):
+        """Drains the queue and exports in batches. If force=True, export everything.
+        Otherwise, export up to `max_batch_size` repeatedly until the queue is empty or below a
+        certain threshold.
+        """
+        while True:
+            items_to_export: list[Span[Any] | Trace] = []
+
+            # Gather a batch of spans up to max_batch_size
+            while not self._queue.empty() and (
+                force or len(items_to_export) < self._max_batch_size
+            ):
+                try:
+                    items_to_export.append(self._queue.get_nowait())
+                except queue.Empty:
+                    # Another thread might have emptied the queue between checks
+                    break
+
+            # If we collected nothing, we're done
+            if not items_to_export:
+                break
+
+            # Export the batch
+            self._exporter.export(items_to_export)
+
+
+# Create a shared global instance:
+_global_exporter = BackendSpanExporter()
+_global_processor = BatchTraceProcessor(_global_exporter)
+
+
+def default_exporter() -> BackendSpanExporter:
+    """The default exporter, which exports traces and spans to the backend in batches."""
+    return _global_exporter
+
+
+def default_processor() -> BatchTraceProcessor:
+    """The default processor, which exports traces and spans to the backend in batches."""
+    return _global_processor

agents/tracing/scope.py

@@ -0,0 +1,45 @@
+# Holds the current active span
+import contextvars
+from typing import TYPE_CHECKING, Any
+
+from .logger import logger
+
+if TYPE_CHECKING:
+    from .spans import Span
+    from .traces import Trace
+
+_current_span: contextvars.ContextVar["Span[Any] | None"] = contextvars.ContextVar(
+    "current_span", default=None
+)
+
+_current_trace: contextvars.ContextVar["Trace | None"] = contextvars.ContextVar(
+    "current_trace", default=None
+)
+
+
+class Scope:
+    @classmethod
+    def get_current_span(cls) -> "Span[Any] | None":
+        return _current_span.get()
+
+    @classmethod
+    def set_current_span(cls, span: "Span[Any] | None") -> "contextvars.Token[Span[Any] | None]":
+        return _current_span.set(span)
+
+    @classmethod
+    def reset_current_span(cls, token: "contextvars.Token[Span[Any] | None]") -> None:
+        _current_span.reset(token)
+
+    @classmethod
+    def get_current_trace(cls) -> "Trace | None":
+        return _current_trace.get()
+
+    @classmethod
+    def set_current_trace(cls, trace: "Trace | None") -> "contextvars.Token[Trace | None]":
+        logger.debug(f"Setting current trace: {trace.trace_id if trace else None}")
+        return _current_trace.set(trace)
+
+    @classmethod
+    def reset_current_trace(cls, token: "contextvars.Token[Trace | None]") -> None:
+        logger.debug("Resetting current trace")
+        _current_trace.reset(token)