mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
"""Call-source context shared by product runtime and benchmark tracking.
|
|
|
|
This module only tags the current async context with the logical source of an
|
|
LLM call. It does not collect usage, emit telemetry events, or replace the
|
|
anonymous statistics/event telemetry under :mod:`openspace.utils.telemetry`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextvars
|
|
from contextlib import contextmanager
|
|
from typing import Iterator
|
|
|
|
|
|
AGENT_SOURCE = "agent"
|
|
|
|
# Default to the main grounding loop so product code does not need to annotate
|
|
# ordinary agent calls. Benchmark token accounting can then count this bucket
|
|
# separately from skill/analysis overhead.
|
|
CALL_SOURCE: contextvars.ContextVar[str] = contextvars.ContextVar(
|
|
"openspace_call_source",
|
|
default=AGENT_SOURCE,
|
|
)
|
|
|
|
|
|
def set_call_source(source: str) -> contextvars.Token:
|
|
"""Set the current call source and return a token for reset."""
|
|
return CALL_SOURCE.set(source)
|
|
|
|
|
|
def reset_call_source(token: contextvars.Token) -> None:
|
|
"""Reset the current call source to the value before ``set_call_source``."""
|
|
CALL_SOURCE.reset(token)
|
|
|
|
|
|
@contextmanager
|
|
def call_source_ctx(source: str) -> Iterator[None]:
|
|
"""Temporarily set the call source inside a ``with`` block."""
|
|
token = set_call_source(source)
|
|
try:
|
|
yield
|
|
finally:
|
|
reset_call_source(token)
|
|
|
|
|
|
def get_call_source(default: str = AGENT_SOURCE) -> str:
|
|
"""Return the current call source label."""
|
|
return CALL_SOURCE.get(default)
|
|
|
|
|
|
__all__ = [
|
|
"AGENT_SOURCE",
|
|
"CALL_SOURCE",
|
|
"call_source_ctx",
|
|
"get_call_source",
|
|
"reset_call_source",
|
|
"set_call_source",
|
|
]
|