cachellm-proxy 0.1.0__py3-none-any.whl

cachellm/cache/service.py

@@ -0,0 +1,332 @@
+"""The cache orchestrator: one lookup path, one store path, one invalidate path.
+
+Everything the proxy knows about caching lives behind this class, so the HTTP
+layer stays thin and the same logic can be reused by an in-process client
+wrapper without a server.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+import numpy as np
+import structlog
+
+from cachellm.cache.analytics import Analytics, NearMiss
+from cachellm.cache.entry import CacheEntry
+from cachellm.cache.exact_store import ExactStore
+from cachellm.cache.keys import (
+    embedding_text,
+    exact_hash,
+    namespace_for,
+    prompt_fingerprint,
+)
+from cachellm.cache.keys import (
+    entry_id as make_entry_id,
+)
+from cachellm.cache.policy import PolicyDecision, decide
+from cachellm.cache.vector_store import VectorStore
+from cachellm.embeddings.base import Embedder
+from cachellm.models import ChatCompletionRequest
+from cachellm.pricing import estimate_cost
+from cachellm.settings import Settings
+
+log = structlog.get_logger(__name__)
+
+LookupStatus = Literal["hit", "miss", "bypass", "shadow_hit"]
+
+
+@dataclass
+class LookupResult:
+    status: LookupStatus
+    decision: PolicyDecision
+    namespace: str = ""
+    similarity: float = 0.0
+    tier: str = ""  # "exact" | "semantic"
+    entry: CacheEntry | None = None
+    embedding: np.ndarray | None = None
+    exact: str = ""
+    cache_text: str = ""
+    lookup_ms: float = 0.0
+    neighbours: list[tuple[str, float]] = field(default_factory=list)
+
+    @property
+    def served_from_cache(self) -> bool:
+        return self.status == "hit" and self.entry is not None
+
+    @property
+    def category(self) -> str:
+        return self.decision.category
+
+    def saved_usd(self, model: str) -> float:
+        if self.entry is None:
+            return 0.0
+        return estimate_cost(model, self.entry.prompt_tokens, self.entry.completion_tokens)
+
+
+class CacheService:
+    def __init__(
+        self,
+        *,
+        settings: Settings,
+        embedder: Embedder,
+        vectors: VectorStore,
+        exact: ExactStore,
+        analytics: Analytics,
+    ) -> None:
+        self.settings = settings
+        self.embedder = embedder
+        self.vectors = vectors
+        self.exact = exact
+        self.analytics = analytics
+
+    # ------------------------------------------------------------------ lookup
+    async def lookup(
+        self,
+        request: ChatCompletionRequest,
+        provider_name: str,
+        *,
+        cache_control: str = "",
+    ) -> LookupResult:
+        started = time.perf_counter()
+        decision = decide(request, self.settings, cache_control=cache_control)
+        if not decision.cacheable:
+            return LookupResult(
+                status="bypass", decision=decision, lookup_ms=(time.perf_counter() - started) * 1000
+            )
+
+        namespace = namespace_for(request, provider_name)
+        raw_text = request.cache_text(include_history=self.settings.cache_multi_turn)
+        exact_key = exact_hash(raw_text)
+
+        # --- level 1: exact match, no embedding needed -----------------------
+        entry_id = await self.exact.get(namespace, exact_key)
+        if entry_id:
+            entry = await self.vectors.get(entry_id)
+            if entry is not None:
+                status: LookupStatus = "shadow_hit" if self.settings.shadow_mode else "hit"
+                return LookupResult(
+                    status=status,
+                    decision=decision,
+                    namespace=namespace,
+                    similarity=1.0,
+                    tier="exact",
+                    entry=entry,
+                    exact=exact_key,
+                    cache_text=raw_text,
+                    lookup_ms=(time.perf_counter() - started) * 1000,
+                )
+            # Pointer outlived its entry (entry TTL shorter, or manual delete).
+            await self.exact.delete(namespace, exact_key)
+
+        # --- level 2: semantic nearest neighbour -----------------------------
+        text_to_embed = embedding_text(raw_text, strip_fillers=self.settings.strip_filler_words)
+        vector = await self.embedder.embed(text_to_embed)
+        matches = await self.vectors.search(namespace, vector, self.settings.top_k)
+        neighbours = [(m.entry_id, round(sim, 4)) for m, sim in matches]
+
+        if matches:
+            best_entry, best_sim = matches[0]
+            if best_sim >= decision.threshold:
+                status = "shadow_hit" if self.settings.shadow_mode else "hit"
+                return LookupResult(
+                    status=status,
+                    decision=decision,
+                    namespace=namespace,
+                    similarity=best_sim,
+                    tier="semantic",
+                    entry=best_entry,
+                    embedding=vector,
+                    exact=exact_key,
+                    cache_text=raw_text,
+                    neighbours=neighbours,
+                    lookup_ms=(time.perf_counter() - started) * 1000,
+                )
+            if best_sim >= decision.threshold - self.settings.near_miss_margin:
+                await self.analytics.record_near_miss(
+                    NearMiss(
+                        prompt=raw_text,
+                        matched_prompt=best_entry.prompt,
+                        similarity=round(best_sim, 4),
+                        threshold=decision.threshold,
+                        category=decision.category,
+                        namespace=namespace,
+                        model=request.model,
+                        at=time.time(),
+                    )
+                )
+                await self.analytics.incr("near_misses")
+
+        return LookupResult(
+            status="miss",
+            decision=decision,
+            namespace=namespace,
+            similarity=matches[0][1] if matches else 0.0,
+            embedding=vector,
+            exact=exact_key,
+            cache_text=raw_text,
+            neighbours=neighbours,
+            lookup_ms=(time.perf_counter() - started) * 1000,
+        )
+
+    # ------------------------------------------------------------------- store
+    async def store(
+        self,
+        *,
+        lookup: LookupResult,
+        request: ChatCompletionRequest,
+        provider_name: str,
+        response_text: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        finish_reason: str,
+        tool_calls: list[dict[str, Any]] | None = None,
+    ) -> str | None:
+        """Persist a completed answer. Only clean, finished responses are stored."""
+        if not lookup.decision.cacheable or lookup.status == "bypass":
+            return None
+
+        complete = finish_reason in ("stop", "end_turn", "eos", "stop_sequence", "complete")
+        if not complete and not self.settings.cache_truncated:
+            # The most common cause of a mysteriously low hit rate: max_tokens is
+            # set below what the model wants to say, every answer comes back
+            # truncated, and nothing is cacheable. Counted so /admin/stats can
+            # say so out loud instead of leaving it to be guessed at.
+            await self.analytics.incr("stores_skipped_truncated")
+            log.debug("skip_store_incomplete", finish_reason=finish_reason)
+            return None
+        if not response_text.strip():
+            await self.analytics.incr("stores_skipped_empty")
+            return None
+
+        vector = lookup.embedding
+        if vector is None:
+            text_to_embed = embedding_text(
+                lookup.cache_text
+                or request.cache_text(include_history=self.settings.cache_multi_turn),
+                strip_fillers=self.settings.strip_filler_words,
+            )
+            vector = await self.embedder.embed(text_to_embed)
+
+        namespace = lookup.namespace or namespace_for(request, provider_name)
+        exact_key = lookup.exact or exact_hash(lookup.cache_text)
+        eid = make_entry_id(namespace, exact_key)
+        entry = CacheEntry(
+            entry_id=eid,
+            namespace=namespace,
+            category=lookup.decision.category,
+            model=request.model,
+            provider=provider_name,
+            prompt=lookup.cache_text or request.last_user_text(),
+            response_text=response_text,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            finish_reason=finish_reason,
+            ttl=lookup.decision.ttl,
+            exact_hash=exact_key,
+            embedding=vector,
+            tool_calls=tool_calls,
+        )
+        await self.vectors.put(entry)
+        await self.exact.put(namespace, exact_key, eid, lookup.decision.ttl)
+        await self.analytics.incr("entries_written")
+        log.debug(
+            "cache_store",
+            entry_id=eid,
+            category=entry.category,
+            fingerprint=prompt_fingerprint(entry.prompt),
+            ttl=entry.ttl,
+        )
+        return eid
+
+    # -------------------------------------------------------------- maintenance
+    async def register_hit(self, lookup: LookupResult, model: str) -> float:
+        """Bump hit bookkeeping and return the modelled dollars saved."""
+        if lookup.entry is None:
+            return 0.0
+        await self.vectors.touch_hit(lookup.entry.entry_id)
+        saved = lookup.saved_usd(model)
+        await self.analytics.bulk(
+            {
+                "hits": 1,
+                f"hits_{lookup.tier}": 1,
+                "tokens_saved_prompt": lookup.entry.prompt_tokens,
+                "tokens_saved_completion": lookup.entry.completion_tokens,
+            },
+            {"usd_saved": saved},
+        )
+        return saved
+
+    async def invalidate(
+        self, *, namespace: str | None = None, model: str | None = None, drop_all: bool = False
+    ) -> int:
+        removed = await self.vectors.invalidate(namespace=namespace, model=model, drop_all=drop_all)
+        await self.analytics.incr("invalidations")
+        log.info(
+            "cache_invalidate", namespace=namespace, model=model, drop_all=drop_all, removed=removed
+        )
+        return removed
+
+    async def stats(self) -> dict[str, Any]:
+        counters = await self.analytics.counters()
+        requests = counters.get("requests", 0.0)
+        hits = counters.get("hits", 0.0)
+        cacheable = counters.get("cacheable_requests", 0.0)
+        misses = counters.get("misses", 0.0)
+        skipped_truncated = counters.get("stores_skipped_truncated", 0.0)
+        diagnostics: list[str] = []
+        if misses and skipped_truncated / misses > 0.2:
+            diagnostics.append(
+                f"{skipped_truncated / misses:.0%} of misses were not cached because the "
+                f"response hit max_tokens. Raise max_tokens so answers finish, or set "
+                f"CACHELLM_CACHE_TRUNCATED=true to cache truncated answers anyway."
+            )
+        if requests and counters.get("bypass", 0.0) / requests > 0.5:
+            diagnostics.append(
+                "Over half of requests bypassed the cache. Check X-Cache-Bypass-Reason: "
+                "temperature, tool calls, JSON mode and multi-turn are skipped by default."
+            )
+        if self.settings.embedding_backend == "hash":
+            diagnostics.append(
+                "Running the hash embedding backend. It is a deterministic test double "
+                "with no understanding of meaning, so the semantic tier will rarely fire. "
+                "Set CACHELLM_EMBEDDING_BACKEND=fastembed for real matching."
+            )
+        elif not self.settings.is_calibrated:
+            diagnostics.append(
+                f"No measured threshold for embedding model "
+                f"{self.settings.embedding_model!r}; using {self.settings.calibrated_threshold}. "
+                f"Run `python -m bench.compare_models` on your own pairs."
+            )
+
+        return {
+            "entries": await self.vectors.count(),
+            "requests": int(requests),
+            "cacheable_requests": int(cacheable),
+            "hits": int(hits),
+            "hits_exact": int(counters.get("hits_exact", 0.0)),
+            "hits_semantic": int(counters.get("hits_semantic", 0.0)),
+            "misses": int(counters.get("misses", 0.0)),
+            "bypassed": int(counters.get("bypass", 0.0)),
+            "shadow_hits": int(counters.get("shadow_hits", 0.0)),
+            "near_misses": int(counters.get("near_misses", 0.0)),
+            "coalesced": int(counters.get("coalesced", 0.0)),
+            "errors": int(counters.get("provider_errors", 0.0)),
+            "hit_rate": round(hits / requests, 4) if requests else 0.0,
+            "hit_rate_of_cacheable": round(hits / cacheable, 4) if cacheable else 0.0,
+            "stores_skipped_truncated": int(skipped_truncated),
+            "stores_skipped_empty": int(counters.get("stores_skipped_empty", 0.0)),
+            "diagnostics": diagnostics,
+            "usd_saved": round(counters.get("usd_saved", 0.0), 6),
+            "usd_spent": round(counters.get("usd_spent", 0.0), 6),
+            "tokens_saved": int(
+                counters.get("tokens_saved_prompt", 0.0)
+                + counters.get("tokens_saved_completion", 0.0)
+            ),
+            "latency_ms": {
+                "hit": await self.analytics.latency_percentiles("hit"),
+                "miss": await self.analytics.latency_percentiles("miss"),
+            },
+        }

cachellm/cache/vector_store.py

@@ -0,0 +1,217 @@
+"""Level 2: semantic nearest-neighbour lookup backed by Redis + RedisVL.
+
+The index is HNSW over cosine distance with a ``namespace`` tag filter, so a
+query only ever considers entries that are allowed to answer it. Redis returns
+cosine *distance*; because every vector is unit length, similarity is simply
+``1 - distance``.
+
+Entries are written with the plain Redis client rather than through RedisVL so
+that the hash write and its TTL land in one pipeline, and so an expiring key
+takes itself out of the index with no sweeper process.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, cast
+
+import numpy as np
+import redis.asyncio as aioredis
+import structlog
+from redis.exceptions import RedisError
+from redisvl.index import AsyncSearchIndex
+from redisvl.query import FilterQuery, VectorQuery
+from redisvl.query.filter import Tag
+
+from cachellm.cache.entry import CacheEntry
+from cachellm.settings import Settings
+
+log = structlog.get_logger(__name__)
+
+RETURN_FIELDS = [
+    "entry_id",
+    "namespace",
+    "category",
+    "model",
+    "provider",
+    "prompt",
+    "response_text",
+    "prompt_tokens",
+    "completion_tokens",
+    "finish_reason",
+    "created_at",
+    "ttl",
+    "hits",
+    "exact_hash",
+    "tool_calls",
+]
+
+
+class VectorStore:
+    def __init__(self, redis: aioredis.Redis, settings: Settings) -> None:
+        self._redis = redis
+        self._settings = settings
+        self._prefix = settings.entry_prefix
+        self._index: AsyncSearchIndex | None = None
+
+    # ------------------------------------------------------------------ schema
+    def schema(self) -> dict[str, Any]:
+        return {
+            "index": {
+                "name": self._settings.index_name,
+                "prefix": self._prefix,
+                "storage_type": "hash",
+            },
+            "fields": [
+                {"name": "namespace", "type": "tag"},
+                {"name": "category", "type": "tag"},
+                {"name": "model", "type": "tag"},
+                {"name": "provider", "type": "tag"},
+                {"name": "created_at", "type": "numeric", "attrs": {"sortable": True}},
+                {
+                    "name": "embedding",
+                    "type": "vector",
+                    "attrs": {
+                        "dims": self._settings.expected_dim(),
+                        "distance_metric": "cosine",
+                        "algorithm": "hnsw",
+                        "datatype": "float32",
+                    },
+                },
+            ],
+        }
+
+    async def connect(self) -> None:
+        index = AsyncSearchIndex.from_dict(self.schema(), redis_client=self._redis)
+        await index.create(overwrite=False)
+        self._index = index
+        log.info(
+            "vector_index_ready",
+            index=self._settings.index_name,
+            dims=self._settings.expected_dim(),
+        )
+
+    @property
+    def index(self) -> AsyncSearchIndex:
+        if self._index is None:
+            raise RuntimeError("VectorStore.connect() has not been awaited")
+        return self._index
+
+    def key(self, entry_id: str) -> str:
+        return f"{self._prefix}{entry_id}"
+
+    # ------------------------------------------------------------------- writes
+    async def put(self, entry: CacheEntry) -> None:
+        key = self.key(entry.entry_id)
+        # redis-py types the mapping key as a wide union, and Mapping is
+        # invariant in its key type, so a plain dict[str, ...] will not match.
+        # The values really are Redis-compatible; this is a variance nit.
+        mapping = cast("Mapping[Any, Any]", entry.to_redis_mapping())
+        async with self._redis.pipeline(transaction=True) as pipe:
+            pipe.hset(key, mapping=mapping)
+            if entry.ttl > 0:
+                pipe.expire(key, entry.ttl)
+            await pipe.execute()
+
+    async def touch_hit(self, entry_id: str) -> None:
+        """Record that an entry served a request. Never resurrects an expired key."""
+        await self._redis.hincrby(self.key(entry_id), "hits", 1)
+
+    async def get(self, entry_id: str) -> CacheEntry | None:
+        raw = await self._redis.hgetall(self.key(entry_id))
+        if not raw:
+            return None
+        return CacheEntry.from_redis_mapping(raw)
+
+    async def delete(self, entry_id: str) -> int:
+        return int(await self._redis.delete(self.key(entry_id)))
+
+    # ------------------------------------------------------------------ queries
+    async def search(
+        self, namespace: str, vector: np.ndarray, k: int
+    ) -> list[tuple[CacheEntry, float]]:
+        query = VectorQuery(
+            vector=np.asarray(vector, dtype=np.float32).tolist(),
+            vector_field_name="embedding",
+            num_results=k,
+            filter_expression=Tag("namespace") == namespace,
+            return_fields=RETURN_FIELDS,
+            return_score=True,
+        )
+        rows = await self.index.query(query)
+        out: list[tuple[CacheEntry, float]] = []
+        for row in rows:
+            distance = float(row.get("vector_distance", 1.0))
+            similarity = 1.0 - distance
+            out.append((CacheEntry.from_redis_mapping(row), similarity))
+        out.sort(key=lambda pair: pair[1], reverse=True)
+        return out
+
+    async def count(self) -> int:
+        try:
+            info = await self.index.info()
+        except RedisError:  # pragma: no cover - index may not exist yet
+            return 0
+        return int(info.get("num_docs", 0) or 0)
+
+    # ------------------------------------------------------------- invalidation
+    async def scan_keys(self, match: str | None = None) -> list[str]:
+        pattern = match or f"{self._prefix}*"
+        keys: list[str] = []
+        async for key in self._redis.scan_iter(match=pattern, count=500):
+            keys.append(key.decode() if isinstance(key, bytes) else str(key))
+        return keys
+
+    async def invalidate(
+        self, *, namespace: str | None = None, model: str | None = None, drop_all: bool = False
+    ) -> int:
+        """Delete entries by namespace, by model, or everything.
+
+        Namespace and model are indexed tags, so this is a query, not a scan of
+        the whole keyspace. That is what makes "the system prompt changed, drop
+        its cache" a one-line operation.
+        """
+        if drop_all:
+            keys = await self.scan_keys()
+            keys += await self._scan_prefix(self._settings.exact_prefix)
+            if not keys:
+                return 0
+            return int(await self._redis.delete(*keys))
+
+        if namespace:
+            expr = Tag("namespace") == namespace
+        elif model:
+            expr = Tag("model") == model
+        else:
+            return 0
+
+        # A filter-only query, not a nearest-neighbour search. Namespace and
+        # model are indexed tags, so "drop everything for this system prompt"
+        # or "this model was upgraded" is an index lookup, never a keyspace
+        # scan. RedisVL escapes tag values, which matters because model ids are
+        # full of dots, slashes and colons.
+        to_delete: list[str] = []
+        query = FilterQuery(
+            filter_expression=expr,
+            return_fields=["entry_id", "namespace", "exact_hash"],
+            num_results=500,
+        )
+        async for page in self.index.paginate(query, page_size=500):
+            for row in page:
+                entry_id = CacheEntry._s(row, "entry_id")
+                ns = CacheEntry._s(row, "namespace")
+                exact = CacheEntry._s(row, "exact_hash")
+                if entry_id:
+                    to_delete.append(self.key(entry_id))
+                if ns and exact:
+                    to_delete.append(f"{self._settings.exact_prefix}{ns}:{exact}")
+
+        if not to_delete:
+            return 0
+        return int(await self._redis.delete(*to_delete))
+
+    async def _scan_prefix(self, prefix: str) -> list[str]:
+        keys: list[str] = []
+        async for key in self._redis.scan_iter(match=f"{prefix}*", count=500):
+            keys.append(key.decode() if isinstance(key, bytes) else str(key))
+        return keys

cachellm/cli.py

@@ -0,0 +1,122 @@
+"""Command line interface."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from cachellm import __version__
+from cachellm.settings import get_settings
+
+app = typer.Typer(
+    name="cachellm",
+    help="A drop-in semantic cache for OpenAI-compatible LLM APIs.",
+    no_args_is_help=True,
+    add_completion=False,
+)
+
+
+@app.command()
+def version() -> None:
+    """Print the installed version."""
+    typer.echo(__version__)
+
+
+@app.command()
+def serve(
+    host: Annotated[str, typer.Option(help="Bind address.")] = "",
+    port: Annotated[int, typer.Option(help="Bind port.")] = 0,
+    reload: Annotated[bool, typer.Option(help="Auto-reload on code changes.")] = False,
+) -> None:
+    """Run the proxy."""
+    import uvicorn
+
+    settings = get_settings()
+    uvicorn.run(
+        "cachellm.api.app:create_app",
+        factory=True,
+        host=host or settings.host,
+        port=port or settings.port,
+        reload=reload,
+        log_config=None,
+    )
+
+
+@app.command()
+def config() -> None:
+    """Show the effective configuration."""
+    settings = get_settings()
+    data = settings.model_dump()
+    data["api_keys"] = f"<{len(settings.client_keys)} key(s) configured>"
+    data["openai_api_key"] = "<set>" if settings.openai_api_key else "<unset>"
+    typer.echo(json.dumps(data, indent=2, default=str))
+
+
+@app.command()
+def stats() -> None:
+    """Print cache statistics straight from Redis."""
+
+    async def run() -> None:
+        from cachellm.api.deps import build_state, shutdown_state
+
+        settings = get_settings()
+        state = await build_state(settings)
+        try:
+            if state.cache is None:
+                typer.secho(f"cache unavailable: {state.degraded_reason}", fg="red")
+                raise typer.Exit(1)
+            typer.echo(json.dumps(await state.cache.stats(), indent=2))
+        finally:
+            await shutdown_state(state)
+
+    asyncio.run(run())
+
+
+@app.command()
+def invalidate(
+    namespace: Annotated[str, typer.Option(help="Namespace hash to drop.")] = "",
+    model: Annotated[str, typer.Option(help="Drop every entry for this model.")] = "",
+    all_entries: Annotated[bool, typer.Option("--all", help="Drop the whole cache.")] = False,
+) -> None:
+    """Remove cache entries by namespace, model, or all of them."""
+    if not (namespace or model or all_entries):
+        typer.secho("pass --namespace, --model or --all", fg="red")
+        raise typer.Exit(2)
+
+    async def run() -> None:
+        from cachellm.api.deps import build_state, shutdown_state
+
+        state = await build_state(get_settings())
+        try:
+            if state.cache is None:
+                typer.secho(f"cache unavailable: {state.degraded_reason}", fg="red")
+                raise typer.Exit(1)
+            removed = await state.cache.invalidate(
+                namespace=namespace or None, model=model or None, drop_all=all_entries
+            )
+            typer.echo(f"removed {removed} keys")
+        finally:
+            await shutdown_state(state)
+
+    asyncio.run(run())
+
+
+@app.command()
+def tune(
+    pairs_file: Annotated[Path, typer.Argument(help="JSONL of {a, b, duplicate} objects.")],
+    output: Annotated[Path, typer.Option(help="Write the sweep to this JSON file.")] = Path(
+        "threshold_sweep.json"
+    ),
+) -> None:
+    """Sweep similarity thresholds against labelled prompt pairs."""
+    from bench.tune_threshold import run_sweep
+
+    asyncio.run(run_sweep(pairs_file, output))
+
+
+if __name__ == "__main__":  # pragma: no cover
+    app()

cachellm/embeddings/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from cachellm.embeddings.base import Embedder
+from cachellm.embeddings.hash_backend import HashEmbedder
+from cachellm.settings import Settings
+
+__all__ = ["Embedder", "HashEmbedder", "build_embedder"]
+
+
+def build_embedder(settings: Settings) -> Embedder:
+    if settings.embedding_backend == "hash":
+        return HashEmbedder(dim=settings.embedding_dim)
+    from cachellm.embeddings.fastembed_backend import FastEmbedEmbedder
+
+    return FastEmbedEmbedder(
+        model_name=settings.embedding_model,
+        dim=settings.expected_dim(),
+        cache_size=settings.embedding_cache_size,
+    )