Files
Artemis-Eapil/artemis/llm/cache_affinity.py
T

50 lines
2.2 KiB
Python

"""Per-conversation routing hints shared by agent and auxiliary Responses calls."""
from collections.abc import Mapping
import hashlib
import re
from typing import Any
from uuid import UUID
def cache_affinity_options(
ctx: Any, role: str | None, model_name: str | None, options: dict | None = None
) -> dict:
"""Keep body and header routing keys aligned without changing model defaults.
An exploration survives multiple execution sessions. Each model role keeps
its own cache scope so background summaries do not compete with the agent.
Set hints per call: raw models/HTTP clients may be shared by concurrent tasks.
Explicit caller keys and headers take precedence (including underscore aliases).
A caller-supplied routing header also supplies a missing body key, so gateways
reading different key sources agree on the same affinity whenever possible.
"""
result = dict(options or {})
existing = result.get("extra_headers")
if "prompt_cache_key" not in result and isinstance(existing, Mapping):
names = {name.lower(): value for name, value in existing.items()}
for name in ("thread-id", "thread_id", "session-id", "session_id"):
value = names.get(name)
if isinstance(value, str) and re.fullmatch(r"[-0-9A-Za-z._:]{1,128}", value):
result["prompt_cache_key"] = value
break
engine = getattr(ctx, "data_engine", None)
session = getattr(ctx, "exploration_id", None) or getattr(engine, "current_session_id", None)
if isinstance(session, (str, UUID)) and str(session):
identity = f"{session}:{role or 'model'}:{model_name}"
result.setdefault(
"prompt_cache_key", "artemis:" + hashlib.sha256(identity.encode()).hexdigest()[:40]
)
key = result.get("prompt_cache_key")
if not isinstance(key, str) or not re.fullmatch(r"[-0-9A-Za-z._:]{1,128}", key):
return result
if existing is not None and not isinstance(existing, Mapping):
return result
headers = dict(existing or {})
names = {name.lower().replace("_", "-") for name in headers}
for name in ("Thread-Id", "Session-Id"):
if name.lower() not in names:
headers[name] = key
result["extra_headers"] = headers
return result