Spaces:
Running on Zero
Running on Zero
File size: 3,860 Bytes
465b477 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """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
|