98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from collections import OrderedDict
|
|
from datetime import datetime, timezone
|
|
from threading import Lock
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from src.config import ModelConfig, PromptCachePolicy
|
|
|
|
__all__ = [
|
|
"GeminiCacheHandle",
|
|
"InMemoryGeminiCacheStore",
|
|
"PromptCachePolicy",
|
|
"build_cache_key",
|
|
"gemini_cache_store",
|
|
]
|
|
|
|
|
|
class GeminiCacheHandle(BaseModel):
|
|
key: str
|
|
cached_content_name: str
|
|
expires_at: datetime
|
|
|
|
|
|
def build_cache_key(
|
|
*,
|
|
config: ModelConfig,
|
|
cache_policy: PromptCachePolicy,
|
|
cacheable_messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]] | None,
|
|
system_instruction: str | None = None,
|
|
tool_config: dict[str, Any] | None = None,
|
|
) -> str:
|
|
"""Deterministic key over the cacheable shape of a request.
|
|
|
|
``system_instruction`` and ``tool_config`` must be part of the key
|
|
because the provider's cached-content handle captures them at creation
|
|
time — two requests that differ only by system prompt or tool
|
|
constraints would otherwise hit the same cached handle and silently get
|
|
the wrong system prompt / tool policy.
|
|
"""
|
|
payload = {
|
|
"transport": config.transport,
|
|
"model": config.model,
|
|
"cache_policy": cache_policy.model_dump(mode="json"),
|
|
"messages": cacheable_messages,
|
|
"tools": tools,
|
|
"system_instruction": system_instruction,
|
|
"tool_config": tool_config,
|
|
}
|
|
encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
|
digest = hashlib.sha256(encoded.encode("utf-8")).hexdigest()
|
|
return f"llm-cache:{cache_policy.key_version}:{digest}"
|
|
|
|
|
|
class InMemoryGeminiCacheStore:
|
|
"""Best-effort local cache for Gemini cached-content handles.
|
|
|
|
Uses LRU eviction with a max entry limit to prevent unbounded growth.
|
|
"""
|
|
|
|
MAX_ENTRIES: int = 1024
|
|
|
|
def __init__(self) -> None:
|
|
self._handles: OrderedDict[str, GeminiCacheHandle] = OrderedDict()
|
|
self._lock: Lock = Lock()
|
|
|
|
def get(self, key: str) -> GeminiCacheHandle | None:
|
|
with self._lock:
|
|
handle = self._handles.get(key)
|
|
if handle is None:
|
|
return None
|
|
if handle.expires_at <= datetime.now(timezone.utc):
|
|
self._handles.pop(key, None)
|
|
return None
|
|
self._handles.move_to_end(key)
|
|
return handle
|
|
|
|
def set(self, handle: GeminiCacheHandle) -> GeminiCacheHandle:
|
|
with self._lock:
|
|
now = datetime.now(timezone.utc)
|
|
expired = [k for k, h in self._handles.items() if h.expires_at <= now]
|
|
for k in expired:
|
|
self._handles.pop(k, None)
|
|
if handle.key in self._handles:
|
|
self._handles.move_to_end(handle.key)
|
|
self._handles[handle.key] = handle
|
|
while len(self._handles) > self.MAX_ENTRIES:
|
|
self._handles.popitem(last=False)
|
|
return handle
|
|
|
|
|
|
gemini_cache_store = InMemoryGeminiCacheStore()
|