Allow replacing AgentRunner and TraceProvider (#720)

Author: pakrym-oai

src/agents/__init__.py

@@ -104,6 +104,7 @@
     handoff_span,
     mcp_tools_span,
     set_trace_processors,
+    set_trace_provider,
     set_tracing_disabled,
     set_tracing_export_api_key,
     speech_group_span,
@@ -246,6 +247,7 @@ def enable_verbose_stdout_logging():
     "guardrail_span",
     "handoff_span",
     "set_trace_processors",
+    "set_trace_provider",
     "set_tracing_disabled",
     "speech_group_span",
     "transcription_span",

src/agents/run.py

@@ -1,5 +1,7 @@
 import atexit
 
+from agents.tracing.provider import DefaultTraceProvider, TraceProvider
+
 from .create import (
     agent_span,
     custom_span,
@@ -18,7 +20,7 @@
 )
 from .processor_interface import TracingProcessor
 from .processors import default_exporter, default_processor
-from .setup import GLOBAL_TRACE_PROVIDER
+from .setup import get_trace_provider, set_trace_provider
 from .span_data import (
     AgentSpanData,
     CustomSpanData,
@@ -45,10 +47,12 @@
     "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",
@@ -67,6 +71,7 @@
     "SpeechSpanData",
     "TranscriptionSpanData",
     "TracingProcessor",
+    "TraceProvider",
     "gen_trace_id",
     "gen_span_id",
     "speech_group_span",
@@ -80,21 +85,21 @@ def add_trace_processor(span_processor: TracingProcessor) -> None:
     """
     Adds a new trace processor. This processor will receive all traces/spans.
     """
-    GLOBAL_TRACE_PROVIDER.register_processor(span_processor)
+    get_trace_provider().register_processor(span_processor)
 
 
 def set_trace_processors(processors: list[TracingProcessor]) -> None:
     """
     Set the list of trace processors. This will replace the current list of processors.
     """
-    GLOBAL_TRACE_PROVIDER.set_processors(processors)
+    get_trace_provider().set_processors(processors)
 
 
 def set_tracing_disabled(disabled: bool) -> None:
     """
     Set whether tracing is globally disabled.
     """
-    GLOBAL_TRACE_PROVIDER.set_disabled(disabled)
+    get_trace_provider().set_disabled(disabled)
 
 
 def set_tracing_export_api_key(api_key: str) -> None:
@@ -104,10 +109,11 @@ def set_tracing_export_api_key(api_key: str) -> None:
     default_exporter().set_api_key(api_key)
 
 
+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())
 
-atexit.register(GLOBAL_TRACE_PROVIDER.shutdown)
+atexit.register(get_trace_provider().shutdown)

src/agents/tracing/__init__.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 from ..logger import logger
-from .setup import GLOBAL_TRACE_PROVIDER
+from .setup import get_trace_provider
 from .span_data import (
     AgentSpanData,
     CustomSpanData,
@@ -56,13 +56,13 @@ def trace(
     Returns:
         The newly created trace object.
     """
-    current_trace = GLOBAL_TRACE_PROVIDER.get_current_trace()
+    current_trace = get_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(
+    return get_trace_provider().create_trace(
         name=workflow_name,
         trace_id=trace_id,
         group_id=group_id,
@@ -73,12 +73,12 @@ def trace(
 
 def get_current_trace() -> Trace | None:
     """Returns the currently active trace, if present."""
-    return GLOBAL_TRACE_PROVIDER.get_current_trace()
+    return get_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()
+    return get_trace_provider().get_current_span()
 
 
 def agent_span(
@@ -108,7 +108,7 @@ def agent_span(
     Returns:
         The newly created agent span.
     """
-    return GLOBAL_TRACE_PROVIDER.create_span(
+    return get_trace_provider().create_span(
         span_data=AgentSpanData(name=name, handoffs=handoffs, tools=tools, output_type=output_type),
         span_id=span_id,
         parent=parent,
@@ -141,7 +141,7 @@ def function_span(
     Returns:
         The newly created function span.
     """
-    return GLOBAL_TRACE_PROVIDER.create_span(
+    return get_trace_provider().create_span(
         span_data=FunctionSpanData(name=name, input=input, output=output),
         span_id=span_id,
         parent=parent,
@@ -183,7 +183,7 @@ def generation_span(
     Returns:
         The newly created generation span.
     """
-    return GLOBAL_TRACE_PROVIDER.create_span(
+    return get_trace_provider().create_span(
         span_data=GenerationSpanData(
             input=input,
             output=output,
@@ -215,7 +215,7 @@ def response_span(
             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(
+    return get_trace_provider().create_span(
         span_data=ResponseSpanData(response=response),
         span_id=span_id,
         parent=parent,
@@ -246,7 +246,7 @@ def handoff_span(
     Returns:
         The newly created handoff span.
     """
-    return GLOBAL_TRACE_PROVIDER.create_span(
+    return get_trace_provider().create_span(
         span_data=HandoffSpanData(from_agent=from_agent, to_agent=to_agent),
         span_id=span_id,
         parent=parent,
@@ -278,7 +278,7 @@ def custom_span(
     Returns:
         The newly created custom span.
     """
-    return GLOBAL_TRACE_PROVIDER.create_span(
+    return get_trace_provider().create_span(
         span_data=CustomSpanData(name=name, data=data or {}),
         span_id=span_id,
         parent=parent,
@@ -306,7 +306,7 @@ def guardrail_span(
             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(
+    return get_trace_provider().create_span(
         span_data=GuardrailSpanData(name=name, triggered=triggered),
         span_id=span_id,
         parent=parent,
@@ -344,7 +344,7 @@ def transcription_span(
     Returns:
         The newly created speech-to-text span.
     """
-    return GLOBAL_TRACE_PROVIDER.create_span(
+    return get_trace_provider().create_span(
         span_data=TranscriptionSpanData(
             input=input,
             input_format=input_format,
@@ -386,7 +386,7 @@ def speech_span(
             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(
+    return get_trace_provider().create_span(
         span_data=SpeechSpanData(
             model=model,
             input=input,
@@ -419,7 +419,7 @@ def speech_group_span(
             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(
+    return get_trace_provider().create_span(
         span_data=SpeechGroupSpanData(input=input),
         span_id=span_id,
         parent=parent,
@@ -447,7 +447,7 @@ def mcp_tools_span(
             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(
+    return get_trace_provider().create_span(
         span_data=MCPListToolsSpanData(server=server, result=result),
         span_id=span_id,
         parent=parent,