Files
Artemis-Eapil/artemis/llm/prefix_cache.py

66 lines
2.5 KiB
Python

"""Local prefix diagnostics; retain hashes only, never prompt/image payloads.
These compare adapted messages, not the gateway's tokenized wire prefix.
Only provider-reported cached_tokens is a measured cache hit.
"""
import hashlib
import json
from langchain_core.messages import BaseMessage
def _digest(value) -> str:
return hashlib.sha256(
json.dumps(value, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")
).hexdigest()
class PromptPrefixTracker:
def __init__(self):
self._previous: dict[str, tuple[str, tuple[str, ...]]] = {}
def observe(self, identity: str, messages, options: dict) -> dict:
identity = _digest(identity)
if not isinstance(messages, (list, tuple)) or not all(
isinstance(message, BaseMessage) for message in messages
):
return {"comparison": "unavailable"}
fingerprints = tuple(
_digest(
{
"role": message.type.removesuffix("MessageChunk").lower(),
"content": message.content,
"name": message.name,
"tool_calls": getattr(message, "tool_calls", None),
"tool_call_id": getattr(message, "tool_call_id", None),
}
)
for message in messages
)
options_hash = _digest(options)
previous = self._previous.get(identity)
# Task-scoped and bounded even if a caller tries many endpoint roles.
if identity not in self._previous and len(self._previous) >= 32:
self._previous.pop(next(iter(self._previous)))
self._previous[identity] = (options_hash, fingerprints)
common = 0
if previous:
for old, new in zip(previous[1], fingerprints):
if old != new:
break
common += 1
schema_changed = previous is not None and previous[0] != options_hash
return {
"comparison": "local_message_fingerprints",
"message_count": len(fingerprints),
"previous_message_count": len(previous[1]) if previous else None,
"shared_prefix_messages": common if previous else None,
"first_changed_message": (
common + 1 if previous and common < len(previous[1]) else None
),
"append_only": bool(previous and common == len(previous[1]) and not schema_changed),
"tools_or_format_changed": schema_changed,
"cold_start": previous is None,
}