25 lines
653 B
Python
25 lines
653 B
Python
"""Context propagated to host-provided HTTP transports for cost accounting."""
|
|
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar, Token
|
|
|
|
_model_call_stage: ContextVar[str] = ContextVar(
|
|
"hy_memory_model_call_stage",
|
|
default="other",
|
|
)
|
|
|
|
|
|
def current_model_call_stage() -> str:
|
|
return _model_call_stage.get()
|
|
|
|
|
|
@contextmanager
|
|
def model_call_stage(stage: str | None) -> Iterator[None]:
|
|
normalized = (stage or "other").strip().lower() or "other"
|
|
token: Token[str] = _model_call_stage.set(normalized)
|
|
try:
|
|
yield
|
|
finally:
|
|
_model_call_stage.reset(token)
|