llm-cache

@dataclass(frozen=True)
class GetResponse:
  entry: CacheEntry | None = None
  embedding: Embedding | None = None


class Cache:
  def __init__(
    self,
    store: CacheStore,
    embedder: Embedder,
    *,
    similarity_threshold: float = 0.95,
  ) -> None:
    self.store = store
    self.embedder = embedder
    self.similarity_threshold = similarity_threshold

  def get(self, scope: str, query: str, fingerprint: str | None = None) -> GetResponse:
    """Look the query up, degrading to a miss if the cache itself is broken.

    A cache that raises turns its own outage into the caller's outage, so a
    StoreError or EmbeddingError is logged and reported as an empty result,
    the caller then does what it would have done on any miss.
    """
    try:
      if fingerprint:
        entry = self.store.get_exact(scope=scope, fingerprint=fingerprint)
        if entry:
          return GetResponse(
            entry=entry,
          )

      embedding = self.embedder.embed(query)
      entries = self.store.search(scope=scope, limit=1, embedding=embedding)

      if entries:
        entry, similarity = entries[0]
        if similarity >= self.similarity_threshold:
          return GetResponse(entry=entry, embedding=embedding)

      return GetResponse(embedding=embedding)
    except CacheError:
      log.warning('cache lookup failed, treating as a miss', exc_info=True)
      return GetResponse()

  def set(self, entry: CacheEntry, embedding: Embedding | None = None) -> None:
    """Store the entry, or give up quietly.

    Failing to write a cache entry costs a future hit and nothing else, so it
    never propagates to the caller.
    """
    try:
      self.store.put(
        entry=entry, embedding=embedding or self.embedder.embed(entry.query)
      )
    except CacheError:
      log.warning('cache write failed, entry not stored', exc_info=True)

import hashlib


def fingerprint(query: str, **params: object) -> str:
  normalized = ' '.join(query.lower().split())
  material = repr((normalized, sorted(params.items())))
  return hashlib.sha256(material.encode()).hexdigest()

from collections.abc import Mapping
from typing import Protocol

from llm_cache.models import CacheEntry, Embedding


class CacheStore(Protocol):
  def get_exact(self, scope: str, fingerprint: str) -> CacheEntry | None: ...
  def search(
    self, scope: str, embedding: Embedding, limit: int
  ) -> list[tuple[CacheEntry, float]]: ...
  def put(self, entry: CacheEntry, embedding: Embedding) -> None: ...
  def touch(self, scope: str, fingerprint: str) -> None: ...
  def touch_many(self, counts: Mapping[tuple[str, str], int]) -> None: ...
  def evict_expired(self) -> int: ...
  def purge(self) -> int: ...
  def clear(self, scope: str) -> int: ...

  @contextmanager
  def _connection(self) -> Generator[psycopg.Connection, None, None]:
    """Borrow a pooled connection, reporting infrastructure faults as StoreError.

    Only the "database is unreachable" family is translated. Programming errors
    — a missing table, mismatched vector dimensions, stay as psycopg errors so
    they surface as bugs instead of being downgraded to a cache miss.
    """
    try:
      with self._pool.connection() as conn:
        yield conn
    except (psycopg.OperationalError, PoolTimeout) as exc:
      raise StoreError('cache store unavailable') from exc

  # get_exact() is the same shape: one indexed lookup on (scope, fingerprint)
  # with the expires_at filter applied in SQL.

  def search(
    self, scope: str, embedding: Embedding, limit: int
  ) -> list[tuple[CacheEntry, float]]:
    """Nearest neighbours in `scope`, best first, as (entry, similarity).

    Returns the top `limit` regardless of how far away they are, deciding what
    counts as close enough is the caller's threshold to apply, not the store's.
    """
    query = sql.SQL(
      'SELECT {columns}, 1 - (embedding <=> %s::vector) AS similarity '
      'FROM {table} '
      'WHERE scope = %s AND (expires_at IS NULL OR expires_at > now()) '
      'ORDER BY embedding <=> %s::vector '
      'LIMIT %s'
    ).format(columns=_COLUMNS, table=self._table)
    vector = list(embedding)
    with self._connection() as conn:
      with conn.cursor(row_factory=dict_row) as cur:
        cur.execute(query, (vector, scope, vector, limit))
        return [
          (CacheEntry(**{k: row[k] for k in _ENTRY_COLUMN_NAMES}), row['similarity'])
          for row in cur.fetchall()
        ]

class Embedder(Protocol):
  """Turns query text into a vector.

  Structural, so anything with a matching `embed` works: an OpenAI client, a
  local sentence-transformers model, or a stub in tests.
  """

  dimensions: int

  def embed(self, text: str) -> Embedding: ...


class OpenAIEmbedder:
  def __init__(
    self,
    *,
    model: str = 'text-embedding-3-small',
    dimensions: int = 1536,
    api_key: str | None = None,
  ) -> None:
    """`dimensions` must match the vector width in the migration.

    pgvector fixes the width at the column, so changing this means a new
    migration, not just a different argument here.
    """
    from openai import OpenAI

    self._client = OpenAI(api_key=api_key)
    self._model = model
    self.dimensions = dimensions

  def embed(self, text: str) -> Embedding:
    response = self._client.embeddings.create(
      model=self._model, input=text, dimensions=self.dimensions
    )
    return response.data[0].embedding

class TouchBuffer:
  """Accumulates hit counts in memory and writes them out on a timer.

  Nothing on the read path consults `hits` or `last_used_at`, so recording a hit
  should never sit between the caller and their response. Counting in memory
  also collapses repeat hits on the same entry into a single `hits + n`, which
  matters because each row update Postgres cannot do in place costs an entry in
  the HNSW index.

  The trade is durability: up to one interval of counts is lost if the process
  dies. These are cache statistics, so nothing reconciles against them.
  """

  def touch(self, scope: str, fingerprint: str) -> None:
    """Record a hit. Never does I/O."""
    with self._lock:
      self._counts[(scope, fingerprint)] += 1

  def flush(self) -> None:
    """Write pending counts. Safe to call directly, e.g. in tests."""
    with self._lock:
      if not self._counts:
        return
      batch = self._counts
      self._counts = Counter()
    try:
      self._store.touch_many(batch)
    except CacheError:
      # Deliberately dropped rather than merged back: a store that stays down
      # would otherwise grow this buffer without bound, and losing hit counts
      # costs nothing but eviction accuracy.
      log.warning('touch flush failed, %d counts dropped', len(batch), exc_info=True)

class Maintenance:
  """Runs the cache's periodic housekeeping on one background thread.

  None of this belongs on the request path. `get_exact` and `search` already
  filter on `expires_at`, so an expired entry is invisible whether or not it has
  been deleted, eviction only reclaims space. Calling it from `set()` would
  make one unlucky request pay for a table-wide DELETE.

  When a TouchBuffer is passed in, its flush becomes one of these jobs and the
  buffer should not also be started on its own; one scheduler is easier to
  reason about than three threads with independent lifetimes.
  """

  def __init__(
    self,
    store: CacheStore,
    *,
    buffer: TouchBuffer | None = None,
    flush_interval: float = 5.0,
    evict_interval: float = 300.0,
    purge_interval: float = 300.0,
  ) -> None:
    self._buffer = buffer
    self._jobs: list[_Job] = []
    if buffer is not None:
      self._jobs.append(_Job('flush', buffer.flush, flush_interval))
    self._jobs.append(_Job('evict_expired', store.evict_expired, evict_interval))
    self._jobs.append(_Job('purge', store.purge, purge_interval))
    self._stopping = threading.Event()
    self._thread: threading.Thread | None = None

class Settings(BaseSettings):
  """Environment-backed configuration, read once at the application edge.

  Nothing inside the package constructs this, `Cache`, `Postgres` and the
  embedders all take explicit arguments, so a library user is never subject to
  whatever happens to be in the process environment. Applications build one of
  these at startup and pass the values down.
  """

  model_config = SettingsConfigDict(
    env_prefix='LLM_CACHE_',
    env_file='.env',
    env_file_encoding='utf-8',
    extra='ignore',
  )

  dsn: str = 'postgresql://llmcache:llmcache@localhost:5432/llmcache'
  table: str = 'llm_cache'
  namespace: str = 'default'

  # Must match vector(n) in the migration; pgvector fixes the width at the
  # column, so a change here without a migration fails at the first write.
  embedding_dimensions: int = Field(default=1536, gt=0)

  # Cosine similarity, so bounded at 1.0. Below roughly 0.9 unrelated questions
  # start matching, which is worse than missing.
  similarity_threshold: float = Field(default=0.95, ge=0.0, le=1.0)

  search_limit: int = Field(default=5, gt=0)

  # Blank in .env means no expiry rather than zero seconds.
  ttl_seconds: OptionalSeconds = None