"""Context variable for propagating user/client/presentation info into LLM calls.""" import contextvars import uuid from dataclasses import dataclass, field from typing import Optional @dataclass class LLMCallContext: user_id: Optional[uuid.UUID] = field(default=None) client_id: Optional[uuid.UUID] = field(default=None) presentation_id: Optional[uuid.UUID] = field(default=None) _llm_context: contextvars.ContextVar[Optional[LLMCallContext]] = contextvars.ContextVar( "llm_call_context", default=None ) def set_llm_context(ctx: LLMCallContext) -> None: _llm_context.set(ctx) def get_llm_context() -> Optional[LLMCallContext]: return _llm_context.get() def clear_llm_context() -> None: _llm_context.set(None)