8000 Allow replacing AgentRunner and TraceProvider (#720) · openai/openai-agents-python@0cf503e · GitHub
[go: up one dir, main page]
Skip to content

Commit 0cf503e

Browse files
pakrym-oaipakrym-oai
authored
Allow replacing AgentRunner and TraceProvider (#720)
1 parent 901d2ac commit 0cf503e

18 files changed

+2309
-2055
lines changed

‎src/agents/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@
104104
handoff_span,
105105
mcp_tools_span,
106106
set_trace_processors,
107+
set_trace_provider,
107108
set_tracing_disabled,
108109
set_tracing_export_api_key,
109110
speech_group_span,
@@ -246,6 +247,7 @@ def enable_verbose_stdout_logging():
246247
"guardrail_span",
247248
"handoff_span",
248249
"set_trace_processors",
250+
"set_trace_provider",
249251
"set_tracing_disabled",
250252
"speech_group_span",
251253
"transcription_span",

‎src/agents/run.py

Lines changed: 206 additions & 95 deletions
Large diffs are not rendered by default.

‎src/agents/tracing/__init__.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import atexit
22

3+
from agents.tracing.provider import DefaultTraceProvider, TraceProvider
4+
35
from .create import (
46
agent_span,
57
custom_span,
@@ -18,7 +20,7 @@
1820
)
1921
from .processor_interface import TracingProcessor
2022
from .processors import default_exporter, default_processor
21-
from .setup import GLOBAL_TRACE_PROVIDER
23+
from .setup import get_trace_provider, set_trace_provider
2224
from .span_data import (
2325
AgentSpanData,
2426
CustomSpanData,
@@ -45,10 +47,12 @@
4547
"generation_span",
4648
"get_current_span",
4749
"get_current_trace",
50+
"get_trace_provider",
4851
"guardrail_span",
4952
"handoff_span",
5053
"response_span",
5154
"set_trace_processors",
55+
"set_trace_provider",
5256
"set_tracing_disabled",
5357
"trace",
5458
"Trace",
@@ -67,6 +71,7 @@
6771
"SpeechSpanData",
6872
"TranscriptionSpanData",
6973
"TracingProcessor",
74+
"TraceProvider",
7075
"gen_trace_id",
7176
"gen_span_id",
7277
"speech_group_span",
@@ -80,21 +85,21 @@ def add_trace_processor(span_processor: TracingProcessor) -> None:
8085
"""
8186
Adds a new trace processor. This processor will receive all traces/spans.
8287
"""
83-
GLOBAL_TRACE_PROVIDER.register_processor(span_processor)
88+
get_trace_provider().register_processor(span_processor)
8489

8590

8691
def set_trace_processors(processors: list[TracingProcessor]) -> None:
8792
"""
8893
Set the list of trace processors. This will replace the current list of processors.
8994
"""
90-
GLOBAL_TRACE_PROVIDER.set_processors(processors)
95+
get_trace_provider().set_processors(processors)
9196

9297

9398
def set_tracing_disabled(disabled: bool) -> None:
9499
"""
95100
Set whether tracing is globally disabled.
96101
"""
97-
GLOBAL_TRACE_PROVIDER.set_disabled(disabled)
102+
get_trace_provider().set_disabled(disabled)
98103

99104

100105
def set_tracing_export_api_key(api_key: str) -> None:
@@ -104,10 +109,11 @@ def set_tracing_export_api_key(api_key: str) -> None:
104109
default_exporter().set_api_key(api_key)
105110

106111

112+
set_trace_provider(DefaultTraceProvider())
107113
# Add the default processor, which exports traces and spans to the backend in batches. You can
108114
# change the default behavior by either:
109115
# 1. calling add_trace_processor(), which adds additional processors, or
110116
# 2. calling set_trace_processors(), which replaces the default processor.
111117
add_trace_processor(default_processor())
112118

113-
atexit.register(GLOBAL_TRACE_PROVIDER.shutdown)
119+
atexit.register(get_trace_provider().shutdown)

‎src/agents/tracing/create.py

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from typing import TYPE_CHECKING, Any
55

66
from ..logger import logger
7-
from .setup import GLOBAL_TRACE_PROVIDER
7+
from .setup import get_trace_provider
88
from .span_data import (
99
AgentSpanData,
1010
CustomSpanData,
@@ -56,13 +56,13 @@ def trace(
5656
Returns:
5757
The newly created trace object.
5858
"""
59-
current_trace = GLOBAL_TRACE_PROVIDER.get_current_trace()
59+
current_trace = get_trace_provider().get_current_trace()
6060
if current_trace:
6161
logger.warning(
6262
"Trace already exists. Creating a new trace, but this is probably a mistake."
6363
)
6464

65-
return GLOBAL_TRACE_PROVIDER.create_trace(
65+
return get_trace_provider().create_trace(
6666
name=workflow_name,
6767
trace_id=trace_id,
6868
group_id=group_id,
@@ -73,12 +73,12 @@ def trace(
7373

7474
def get_current_trace() -> Trace | None:
7575
"""Returns the currently active trace, if present."""
76-
return GLOBAL_TRACE_PROVIDER.get_current_trace()
76+
return get_trace_provider().get_current_trace()
7777

7878

7979
def get_current_span() -> Span[Any] | None:
8080
"""Returns the currently active span, if present."""
81-
return GLOBAL_TRACE_PROVIDER.get_current_span()
81+
return get_trace_provider().get_current_span()
8282

8383

8484
def agent_span(
@@ -108,7 +108,7 @@ def agent_span(
108108
Returns:
109109
The newly created agent span.
110110
"""
111-
return GLOBAL_TRACE_PROVIDER.create_span(
111+
return get_trace_provider().create_span(
112112
span_data=AgentSpanData(name=name, handoffs=handoffs, tools=tools, output_type=output_type),
113113
span_id=span_id,
114114
parent=parent,
@@ -141,7 +141,7 @@ def function_span(
141141
Returns:
142142
The newly created function span.
143143
"""
144-
return GLOBAL_TRACE_PROVIDER.create_span(
144+
return get_trace_provider().create_span(
145145
span_data=FunctionSpanData(name=name, input=input, output=output),
146146
span_id=span_id,
147147
parent=parent,
@@ -183,7 +183,7 @@ def generation_span(
183183
Returns:
184184
The newly created generation span.
185185
"""
186-
return GLOBAL_TRACE_PROVIDER.create_span(
186+
return get_trace_provider().create_span(
187187
span_data=GenerationSpanData(
188188
input=input,
189189
output=output,
@@ -215,7 +215,7 @@ def response_span(
215215
trace/span as the parent.
216216
disabled: If True, we will return a Span but the Span will not be recorded.
217217
"""
218-
return GLOBAL_TRACE_PROVIDER.create_span(
218+
return get_trace_provider().create_span(
219219
span_data=ResponseSpanData(response=response),
220220
span_id=span_id,
221221
parent=parent,
@@ -246,7 +246,7 @@ def handoff_span(
246246
Returns:
247247
The newly created handoff span.
248248
"""
249-
return GLOBAL_TRACE_PROVIDER.create_span(
249+
return get_trace_provider().create_span(
250250
span_data=HandoffSpanData(from_agent=from_agent, to_agent=to_agent),
251251
span_id=span_id,
252252
parent=parent,
@@ -278,7 +278,7 @@ def custom_span(
278278
Returns:
279279
The newly created custom span.
280280
"""
281-
return GLOBAL_TRACE_PROVIDER.create_span(
281+
return get_trace_provider().create_span(
282282
span_data=CustomSpanData(name=name, data=data or {}),
283283
span_id=span_id,
284284
parent=parent,
@@ -306,7 +306,7 @@ def guardrail_span(
306306
trace/span as the parent.
307307
disabled: If True, we will return a Span but the Span will not be recorded.
308308
"""
309-
return GLOBAL_TRACE_PROVIDER.create_span(
309+
return get_trace_provider().create_span(
310310
span_data=GuardrailSpanData(name=name, triggered=triggered),
311311
span_id=span_id,
312312
parent=parent,
@@ -344,7 +344,7 @@ def transcription_span(
344344
Returns:
345345
The newly created speech-to-text span.
346346
"""
347-
return GLOBAL_TRACE_PROVIDER.create_span(
347+
return get_trace_provider().create_span(
348348
span_data=TranscriptionSpanData(
349349
input=input,
350350
input_format=input_format,
@@ -386,7 +386,7 @@ def speech_span(
386386
trace/span as the parent.
387387
disabled: If True, we will return a Span but the Span will not be recorded.
388388
"""
389-
return GLOBAL_TRACE_PROVIDER.create_span(
389+
return get_trace_provider().create_span(
390390
span_data=SpeechSpanData(
391391
model=model,
392392
input=input,
@@ -419,7 +419,7 @@ def speech_group_span(
419419
trace/span as the parent.
420420
disabled: If True, we will return a Span but the Span will not be recorded.
421421
"""
422-
return GLOBAL_TRACE_PROVIDER.create_span(
422+
return get_trace_provider().create_span(
423423
span_data=SpeechGroupSpanData(input=input),
424424
span_id=span_id,
425425
parent=parent,
@@ -447,7 +447,7 @@ def mcp_tools_span(
447447
trace/span as the parent.
448448
disabled: If True, we will return a Span but the Span will not be recorded.
449449
"""
450-
return GLOBAL_TRACE_PROVIDER.create_span(
450+
return get_trace_provider().create_span(
451451
span_data=MCPListToolsSpanData(server=server, result=result),
452452
span_id=span_id,
453453
parent=parent,

0 commit comments

Comments
 (0)
0