Skip to content

Commit 97b4f9a

Browse files
authored
Update run.py
1 parent 7bf5202 commit 97b4f9a

File tree

1 file changed

+24
-5
lines changed

1 file changed

+24
-5
lines changed

src/agents/run.py

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
import asyncio
4+
import contextvars
45
import copy
56
import inspect
67
from dataclasses import dataclass, field
@@ -51,11 +52,14 @@
5152
from .stream_events import AgentUpdatedStreamEvent, RawResponsesStreamEvent
5253
from .tool import Tool
5354
from .tracing import Span, SpanError, agent_span, get_current_trace, trace
54-
from .tracing.scope import Scope
5555
from .tracing.span_data import AgentSpanData
5656
from .usage import Usage
5757
from .util import _coro, _error_tracing
5858

59+
_current_run_config: contextvars.ContextVar[RunConfig | None] = contextvars.ContextVar(
60+
"current_run_config", default=None
61+
)
62+
5963
DEFAULT_MAX_TURNS = 10
6064

6165
DEFAULT_AGENT_RUNNER: AgentRunner = None # type: ignore
@@ -80,6 +84,21 @@ def get_default_agent_runner() -> AgentRunner:
8084
return DEFAULT_AGENT_RUNNER
8185

8286

87+
def get_current_run_config() -> RunConfig | None:
88+
"""Get the current run config from context."""
89+
return _current_run_config.get()
90+
91+
92+
def set_current_run_config(run_config: RunConfig | None) -> contextvars.Token[RunConfig | None]:
93+
"""Set the current run config in context."""
94+
return _current_run_config.set(run_config)
95+
96+
97+
def reset_current_run_config(token: contextvars.Token[RunConfig | None]) -> None:
98+
"""Reset the current run config in context."""
99+
_current_run_config.reset(token)
100+
101+
83102
@dataclass
84103
class RunConfig:
85104
"""Configures settings for the entire agent run."""
@@ -341,7 +360,7 @@ async def run(
341360
# Set the run_config context variable if enabled
342361
run_config_token = None
343362
if run_config.pass_run_config_to_sub_agents:
344-
run_config_token = Scope.set_current_run_config(run_config)
363+
run_config_token = set_current_run_config(run_config)
345364

346365
try:
347366
with TraceCtxManager(
@@ -495,7 +514,7 @@ async def run(
495514
finally:
496515
# Always clean up the context variable
497516
if run_config_token is not None:
498-
Scope.reset_current_run_config(run_config_token)
517+
reset_current_run_config(run_config_token)
499518

500519
def run_sync(
501520
self,
@@ -539,7 +558,7 @@ def run_streamed(
539558
# Set the run_config context variable if enabled
540559
run_config_token = None
541560
if run_config.pass_run_config_to_sub_agents:
542-
run_config_token = Scope.set_current_run_config(run_config)
561+
run_config_token = set_current_run_config(run_config)
543562

544563
try:
545564
# If there's already a trace, we don't create a new one. In addition, we can't end the
@@ -596,7 +615,7 @@ def run_streamed(
596615
finally:
597616
# Always reset the context variable
598617
if run_config_token is not None:
599-
Scope.reset_current_run_config(run_config_token)
618+
reset_current_run_config(run_config_token)
600619

601620
@classmethod
602621
async def _run_input_guardrails_with_queue(

0 commit comments

Comments
 (0)