diff --git a/src/agents/tracing/__init__.py b/src/agents/tracing/__init__.py index 64b0bd71f..ddfdae3e1 100644 --- a/src/agents/tracing/__init__.py +++ b/src/agents/tracing/__init__.py @@ -1,119 +1,161 @@ -import atexit +""" +OpenAI Agents Tracing -from agents.tracing.provider import DefaultTraceProvider, TraceProvider +Configuration: + - AGENT_TRACE_PROVIDER: Optional "module.ClassName" specifying a custom TraceProvider. + If unset, DefaultTraceProvider is used. + - Use `set_tracing_export_api_key()` to configure the API key for the built-in exporter. -from .create import ( - agent_span, - custom_span, - function_span, - generation_span, - get_current_span, - get_current_trace, - guardrail_span, - handoff_span, - mcp_tools_span, - response_span, - speech_group_span, - speech_span, - trace, - transcription_span, -) +Runtime APIs: + - add_trace_processor(span_processor: TracingProcessor) + - set_trace_processors(processors: List[TracingProcessor]) + - set_tracing_disabled(disabled: bool) + - set_tracing_export_api_key(api_key: str) +""" + +# Enable postponed evaluation of annotations for Python <3.10 compatibility +from __future__ import annotations + +import atexit +import os +import logging +from typing import List # Removed unused Optional + +# Critical imports explicitly exposed in __all__ +from .provider import TraceProvider # Now directly available for importers +from .provider import DefaultTraceProvider from .processor_interface import TracingProcessor -from .processors import default_exporter, default_processor -from .setup import get_trace_provider, set_trace_provider from .span_data import ( - AgentSpanData, - CustomSpanData, - FunctionSpanData, - GenerationSpanData, - GuardrailSpanData, - HandoffSpanData, - MCPListToolsSpanData, - ResponseSpanData, - SpanData, - SpeechGroupSpanData, - SpeechSpanData, - TranscriptionSpanData, + AgentSpanData, CustomSpanData, FunctionSpanData, GenerationSpanData, + GuardrailSpanData, HandoffSpanData, MCPListToolsSpanData, + ResponseSpanData, SpanData, SpeechGroupSpanData, SpeechSpanData, + TranscriptionSpanData ) from .spans import Span, SpanError from .traces import Trace from .util import gen_span_id, gen_trace_id +from .create import ( + agent_span, custom_span, function_span, generation_span, get_current_span, + get_current_trace, guardrail_span, handoff_span, mcp_tools_span, + response_span, speech_group_span, speech_span, trace, transcription_span +) __all__ = [ "add_trace_processor", - "agent_span", - "custom_span", - "function_span", - "generation_span", - "get_current_span", - "get_current_trace", - "get_trace_provider", - "guardrail_span", - "handoff_span", - "response_span", "set_trace_processors", - "set_trace_provider", "set_tracing_disabled", - "trace", - "Trace", - "SpanError", - "Span", - "SpanData", - "AgentSpanData", - "CustomSpanData", - "FunctionSpanData", - "GenerationSpanData", - "GuardrailSpanData", - "HandoffSpanData", - "MCPListToolsSpanData", - "ResponseSpanData", - "SpeechGroupSpanData", - "SpeechSpanData", - "TranscriptionSpanData", - "TracingProcessor", - "TraceProvider", - "gen_trace_id", - "gen_span_id", - "speech_group_span", - "speech_span", - "transcription_span", - "mcp_tools_span", + "set_tracing_export_api_key", + "agent_span", "custom_span", "function_span", "generation_span", + "get_current_span", "get_current_trace", + "guardrail_span", "handoff_span", "response_span", "speech_group_span", + "speech_span", "trace", "transcription_span", "mcp_tools_span", + "Trace", "Span", "SpanError", + "SpanData", "AgentSpanData", "CustomSpanData", "FunctionSpanData", + "GenerationSpanData", "GuardrailSpanData", "HandoffSpanData", + "MCPListToolsSpanData", "ResponseSpanData", "SpeechGroupSpanData", + "SpeechSpanData", "TranscriptionSpanData", + "TracingProcessor", "TraceProvider", "gen_trace_id", "gen_span_id", ] - -def add_trace_processor(span_processor: TracingProcessor) -> None: +def _load_provider() -> TraceProvider: """ - Adds a new trace processor. This processor will receive all traces/spans. + Load a TraceProvider implementation based on AGENT_TRACE_PROVIDER env var, + or fall back to DefaultTraceProvider. """ - get_trace_provider().register_processor(span_processor) + from .provider import DefaultTraceProvider, TraceProvider + spec = os.getenv("AGENT_TRACE_PROVIDER") + if not spec: + provider = DefaultTraceProvider() + logging.debug("Using default trace provider: %s", type(provider).__name__) + return provider -def set_trace_processors(processors: list[TracingProcessor]) -> None: - """ - Set the list of trace processors. This will replace the current list of processors. - """ - get_trace_provider().set_processors(processors) + module_name, _, class_name = spec.rpartition(".") + try: + mod = __import__(module_name, fromlist=[class_name]) + provider_cls = getattr(mod, class_name) + except ImportError as exc: + raise RuntimeError( + f"Custom provider '{spec}' not found. Check your AGENT_TRACE_PROVIDER value." + ) from exc + except Exception as exc: + logging.warning( + "Error loading AGENT_TRACE_PROVIDER=%r: %s", spec, exc + ) + return DefaultTraceProvider() + if not issubclass(provider_cls, TraceProvider): + raise TypeError(f"{spec!r} must inherit from TraceProvider") -def set_tracing_disabled(disabled: bool) -> None: - """ - Set whether tracing is globally disabled. - """ - get_trace_provider().set_disabled(disabled) + provider = provider_cls() + logging.debug("Using custom trace provider: %s", spec) + return provider +# Bootstrap the global provider +from .setup import set_trace_provider, get_trace_provider +try: + provider = _load_provider() + set_trace_provider(provider) +except Exception as exc: + logging.warning("Failed to set trace provider: %s", exc) -def set_tracing_export_api_key(api_key: str) -> None: - """ - Set the OpenAI API key for the backend exporter. - """ - default_exporter().set_api_key(api_key) +def add_trace_processor(span_processor: TracingProcessor) -> None: + """Register a new trace processor to receive all spans.""" + try: + from .setup import get_trace_provider as _get + _get().register_processor(span_processor) + except Exception as exc: + logging.warning("Failed to register trace processor: %s", exc) + +def set_trace_processors(processors: List[TracingProcessor]) -> None: + """Replace all trace processors with the given list.""" + if not processors: + logging.warning("Empty processor list provided to set_trace_processors()") + try: + from .setup import get_trace_provider as _get + _get().set_processors(processors) + except Exception as exc: + logging.warning("Failed to set trace processors: %s", exc) +def set_tracing_disabled(disabled: bool) -> None: + """Globally enable or disable tracing at runtime.""" + try: + from .setup import get_trace_provider as _get + _get().set_disabled(disabled) + except Exception as exc: + logging.warning("Failed to toggle tracing: %s", exc) + +def set_tracing_export_api_key(api_key: str) -> None: + """Configure the OpenAI API key for the built-in tracing exporter.""" + try: + from .processors import default_exporter + exporter = default_exporter() + if exporter is None: + raise RuntimeError( + "No tracing exporter available; did you install the optional extras?" + ) + exporter.set_api_key(api_key) + except Exception as exc: + logging.warning("Failed to set tracing export API key: %s", exc) + if isinstance(exc, RuntimeError): + raise -set_trace_provider(DefaultTraceProvider()) -# Add the default processor, which exports traces and spans to the backend in batches. You can -# change the default behavior by either: -# 1. calling add_trace_processor(), which adds additional processors, or -# 2. calling set_trace_processors(), which replaces the default processor. -add_trace_processor(default_processor()) +# Install default processor (only if it succeeds) +try: + from .processors import default_processor + proc = default_processor() + if proc is not None: + add_trace_processor(proc) +except Exception as exc: + logging.warning("Failed to install default trace processor: %s", exc) -atexit.register(get_trace_provider().shutdown) +# Ensure provider shutdown at exit (with safety checks) +try: + provider = get_trace_provider() + shutdown_fn = getattr(provider, "shutdown", None) + if callable(shutdown_fn): + atexit.register(shutdown_fn) + else: + logging.debug("Trace provider has no shutdown method - skipping cleanup registration") +except Exception as exc: + logging.warning("Failed to register shutdown hook: %s", exc)