cachellm-proxy 0.1.0__py3-none-any.whl

cachellm/embeddings/base.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import abc
+
+import numpy as np
+
+
+class Embedder(abc.ABC):
+    """Turns text into a unit-length float32 vector.
+
+    Vectors are L2-normalised on the way out, which means the cosine distance
+    Redis reports can be turned into a similarity with ``1 - distance`` and no
+    extra maths at query time.
+    """
+
+    name: str
+    dim: int
+
+    @abc.abstractmethod
+    async def embed(self, text: str) -> np.ndarray: ...
+
+    @abc.abstractmethod
+    async def embed_batch(self, texts: list[str]) -> list[np.ndarray]: ...
+
+    async def warmup(self) -> None:
+        await self.embed("warmup")
+
+    @staticmethod
+    def normalise(vector: np.ndarray) -> np.ndarray:
+        vec = np.asarray(vector, dtype=np.float32)
+        norm = float(np.linalg.norm(vec))
+        if norm == 0.0:
+            return vec
+        return (vec / norm).astype(np.float32)
+
+    @staticmethod
+    def to_bytes(vector: np.ndarray) -> bytes:
+        return np.asarray(vector, dtype=np.float32).tobytes()

cachellm/embeddings/fastembed_backend.py

@@ -0,0 +1,75 @@
+"""Local ONNX embeddings via fastembed.
+
+Running the embedding model *inside* the proxy is the single most important
+latency decision in the project. A hosted embedding API costs 80-200 ms from
+India, which would become the floor for every cache hit and destroy the whole
+point. bge-small on CPU is single-digit milliseconds and costs nothing.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections import OrderedDict
+from typing import Any
+
+import numpy as np
+
+from cachellm.embeddings.base import Embedder
+
+
+class FastEmbedEmbedder(Embedder):
+    def __init__(self, model_name: str, dim: int, cache_size: int = 2048) -> None:
+        self.name = model_name
+        self.dim = dim
+        self._cache_size = cache_size
+        self._cache: OrderedDict[str, np.ndarray] = OrderedDict()
+        self._model: Any = None
+        self._lock = asyncio.Lock()
+
+    async def _ensure_model(self) -> Any:
+        """Load the ONNX model once, off the event loop.
+
+        Loading takes a second or two and is CPU-bound, so it goes to a worker
+        thread; the lock stops a burst of concurrent first requests from each
+        loading their own copy.
+        """
+        if self._model is None:
+            async with self._lock:
+                if self._model is None:
+                    from fastembed import TextEmbedding
+
+                    def load() -> Any:
+                        return TextEmbedding(model_name=self.name)
+
+                    self._model = await asyncio.to_thread(load)
+        return self._model
+
+    def _cache_get(self, text: str) -> np.ndarray | None:
+        vec = self._cache.get(text)
+        if vec is not None:
+            self._cache.move_to_end(text)
+        return vec
+
+    def _cache_put(self, text: str, vector: np.ndarray) -> None:
+        self._cache[text] = vector
+        self._cache.move_to_end(text)
+        while len(self._cache) > self._cache_size:
+            self._cache.popitem(last=False)
+
+    async def embed(self, text: str) -> np.ndarray:
+        cached = self._cache_get(text)
+        if cached is not None:
+            return cached
+        vectors = await self.embed_batch([text])
+        return vectors[0]
+
+    async def embed_batch(self, texts: list[str]) -> list[np.ndarray]:
+        if not texts:
+            return []
+        model = await self._ensure_model()
+        pending = [t for t in texts if self._cache_get(t) is None]
+        if pending:
+            raw = await asyncio.to_thread(lambda: list(model.embed(pending)))
+            for text, vector in zip(pending, raw, strict=True):
+                self._cache_put(text, self.normalise(np.asarray(vector, dtype=np.float32)))
+        return [self._cache[t] for t in texts]

cachellm/embeddings/hash_backend.py

@@ -0,0 +1,42 @@
+"""Deterministic hashing embedder.
+
+No model download, no network, identical output on every machine, which makes
+it the right backend for unit tests and CI. It uses the classic hashing trick
+(signed token buckets), so it still produces meaningful lexical similarity:
+shared words pull vectors together. It does NOT understand meaning, so it is
+not a substitute for a real model in the evaluation numbers.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+
+import numpy as np
+
+from cachellm.embeddings.base import Embedder
+
+_TOKEN = re.compile(r"[a-z0-9']+")
+
+
+class HashEmbedder(Embedder):
+    def __init__(self, dim: int = 384) -> None:
+        self.name = "hash-embedder"
+        self.dim = dim
+
+    def _vector(self, text: str) -> np.ndarray:
+        vec = np.zeros(self.dim, dtype=np.float32)
+        tokens = _TOKEN.findall(text.lower())
+        for token in tokens:
+            digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
+            value = int.from_bytes(digest, "big")
+            bucket = value % self.dim
+            sign = 1.0 if (value >> 63) & 1 else -1.0
+            vec[bucket] += sign
+        return self.normalise(vec)
+
+    async def embed(self, text: str) -> np.ndarray:
+        return self._vector(text)
+
+    async def embed_batch(self, texts: list[str]) -> list[np.ndarray]:
+        return [self._vector(t) for t in texts]

cachellm/errors.py

@@ -0,0 +1,72 @@
+"""OpenAI-shaped error envelope.
+
+A drop-in proxy has to fail the way the thing it replaces fails, otherwise
+client SDK error handling breaks. The official clients read
+``{"error": {"message", "type", "param", "code"}}``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import HTTPException
+
+
+class CacheLLMError(HTTPException):
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        err_type: str = "invalid_request_error",
+        code: str | None = None,
+        param: str | None = None,
+    ) -> None:
+        super().__init__(status_code=status_code, detail=message)
+        self.message = message
+        self.err_type = err_type
+        self.code = code
+        self.param = param
+
+    def envelope(self) -> dict[str, Any]:
+        return {
+            "error": {
+                "message": self.message,
+                "type": self.err_type,
+                "param": self.param,
+                "code": self.code,
+            }
+        }
+
+
+class AuthError(CacheLLMError):
+    def __init__(self, message: str = "Invalid API key provided.") -> None:
+        super().__init__(401, message, "invalid_request_error", code="invalid_api_key")
+
+
+class UpstreamError(CacheLLMError):
+    def __init__(self, message: str, status_code: int = 502) -> None:
+        super().__init__(status_code, message, "api_error", code="upstream_error")
+
+
+class ModelNotFoundError(CacheLLMError):
+    def __init__(self, model: str) -> None:
+        super().__init__(
+            404,
+            f"The model `{model}` does not exist or you do not have access to it.",
+            "invalid_request_error",
+            code="model_not_found",
+            param="model",
+        )
+
+
+class CacheMissError(CacheLLMError):
+    """Raised when a client sent ``X-Cache-Control: only-if-cached`` and missed."""
+
+    def __init__(self) -> None:
+        super().__init__(
+            504,
+            "No cached response above the similarity threshold and the request "
+            "was sent with X-Cache-Control: only-if-cached.",
+            "api_error",
+            code="cache_miss",
+        )

cachellm/logging_setup.py

@@ -0,0 +1,56 @@
+"""structlog configuration.
+
+Prompt and completion text is treated as sensitive by default. A proxy sees
+every question every user asks, so logging it wholesale is a privacy incident
+waiting to happen. Set ``CACHELLM_LOG_PROMPTS=true`` only in local development.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from typing import Any
+
+import structlog
+
+from cachellm.settings import Settings
+
+SENSITIVE_KEYS = {"prompt", "matched_prompt", "response_text", "messages", "content", "answer"}
+
+
+def _redact_prompts(_logger: Any, _name: str, event_dict: dict[str, Any]) -> dict[str, Any]:
+    for key in list(event_dict):
+        if key in SENSITIVE_KEYS:
+            value = event_dict[key]
+            if isinstance(value, str):
+                event_dict[key] = f"<redacted {len(value)} chars>"
+            else:
+                event_dict[key] = "<redacted>"
+    return event_dict
+
+
+def configure_logging(settings: Settings) -> None:
+    level = getattr(logging, settings.log_level.upper(), logging.INFO)
+    logging.basicConfig(format="%(message)s", stream=sys.stdout, level=level)
+
+    processors: list[Any] = [
+        structlog.contextvars.merge_contextvars,
+        structlog.processors.add_log_level,
+        structlog.processors.TimeStamper(fmt="iso", utc=True),
+    ]
+    if not settings.log_prompts:
+        processors.append(_redact_prompts)
+    processors.append(structlog.processors.StackInfoRenderer())
+    processors.append(structlog.processors.format_exc_info)
+    processors.append(
+        structlog.processors.JSONRenderer()
+        if settings.log_json
+        else structlog.dev.ConsoleRenderer(colors=sys.stdout.isatty())
+    )
+
+    structlog.configure(
+        processors=processors,
+        wrapper_class=structlog.make_filtering_bound_logger(level),
+        logger_factory=structlog.PrintLoggerFactory(),
+        cache_logger_on_first_use=True,
+    )

cachellm/models.py

@@ -0,0 +1,181 @@
+"""Request and response schemas mirroring the OpenAI chat completions contract.
+
+Unknown fields are preserved rather than rejected: OpenAI adds parameters
+regularly and a proxy that 422s on an unrecognised key is not a drop-in.
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+Role = Literal["system", "user", "assistant", "tool", "developer", "function"]
+
+
+class ChatMessage(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    role: Role
+    content: str | list[dict[str, Any]] | None = None
+    name: str | None = None
+    tool_calls: list[dict[str, Any]] | None = None
+    tool_call_id: str | None = None
+
+    def as_text(self) -> str:
+        """Flatten multi-part content down to plain text for hashing/embedding."""
+        if self.content is None:
+            return ""
+        if isinstance(self.content, str):
+            return self.content
+        parts: list[str] = []
+        for part in self.content:
+            if not isinstance(part, dict):
+                continue
+            if part.get("type") in (None, "text") and isinstance(part.get("text"), str):
+                parts.append(part["text"])
+        return "\n".join(parts)
+
+    def has_non_text_parts(self) -> bool:
+        if isinstance(self.content, list):
+            return any(
+                isinstance(p, dict) and p.get("type") not in (None, "text") for p in self.content
+            )
+        return False
+
+
+class ChatCompletionRequest(BaseModel):
+    model_config = ConfigDict(extra="allow", protected_namespaces=())
+
+    model: str
+    messages: list[ChatMessage]
+    temperature: float | None = None
+    top_p: float | None = None
+    n: int | None = None
+    stream: bool | None = None
+    stream_options: dict[str, Any] | None = None
+    stop: str | list[str] | None = None
+    max_tokens: int | None = None
+    max_completion_tokens: int | None = None
+    presence_penalty: float | None = None
+    frequency_penalty: float | None = None
+    seed: int | None = None
+    user: str | None = None
+    response_format: dict[str, Any] | None = None
+    tools: list[dict[str, Any]] | None = None
+    tool_choice: str | dict[str, Any] | None = None
+
+    # ------------------------------------------------------------------ helpers
+    @property
+    def effective_temperature(self) -> float:
+        return 1.0 if self.temperature is None else float(self.temperature)
+
+    @property
+    def effective_max_tokens(self) -> int | None:
+        return self.max_completion_tokens or self.max_tokens
+
+    def system_text(self) -> str:
+        return "\n".join(
+            m.as_text() for m in self.messages if m.role in ("system", "developer")
+        ).strip()
+
+    def conversation_messages(self) -> list[ChatMessage]:
+        return [m for m in self.messages if m.role not in ("system", "developer")]
+
+    def last_user_text(self) -> str:
+        for message in reversed(self.messages):
+            if message.role == "user":
+                return message.as_text().strip()
+        return ""
+
+    def user_turn_count(self) -> int:
+        return sum(1 for m in self.messages if m.role == "user")
+
+    def cache_text(self, include_history: bool) -> str:
+        """The text that gets embedded and hashed."""
+        if not include_history:
+            return self.last_user_text()
+        turns = [f"{m.role}: {m.as_text().strip()}" for m in self.conversation_messages()]
+        return "\n".join(t for t in turns if t.strip())
+
+
+class Usage(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+class ChoiceMessage(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    role: Literal["assistant"] = "assistant"
+    content: str | None = None
+    tool_calls: list[dict[str, Any]] | None = None
+
+
+class Choice(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    index: int = 0
+    message: ChoiceMessage
+    finish_reason: str | None = "stop"
+    logprobs: None = None
+
+
+class ChatCompletionResponse(BaseModel):
+    model_config = ConfigDict(extra="allow", protected_namespaces=())
+
+    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:24]}")
+    object: Literal["chat.completion"] = "chat.completion"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    choices: list[Choice]
+    usage: Usage = Field(default_factory=Usage)
+    system_fingerprint: str | None = None
+
+    @classmethod
+    def from_text(
+        cls,
+        *,
+        model: str,
+        text: str,
+        prompt_tokens: int = 0,
+        completion_tokens: int = 0,
+        finish_reason: str = "stop",
+    ) -> ChatCompletionResponse:
+        return cls(
+            model=model,
+            choices=[
+                Choice(
+                    index=0,
+                    message=ChoiceMessage(role="assistant", content=text),
+                    finish_reason=finish_reason,
+                )
+            ],
+            usage=Usage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
+            ),
+        )
+
+    def text(self) -> str:
+        if not self.choices:
+            return ""
+        return self.choices[0].message.content or ""
+
+
+class ModelCard(BaseModel):
+    id: str
+    object: Literal["model"] = "model"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    owned_by: str = "cachellm"
+
+
+class ModelList(BaseModel):
+    object: Literal["list"] = "list"
+    data: list[ModelCard]

cachellm/observability/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from cachellm.observability.metrics import Metrics, get_metrics, reset_metrics
+from cachellm.observability.tracing import setup_tracing, span
+
+__all__ = ["Metrics", "get_metrics", "reset_metrics", "setup_tracing", "span"]

cachellm/observability/metrics.py

@@ -0,0 +1,147 @@
+"""Prometheus metrics.
+
+The metric that matters most is ``cachellm_cost_usd_total{kind="saved"}``. It is
+computed once, at serve time, from the token counts stored on the cache entry
+multiplied by that model's price. Everything downstream (the Grafana panel, the
+admin endpoint, the README headline) reads the same number, so they can never
+disagree with each other.
+
+Label cardinality is kept deliberately low: model and provider are bounded
+sets, category has six values, result has five. Prompt text never becomes a
+label.
+"""
+
+from __future__ import annotations
+
+from prometheus_client import (
+    CONTENT_TYPE_LATEST,
+    CollectorRegistry,
+    Counter,
+    Gauge,
+    Histogram,
+    generate_latest,
+)
+
+# Buckets tuned for this workload: cache hits land in single-digit milliseconds,
+# provider calls in hundreds of milliseconds to seconds.
+LATENCY_BUCKETS = (0.001, 0.002, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0)
+SIMILARITY_BUCKETS = (0.5, 0.6, 0.7, 0.75, 0.8, 0.85, 0.88, 0.9, 0.92, 0.94, 0.96, 0.98, 0.99, 1.0)
+
+
+class Metrics:
+    def __init__(self, registry: CollectorRegistry | None = None) -> None:
+        self.registry = registry or CollectorRegistry()
+
+        self.requests = Counter(
+            "cachellm_requests_total",
+            "Chat completion requests handled, by cache outcome.",
+            ["result", "category", "provider", "model"],
+            registry=self.registry,
+        )
+        self.request_duration = Histogram(
+            "cachellm_request_duration_seconds",
+            "End-to-end request duration.",
+            ["result"],
+            buckets=LATENCY_BUCKETS,
+            registry=self.registry,
+        )
+        self.lookup_duration = Histogram(
+            "cachellm_lookup_duration_seconds",
+            "Time spent deciding hit or miss (policy, embedding, vector search).",
+            ["tier"],
+            buckets=LATENCY_BUCKETS,
+            registry=self.registry,
+        )
+        self.embed_duration = Histogram(
+            "cachellm_embed_duration_seconds",
+            "Time spent embedding the prompt.",
+            buckets=LATENCY_BUCKETS,
+            registry=self.registry,
+        )
+        self.similarity = Histogram(
+            "cachellm_similarity",
+            "Best similarity score seen at lookup time.",
+            ["result"],
+            buckets=SIMILARITY_BUCKETS,
+            registry=self.registry,
+        )
+        self.tokens = Counter(
+            "cachellm_tokens_total",
+            "Tokens, split by whether they were spent upstream or avoided.",
+            ["kind", "direction", "model"],
+            registry=self.registry,
+        )
+        self.cost = Counter(
+            "cachellm_cost_usd_total",
+            "Modelled USD, split into spent upstream and saved by the cache.",
+            ["kind", "model"],
+            registry=self.registry,
+        )
+        self.entries = Gauge(
+            "cachellm_cache_entries",
+            "Documents currently in the vector index.",
+            registry=self.registry,
+        )
+        self.provider_errors = Counter(
+            "cachellm_provider_errors_total",
+            "Upstream provider failures.",
+            ["provider"],
+            registry=self.registry,
+        )
+        self.coalesced = Counter(
+            "cachellm_coalesced_total",
+            "Requests that joined an in-flight identical call instead of duplicating it.",
+            registry=self.registry,
+        )
+        self.near_misses = Counter(
+            "cachellm_near_miss_total",
+            "Lookups that landed just below the similarity threshold.",
+            ["category"],
+            registry=self.registry,
+        )
+        self.shadow_hits = Counter(
+            "cachellm_shadow_hits_total",
+            "Requests the cache would have served while running in shadow mode.",
+            ["category"],
+            registry=self.registry,
+        )
+
+    # ------------------------------------------------------------------ helpers
+    def observe_tokens(
+        self, model: str, prompt_tokens: int, completion_tokens: int, kind: str
+    ) -> None:
+        if prompt_tokens:
+            self.tokens.labels(kind=kind, direction="input", model=model).inc(prompt_tokens)
+        if completion_tokens:
+            self.tokens.labels(kind=kind, direction="output", model=model).inc(completion_tokens)
+
+    def observe_cost(self, model: str, usd: float, kind: str) -> None:
+        if usd:
+            self.cost.labels(kind=kind, model=model).inc(usd)
+
+    def render(self) -> tuple[bytes, str]:
+        """Serialise for a Prometheus scrape.
+
+        The content type has to match the body. `generate_latest` emits the
+        Prometheus text format; advertising OpenMetrics alongside it makes
+        Prometheus reject the scrape with "data does not end with # EOF",
+        because OpenMetrics requires that terminator and the text format has
+        no such thing. Nothing but a real scrape catches this.
+        """
+        return generate_latest(self.registry), CONTENT_TYPE_LATEST
+
+
+_metrics: Metrics | None = None
+
+
+def get_metrics() -> Metrics:
+    global _metrics
+    if _metrics is None:
+        _metrics = Metrics()
+    return _metrics
+
+
+def reset_metrics() -> None:
+    """Test helper: start from a clean registry."""
+    global _metrics
+    _metrics = None