mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-23 08:52:56 +00:00
24 lines
680 B
Python
24 lines
680 B
Python
|
from contextlib import contextmanager
|
||
|
|
from contextvars import ContextVar
|
||
|
|
from typing import Any, ClassVar, Generator
|
||
|
|
|
||
|
|
|
||
|
|
class RunContext:
|
||
|
|
RUN_CONTEXT_VAR: ClassVar[ContextVar] = ContextVar("RUN_CONTEXT_VAR")
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
@contextmanager
|
||
|
|
def populate_context(cls, run_id) -> Generator[None, Any, None]:
|
||
|
|
token = RunContext.RUN_CONTEXT_VAR.set(run_id)
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
RunContext.RUN_CONTEXT_VAR.reset(token)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def current_run_id(cls) -> str:
|
||
|
|
try:
|
||
|
|
return cls.RUN_CONTEXT_VAR.get()
|
||
|
|
except LookupError as e:
|
||
|
|
raise RuntimeError("Error getting run id") from e
|