Spaces:
Running on Zero
Running on Zero
| """Content-addressed on-disk cache of raw model calls. | |
| Every request/response pair is stored as one JSON file keyed by the SHA-256 of the | |
| canonical request. The cache doubles as the raw record of every model call an | |
| experiment made, so re-running an experiment replays identical outputs and a third | |
| party can audit exactly what the model saw and said. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import os | |
| import tempfile | |
| import threading | |
| from datetime import UTC, datetime | |
| from pathlib import Path | |
| from .base import LLMClient, LLMError, LLMRequest, LLMResponse, ModelSpec | |
| def request_key(spec: ModelSpec, request: LLMRequest) -> str: | |
| canonical = json.dumps(request.cache_payload(spec), sort_keys=True, ensure_ascii=False) | |
| return hashlib.sha256(canonical.encode("utf-8")).hexdigest() | |
| class CachedClient: | |
| """Wraps a client; ``offline=True`` turns any cache miss into an error (pure replay).""" | |
| def __init__(self, inner: LLMClient, cache_dir: str | Path, offline: bool = False): | |
| self.inner = inner | |
| self.spec = inner.spec | |
| self.root = Path(cache_dir) | |
| self.offline = offline | |
| self._lock = threading.Lock() | |
| # Identical requests issued concurrently (e.g. segment and provseg asking for the same | |
| # note) wait for one provider call instead of each paying for, and caching, their own. | |
| self._inflight: dict[str, threading.Lock] = {} | |
| def _path(self, key: str) -> Path: | |
| return self.root / key[:2] / f"{key}.json" | |
| def get(self, key: str) -> LLMResponse | None: | |
| p = self._path(key) | |
| if not p.exists(): | |
| return None | |
| rec = json.loads(p.read_text(encoding="utf-8")) | |
| r = rec["response"] | |
| return LLMResponse( | |
| text=r["text"], | |
| input_tokens=r["input_tokens"], | |
| output_tokens=r["output_tokens"], | |
| latency_s=r["latency_s"], | |
| model=r["model"], | |
| stop_reason=r.get("stop_reason"), | |
| cached=True, | |
| cache_key=key, | |
| ) | |
| def put(self, key: str, request: LLMRequest, resp: LLMResponse) -> None: | |
| p = self._path(key) | |
| p.parent.mkdir(parents=True, exist_ok=True) | |
| rec = { | |
| "key": key, | |
| "created_utc": datetime.now(UTC).isoformat(), | |
| "request": request.cache_payload(self.spec), | |
| "response": { | |
| "text": resp.text, | |
| "input_tokens": resp.input_tokens, | |
| "output_tokens": resp.output_tokens, | |
| "latency_s": resp.latency_s, | |
| "model": resp.model, | |
| "stop_reason": resp.stop_reason, | |
| }, | |
| } | |
| fd, tmp = tempfile.mkstemp(dir=p.parent, suffix=".tmp") | |
| with os.fdopen(fd, "w", encoding="utf-8") as f: | |
| json.dump(rec, f, ensure_ascii=False, indent=1) | |
| with self._lock: | |
| os.replace(tmp, p) | |
| def invalidate(self, key: str) -> None: | |
| """Forget a stored response (used when it turned out to be unusable).""" | |
| with self._lock: | |
| self._path(key).unlink(missing_ok=True) | |
| def complete(self, request: LLMRequest) -> LLMResponse: | |
| key = request_key(self.spec, request) | |
| hit = self.get(key) | |
| if hit is not None: | |
| return hit | |
| if self.offline: | |
| raise LLMError(f"offline replay: no cached response for key {key[:12]} (model {self.spec.model})") | |
| with self._lock: | |
| key_lock = self._inflight.setdefault(key, threading.Lock()) | |
| with key_lock: | |
| hit = self.get(key) # another thread may have completed it while we waited | |
| if hit is not None: | |
| return hit | |
| resp = self.inner.complete(request) | |
| self.put(key, request, resp) | |
| resp.cache_key = key | |
| return resp | |