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

agents/tracing/setup.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import os
+import threading
+from typing import Any
+
+from . import util
+from .logger import logger
+from .processor_interface import TracingProcessor
+from .scope import Scope
+from .spans import NoOpSpan, Span, SpanImpl, TSpanData
+from .traces import NoOpTrace, Trace, TraceImpl
+
+
+class SynchronousMultiTracingProcessor(TracingProcessor):
+    """
+    Forwards all calls to a list of TracingProcessors, in order of registration.
+    """
+
+    def __init__(self):
+        # Using a tuple to avoid race conditions when iterating over processors
+        self._processors: tuple[TracingProcessor, ...] = ()
+        self._lock = threading.Lock()
+
+    def add_tracing_processor(self, tracing_processor: TracingProcessor):
+        """
+        Add a processor to the list of processors. Each processor will receive all traces/spans.
+        """
+        with self._lock:
+            self._processors += (tracing_processor,)
+
+    def set_processors(self, processors: list[TracingProcessor]):
+        """
+        Set the list of processors. This will replace the current list of processors.
+        """
+        with self._lock:
+            self._processors = tuple(processors)
+
+    def on_trace_start(self, trace: Trace) -> None:
+        """
+        Called when a trace is started.
+        """
+        for processor in self._processors:
+            processor.on_trace_start(trace)
+
+    def on_trace_end(self, trace: Trace) -> None:
+        """
+        Called when a trace is finished.
+        """
+        for processor in self._processors:
+            processor.on_trace_end(trace)
+
+    def on_span_start(self, span: Span[Any]) -> None:
+        """
+        Called when a span is started.
+        """
+        for processor in self._processors:
+            processor.on_span_start(span)
+
+    def on_span_end(self, span: Span[Any]) -> None:
+        """
+        Called when a span is finished.
+        """
+        for processor in self._processors:
+            processor.on_span_end(span)
+
+    def shutdown(self) -> None:
+        """
+        Called when the application stops.
+        """
+        for processor in self._processors:
+            logger.debug(f"Shutting down trace processor {processor}")
+            processor.shutdown()
+
+    def force_flush(self):
+        """
+        Force the processors to flush their buffers.
+        """
+        for processor in self._processors:
+            processor.force_flush()
+
+
+class TraceProvider:
+    def __init__(self):
+        self._multi_processor = SynchronousMultiTracingProcessor()
+        self._disabled = os.environ.get("OPENAI_AGENTS_DISABLE_TRACING", "false").lower() in (
+            "true",
+            "1",
+        )
+
+    def register_processor(self, processor: TracingProcessor):
+        """
+        Add a processor to the list of processors. Each processor will receive all traces/spans.
+        """
+        self._multi_processor.add_tracing_processor(processor)
+
+    def set_processors(self, processors: list[TracingProcessor]):
+        """
+        Set the list of processors. This will replace the current list of processors.
+        """
+        self._multi_processor.set_processors(processors)
+
+    def get_current_trace(self) -> Trace | None:
+        """
+        Returns the currently active trace, if any.
+        """
+        return Scope.get_current_trace()
+
+    def get_current_span(self) -> Span[Any] | None:
+        """
+        Returns the currently active span, if any.
+        """
+        return Scope.get_current_span()
+
+    def set_disabled(self, disabled: bool) -> None:
+        """
+        Set whether tracing is disabled.
+        """
+        self._disabled = disabled
+
+    def create_trace(
+        self,
+        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.
+        """
+        if self._disabled or disabled:
+            logger.debug(f"Tracing is disabled. Not creating trace {name}")
+            return NoOpTrace()
+
+        trace_id = trace_id or util.gen_trace_id()
+
+        logger.debug(f"Creating trace {name} with id {trace_id}")
+
+        return TraceImpl(
+            name=name,
+            trace_id=trace_id,
+            group_id=group_id,
+            metadata=metadata,
+            processor=self._multi_processor,
+        )
+
+    def create_span(
+        self,
+        span_data: TSpanData,
+        span_id: str | None = None,
+        parent: Trace | Span[Any] | None = None,
+        disabled: bool = False,
+    ) -> Span[TSpanData]:
+        """
+        Create a new span.
+        """
+        if self._disabled or disabled:
+            logger.debug(f"Tracing is disabled. Not creating span {span_data}")
+            return NoOpSpan(span_data)
+
+        if not parent:
+            current_span = Scope.get_current_span()
+            current_trace = Scope.get_current_trace()
+            if current_trace is None:
+                logger.error(
+                    "No active trace. Make sure to start a trace with `trace()` first"
+                    "Returning NoOpSpan."
+                )
+                return NoOpSpan(span_data)
+            elif isinstance(current_trace, NoOpTrace) or isinstance(current_span, NoOpSpan):
+                logger.debug(
+                    f"Parent {current_span} or {current_trace} is no-op, returning NoOpSpan"
+                )
+                return NoOpSpan(span_data)
+
+            parent_id = current_span.span_id if current_span else None
+            trace_id = current_trace.trace_id
+
+        elif isinstance(parent, Trace):
+            if isinstance(parent, NoOpTrace):
+                logger.debug(f"Parent {parent} is no-op, returning NoOpSpan")
+                return NoOpSpan(span_data)
+            trace_id = parent.trace_id
+            parent_id = None
+        elif isinstance(parent, Span):
+            if isinstance(parent, NoOpSpan):
+                logger.debug(f"Parent {parent} is no-op, returning NoOpSpan")
+                return NoOpSpan(span_data)
+            parent_id = parent.span_id
+            trace_id = parent.trace_id
+
+        logger.debug(f"Creating span {span_data} with id {span_id}")
+
+        return SpanImpl(
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_id=parent_id,
+            processor=self._multi_processor,
+            span_data=span_data,
+        )
+
+    def shutdown(self) -> None:
+        try:
+            logger.debug("Shutting down trace provider")
+            self._multi_processor.shutdown()
+        except Exception as e:
+            logger.error(f"Error shutting down trace provider: {e}")
+
+
+GLOBAL_TRACE_PROVIDER = TraceProvider()

agents/tracing/span_data.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import abc
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from openai.types.responses import Response, ResponseInputItemParam
+
+
+class SpanData(abc.ABC):
+    @abc.abstractmethod
+    def export(self) -> dict[str, Any]:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def type(self) -> str:
+        pass
+
+
+class AgentSpanData(SpanData):
+    __slots__ = ("name", "handoffs", "tools", "output_type")
+
+    def __init__(
+        self,
+        name: str,
+        handoffs: list[str] | None = None,
+        tools: list[str] | None = None,
+        output_type: str | None = None,
+    ):
+        self.name = name
+        self.handoffs: list[str] | None = handoffs
+        self.tools: list[str] | None = tools
+        self.output_type: str | None = output_type
+
+    @property
+    def type(self) -> str:
+        return "agent"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "name": self.name,
+            "handoffs": self.handoffs,
+            "tools": self.tools,
+            "output_type": self.output_type,
+        }
+
+
+class FunctionSpanData(SpanData):
+    __slots__ = ("name", "input", "output")
+
+    def __init__(self, name: str, input: str | None, output: str | None):
+        self.name = name
+        self.input = input
+        self.output = output
+
+    @property
+    def type(self) -> str:
+        return "function"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "name": self.name,
+            "input": self.input,
+            "output": self.output,
+        }
+
+
+class GenerationSpanData(SpanData):
+    __slots__ = (
+        "input",
+        "output",
+        "model",
+        "model_config",
+        "usage",
+    )
+
+    def __init__(
+        self,
+        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,
+    ):
+        self.input = input
+        self.output = output
+        self.model = model
+        self.model_config = model_config
+        self.usage = usage
+
+    @property
+    def type(self) -> str:
+        return "generation"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "input": self.input,
+            "output": self.output,
+            "model": self.model,
+            "model_config": self.model_config,
+            "usage": self.usage,
+        }
+
+
+class ResponseSpanData(SpanData):
+    __slots__ = ("response", "input")
+
+    def __init__(
+        self,
+        response: Response | None = None,
+        input: str | list[ResponseInputItemParam] | None = None,
+    ) -> None:
+        self.response = response
+        # This is not used by the OpenAI trace processors, but is useful for other tracing
+        # processor implementations
+        self.input = input
+
+    @property
+    def type(self) -> str:
+        return "response"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "response_id": self.response.id if self.response else None,
+        }
+
+
+class HandoffSpanData(SpanData):
+    __slots__ = ("from_agent", "to_agent")
+
+    def __init__(self, from_agent: str | None, to_agent: str | None):
+        self.from_agent = from_agent
+        self.to_agent = to_agent
+
+    @property
+    def type(self) -> str:
+        return "handoff"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "from_agent": self.from_agent,
+            "to_agent": self.to_agent,
+        }
+
+
+class CustomSpanData(SpanData):
+    __slots__ = ("name", "data")
+
+    def __init__(self, name: str, data: dict[str, Any]):
+        self.name = name
+        self.data = data
+
+    @property
+    def type(self) -> str:
+        return "custom"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "name": self.name,
+            "data": self.data,
+        }
+
+
+class GuardrailSpanData(SpanData):
+    __slots__ = ("name", "triggered")
+
+    def __init__(self, name: str, triggered: bool = False):
+        self.name = name
+        self.triggered = triggered
+
+    @property
+    def type(self) -> str:
+        return "guardrail"
+
+    def export(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "name": self.name,
+            "triggered": self.triggered,
+        }

agents/tracing/spans.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import abc
+import contextvars
+from typing import Any, Generic, TypeVar
+
+from typing_extensions import TypedDict
+
+from . import util
+from .logger import logger
+from .processor_interface import TracingProcessor
+from .scope import Scope
+from .span_data import SpanData
+
+TSpanData = TypeVar("TSpanData", bound=SpanData)
+
+
+class SpanError(TypedDict):
+    message: str
+    data: dict[str, Any] | None
+
+
+class Span(abc.ABC, Generic[TSpanData]):
+    @property
+    @abc.abstractmethod
+    def trace_id(self) -> str:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def span_id(self) -> str:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def span_data(self) -> TSpanData:
+        pass
+
+    @abc.abstractmethod
+    def start(self, mark_as_current: bool = False):
+        """
+        Start the span.
+
+        Args:
+            mark_as_current: If true, the span will be marked as the current span.
+        """
+        pass
+
+    @abc.abstractmethod
+    def finish(self, reset_current: bool = False) -> None:
+        """
+        Finish the span.
+
+        Args:
+            reset_current: If true, the span will be reset as the current span.
+        """
+        pass
+
+    @abc.abstractmethod
+    def __enter__(self) -> Span[TSpanData]:
+        pass
+
+    @abc.abstractmethod
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def parent_id(self) -> str | None:
+        pass
+
+    @abc.abstractmethod
+    def set_error(self, error: SpanError) -> None:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def error(self) -> SpanError | None:
+        pass
+
+    @abc.abstractmethod
+    def export(self) -> dict[str, Any] | None:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def started_at(self) -> str | None:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def ended_at(self) -> str | None:
+        pass
+
+
+class NoOpSpan(Span[TSpanData]):
+    __slots__ = ("_span_data", "_prev_span_token")
+
+    def __init__(self, span_data: TSpanData):
+        self._span_data = span_data
+        self._prev_span_token: contextvars.Token[Span[TSpanData] | None] | None = None
+
+    @property
+    def trace_id(self) -> str:
+        return "no-op"
+
+    @property
+    def span_id(self) -> str:
+        return "no-op"
+
+    @property
+    def span_data(self) -> TSpanData:
+        return self._span_data
+
+    @property
+    def parent_id(self) -> str | None:
+        return None
+
+    def start(self, mark_as_current: bool = False):
+        if mark_as_current:
+            self._prev_span_token = Scope.set_current_span(self)
+
+    def finish(self, reset_current: bool = False) -> None:
+        if reset_current and self._prev_span_token is not None:
+            Scope.reset_current_span(self._prev_span_token)
+            self._prev_span_token = None
+
+    def __enter__(self) -> Span[TSpanData]:
+        self.start(mark_as_current=True)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        reset_current = True
+        if exc_type is GeneratorExit:
+            logger.debug("GeneratorExit, skipping span reset")
+            reset_current = False
+
+        self.finish(reset_current=reset_current)
+
+    def set_error(self, error: SpanError) -> None:
+        pass
+
+    @property
+    def error(self) -> SpanError | None:
+        return None
+
+    def export(self) -> dict[str, Any] | None:
+        return None
+
+    @property
+    def started_at(self) -> str | None:
+        return None
+
+    @property
+    def ended_at(self) -> str | None:
+        return None
+
+
+class SpanImpl(Span[TSpanData]):
+    __slots__ = (
+        "_trace_id",
+        "_span_id",
+        "_parent_id",
+        "_started_at",
+        "_ended_at",
+        "_error",
+        "_prev_span_token",
+        "_processor",
+        "_span_data",
+    )
+
+    def __init__(
+        self,
+        trace_id: str,
+        span_id: str | None,
+        parent_id: str | None,
+        processor: TracingProcessor,
+        span_data: TSpanData,
+    ):
+        self._trace_id = trace_id
+        self._span_id = span_id or util.gen_span_id()
+        self._parent_id = parent_id
+        self._started_at: str | None = None
+        self._ended_at: str | None = None
+        self._processor = processor
+        self._error: SpanError | None = None
+        self._prev_span_token: contextvars.Token[Span[TSpanData] | None] | None = None
+        self._span_data = span_data
+
+    @property
+    def trace_id(self) -> str:
+        return self._trace_id
+
+    @property
+    def span_id(self) -> str:
+        return self._span_id
+
+    @property
+    def span_data(self) -> TSpanData:
+        return self._span_data
+
+    @property
+    def parent_id(self) -> str | None:
+        return self._parent_id
+
+    def start(self, mark_as_current: bool = False):
+        if self.started_at is not None:
+            logger.warning("Span already started")
+            return
+
+        self._started_at = util.time_iso()
+        self._processor.on_span_start(self)
+        if mark_as_current:
+            self._prev_span_token = Scope.set_current_span(self)
+
+    def finish(self, reset_current: bool = False) -> None:
+        if self.ended_at is not None:
+            logger.warning("Span already finished")
+            return
+
+        self._ended_at = util.time_iso()
+        self._processor.on_span_end(self)
+        if reset_current and self._prev_span_token is not None:
+            Scope.reset_current_span(self._prev_span_token)
+            self._prev_span_token = None
+
+    def __enter__(self) -> Span[TSpanData]:
+        self.start(mark_as_current=True)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        reset_current = True
+        if exc_type is GeneratorExit:
+            logger.debug("GeneratorExit, skipping span reset")
+            reset_current = False
+
+        self.finish(reset_current=reset_current)
+
+    def set_error(self, error: SpanError) -> None:
+        self._error = error
+
+    @property
+    def error(self) -> SpanError | None:
+        return self._error
+
+    @property
+    def started_at(self) -> str | None:
+        return self._started_at
+
+    @property
+    def ended_at(self) -> str | None:
+        return self._ended_at
+
+    def export(self) -> dict[str, Any] | None:
+        return {
+            "object": "trace.span",
+            "id": self.span_id,
+            "trace_id": self.trace_id,
+            "parent_id": self._parent_id,
+            "started_at": self._started_at,
+            "ended_at": self._ended_at,
+            "span_data": self.span_data.export(),
+            "error": self._error,
+        }