248 lines
8.5 KiB
Python
248 lines
8.5 KiB
Python
"""A disk cache so a corpus is never paid for twice.
|
|
|
|
**Why the key is content-addressed.** A vector is a pure function of three
|
|
things: the model, the input kind, and the exact bytes embedded. Nothing else
|
|
about the record changes the answer. `docs/v1-delivery-plan.md` §4.A proposed
|
|
keying on `chunk_id` + sha256; measured against the real corpus that would
|
|
charge twice for identical text — `chunks.jsonl` holds 15,066 records but only
|
|
14,869 distinct texts, so 197 records (1.31%) are repeats of a text already
|
|
embedded. The key here is `(model_id, input_kind, text_sha256)`, which collapses
|
|
those and, more importantly, cannot silently serve a stale vector after a chunk's
|
|
text is edited: an edit changes the digest, so it is a miss.
|
|
|
|
Traceability is not lost by dropping `chunk_id` from the key. Every cached
|
|
record carries the same sha256 rule (`ports.text_digest`) that produced it, so a
|
|
vector is matched back to its chunk by re-digesting that chunk's text. Pairing
|
|
vectors to chunk records is `load/`'s job, not the cache's.
|
|
|
|
**Why the index holds offsets, not vectors.** 15,066 vectors of 1,024 floats do
|
|
not belong in memory all at once; a list of that many Python floats is tens of
|
|
kilobytes each. Startup scans the file once to map key to byte offset, and a
|
|
`get` seeks and parses exactly one line.
|
|
|
|
The file is append-only. A key already present is never rewritten, so the file
|
|
is a log that can be inspected, truncated, or resumed after an interrupted run
|
|
without a repair step.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, Iterable, List, Optional, Sequence, Tuple
|
|
|
|
from .ports import (
|
|
INPUT_KINDS,
|
|
EmbeddingBatch,
|
|
EmbeddingProvider,
|
|
EmbeddingVector,
|
|
text_digest,
|
|
)
|
|
|
|
CacheKey = Tuple[str, str, str]
|
|
|
|
_KEY_FIELDS = ("model_id", "input_kind", "text_sha256")
|
|
|
|
|
|
def cache_key(model_id: str, input_kind: str, text: str) -> CacheKey:
|
|
return (model_id, input_kind, text_digest(text))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CacheStats:
|
|
hits: int = 0
|
|
misses: int = 0
|
|
|
|
@property
|
|
def lookups(self) -> int:
|
|
return self.hits + self.misses
|
|
|
|
@property
|
|
def hit_rate(self) -> float:
|
|
return self.hits / self.lookups if self.lookups else 0.0
|
|
|
|
|
|
class EmbeddingCache:
|
|
"""Append-only JSONL of vectors, indexed by byte offset."""
|
|
|
|
def __init__(self, path: os.PathLike | str) -> None:
|
|
self._path = Path(path)
|
|
self._offsets: Dict[CacheKey, int] = {}
|
|
self._hits = 0
|
|
self._misses = 0
|
|
if self._path.exists():
|
|
self._build_index()
|
|
|
|
@property
|
|
def path(self) -> Path:
|
|
return self._path
|
|
|
|
@property
|
|
def stats(self) -> CacheStats:
|
|
return CacheStats(hits=self._hits, misses=self._misses)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._offsets)
|
|
|
|
def __contains__(self, key: CacheKey) -> bool:
|
|
return key in self._offsets
|
|
|
|
def _build_index(self) -> None:
|
|
with self._path.open("rb") as handle:
|
|
offset = 0
|
|
for raw in handle:
|
|
line = raw.decode("utf-8").strip()
|
|
if line:
|
|
record = json.loads(line)
|
|
self._offsets[self._key_of(record)] = offset
|
|
offset += len(raw)
|
|
|
|
@staticmethod
|
|
def _key_of(record: dict) -> CacheKey:
|
|
missing = [f for f in _KEY_FIELDS if not record.get(f)]
|
|
if missing:
|
|
raise ValueError(f"cache record is missing key fields: {missing}")
|
|
return (
|
|
record["model_id"],
|
|
record["input_kind"],
|
|
record["text_sha256"],
|
|
)
|
|
|
|
def get(self, key: CacheKey) -> Optional[EmbeddingVector]:
|
|
offset = self._offsets.get(key)
|
|
if offset is None:
|
|
self._misses += 1
|
|
return None
|
|
with self._path.open("rb") as handle:
|
|
handle.seek(offset)
|
|
record = json.loads(handle.readline().decode("utf-8"))
|
|
self._hits += 1
|
|
return _vector_from_record(record)
|
|
|
|
def put(self, vector: EmbeddingVector) -> bool:
|
|
"""Append a vector. Returns False if the key was already stored."""
|
|
key = (vector.model_id, vector.input_kind, vector.text_sha256)
|
|
if key in self._offsets:
|
|
return False
|
|
self._path.parent.mkdir(parents=True, exist_ok=True)
|
|
line = json.dumps(_record_from_vector(vector), ensure_ascii=False) + "\n"
|
|
encoded = line.encode("utf-8")
|
|
with self._path.open("ab") as handle:
|
|
offset = handle.tell()
|
|
handle.write(encoded)
|
|
self._offsets[key] = offset
|
|
return True
|
|
|
|
def put_many(self, vectors: Iterable[EmbeddingVector]) -> int:
|
|
return sum(1 for vector in vectors if self.put(vector))
|
|
|
|
|
|
def _record_from_vector(vector: EmbeddingVector) -> dict:
|
|
return {
|
|
"model_id": vector.model_id,
|
|
"input_kind": vector.input_kind,
|
|
"text_sha256": vector.text_sha256,
|
|
"provider": vector.provider,
|
|
"dimensions": vector.dimensions,
|
|
"normalized": vector.normalized,
|
|
"input_token_count": vector.input_token_count,
|
|
"values": vector.values,
|
|
}
|
|
|
|
|
|
def _vector_from_record(record: dict) -> EmbeddingVector:
|
|
values: List[float] = record["values"]
|
|
declared = record["dimensions"]
|
|
if len(values) != declared:
|
|
raise ValueError(
|
|
f"cached vector for {record['text_sha256'][:12]} has {len(values)} "
|
|
f"values but declares {declared} dimensions"
|
|
)
|
|
return EmbeddingVector(
|
|
values=values,
|
|
text_sha256=record["text_sha256"],
|
|
provider=record["provider"],
|
|
model_id=record["model_id"],
|
|
dimensions=declared,
|
|
input_kind=record["input_kind"],
|
|
normalized=record.get("normalized"),
|
|
input_token_count=record.get("input_token_count"),
|
|
)
|
|
|
|
|
|
class CachingEmbeddingProvider(EmbeddingProvider):
|
|
"""Wraps a provider so only uncached texts reach it.
|
|
|
|
A decorator rather than a change to each adapter: the three existing
|
|
providers stay unaware that a cache exists, and a fourth needs no cache code
|
|
to benefit. `request_count` counts requests the *inner* provider actually
|
|
made, which is what makes "a second run costs nothing" a checkable claim
|
|
rather than an assertion — a fully cached run reports zero.
|
|
"""
|
|
|
|
def __init__(self, inner: EmbeddingProvider, cache: EmbeddingCache) -> None:
|
|
self._inner = inner
|
|
self._cache = cache
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return f"cached:{self._inner.name}"
|
|
|
|
@property
|
|
def model_id(self) -> str:
|
|
return self._inner.model_id
|
|
|
|
@property
|
|
def dimensions(self) -> int:
|
|
return self._inner.dimensions
|
|
|
|
@property
|
|
def max_batch_size(self) -> int:
|
|
return self._inner.max_batch_size
|
|
|
|
@property
|
|
def cache(self) -> EmbeddingCache:
|
|
return self._cache
|
|
|
|
def _embed_batch(
|
|
self, texts: Sequence[str], input_kind: str
|
|
) -> List[EmbeddingVector]:
|
|
return list(self._inner.embed(list(texts), input_kind).vectors)
|
|
|
|
def embed(self, texts: Sequence[str], input_kind: str) -> EmbeddingBatch:
|
|
if input_kind not in INPUT_KINDS:
|
|
raise ValueError(
|
|
f"input_kind must be one of {INPUT_KINDS}, got {input_kind!r}"
|
|
)
|
|
if any(not t.strip() for t in texts):
|
|
raise ValueError("refusing to embed an empty or whitespace-only text")
|
|
|
|
resolved: List[Optional[EmbeddingVector]] = []
|
|
pending: Dict[str, List[int]] = {}
|
|
for position, text in enumerate(texts):
|
|
hit = self._cache.get(cache_key(self.model_id, input_kind, text))
|
|
resolved.append(hit)
|
|
if hit is None:
|
|
pending.setdefault(text, []).append(position)
|
|
|
|
requests = 0
|
|
latency_ms = 0.0
|
|
if pending:
|
|
wanted = list(pending)
|
|
batch = self._inner.embed(wanted, input_kind)
|
|
requests = batch.request_count
|
|
latency_ms = batch.latency_ms
|
|
for text, vector in zip(wanted, batch.vectors, strict=True):
|
|
self._cache.put(vector)
|
|
for position in pending[text]:
|
|
resolved[position] = vector
|
|
|
|
if any(vector is None for vector in resolved):
|
|
raise ValueError("cache resolution left a text without a vector")
|
|
return EmbeddingBatch(
|
|
vectors=[vector for vector in resolved if vector is not None],
|
|
request_count=requests,
|
|
latency_ms=latency_ms,
|
|
)
|