cachellm-proxy 0.1.0__py3-none-any.whl

cachellm/providers/openai_compat.py

@@ -0,0 +1,131 @@
+"""Any OpenAI-compatible upstream: OpenAI itself, Groq, Together, OpenRouter,
+vLLM, LM Studio, llama.cpp server.
+
+One adapter covers all of them because they all speak the same
+``POST /chat/completions``. That is the whole reason the OpenAI shape became
+the industry contract, and it is why CacheLLM speaks it on the front door too.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+import structlog
+
+from cachellm.errors import UpstreamError
+from cachellm.models import ChatCompletionRequest
+from cachellm.providers.base import (
+    Provider,
+    ProviderResult,
+    StreamEvent,
+    normalise_finish_reason,
+)
+from cachellm.settings import Settings
+
+log = structlog.get_logger(__name__)
+
+
+class OpenAICompatProvider(Provider):
+    name = "openai"
+
+    def __init__(
+        self, settings: Settings, base_url: str | None = None, api_key: str | None = None
+    ) -> None:
+        self._settings = settings
+        self._base_url = (base_url or settings.openai_base_url).rstrip("/")
+        self._api_key = api_key if api_key is not None else settings.openai_api_key
+        self._client: httpx.AsyncClient | None = None
+
+    def _http(self) -> httpx.AsyncClient:
+        if self._client is None:
+            headers = {"Content-Type": "application/json"}
+            if self._api_key:
+                headers["Authorization"] = f"Bearer {self._api_key}"
+            self._client = httpx.AsyncClient(
+                base_url=self._base_url,
+                headers=headers,
+                timeout=httpx.Timeout(self._settings.request_timeout, connect=10.0),
+            )
+        return self._client
+
+    def _payload(self, request: ChatCompletionRequest, stream: bool) -> dict[str, Any]:
+        body = request.model_dump(exclude_none=True, by_alias=True)
+        body["model"] = self.resolve_model(request.model)
+        body["stream"] = stream
+        body.pop("stream_options", None)
+        if stream:
+            body["stream_options"] = {"include_usage": True}
+        return body
+
+    async def complete(self, request: ChatCompletionRequest) -> ProviderResult:
+        try:
+            response = await self._http().post(
+                "/chat/completions", json=self._payload(request, False)
+            )
+            response.raise_for_status()
+        except httpx.HTTPStatusError as exc:
+            detail = exc.response.text[:400]
+            log.warning("openai_http_error", status=exc.response.status_code, detail=detail)
+            raise UpstreamError(
+                f"Upstream returned {exc.response.status_code}: {detail}",
+                status_code=exc.response.status_code,
+            ) from exc
+        except httpx.HTTPError as exc:
+            raise UpstreamError(f"Upstream request failed: {exc}") from exc
+
+        data = response.json()
+        choice = (data.get("choices") or [{}])[0]
+        message = choice.get("message", {}) or {}
+        usage = data.get("usage", {}) or {}
+        return ProviderResult(
+            text=message.get("content") or "",
+            model=data.get("model", self.resolve_model(request.model)),
+            prompt_tokens=int(usage.get("prompt_tokens", 0) or 0),
+            completion_tokens=int(usage.get("completion_tokens", 0) or 0),
+            finish_reason=normalise_finish_reason(choice.get("finish_reason")),
+            tool_calls=message.get("tool_calls"),
+        )
+
+    async def stream(self, request: ChatCompletionRequest) -> AsyncIterator[StreamEvent]:
+        payload = self._payload(request, True)
+        try:
+            async with self._http().stream("POST", "/chat/completions", json=payload) as response:
+                if response.status_code >= 400:
+                    body = (await response.aread()).decode("utf-8", "replace")[:400]
+                    raise UpstreamError(
+                        f"Upstream returned {response.status_code}: {body}",
+                        status_code=response.status_code,
+                    )
+                async for line in response.aiter_lines():
+                    if not line.startswith("data:"):
+                        continue
+                    data = line[5:].strip()
+                    if data == "[DONE]":
+                        break
+                    try:
+                        chunk = json.loads(data)
+                    except json.JSONDecodeError:
+                        continue
+                    usage = chunk.get("usage") or {}
+                    if usage:
+                        yield StreamEvent(
+                            prompt_tokens=int(usage.get("prompt_tokens", 0) or 0),
+                            completion_tokens=int(usage.get("completion_tokens", 0) or 0),
+                        )
+                    for choice in chunk.get("choices", []) or []:
+                        delta = (choice.get("delta") or {}).get("content") or ""
+                        finish = choice.get("finish_reason")
+                        if delta:
+                            yield StreamEvent(delta=delta)
+                        if finish:
+                            yield StreamEvent(finish_reason=normalise_finish_reason(finish))
+        except httpx.HTTPError as exc:
+            raise UpstreamError(f"Upstream stream failed: {exc}") from exc
+
+    async def close(self) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None

cachellm/providers/registry.py

@@ -0,0 +1,96 @@
+"""Model to provider routing.
+
+Two ways to reach a provider:
+
+* explicit prefix, ``bedrock/us.amazon.nova-lite-v1:0`` or ``openai/gpt-4o-mini``,
+  which is unambiguous and is what the docs recommend;
+* bare model id, which is what a drop-in client actually sends, routed by
+  recognising vendor prefixes and otherwise falling back to the configured
+  default provider.
+
+The provider name is part of the cache namespace, so a Bedrock answer can never
+be served to an OpenAI request even when the model strings look alike.
+"""
+
+from __future__ import annotations
+
+import structlog
+
+from cachellm.providers.base import Provider
+from cachellm.providers.bedrock import BedrockProvider, looks_like_bedrock_model
+from cachellm.providers.fake import FakeProvider
+from cachellm.providers.openai_compat import OpenAICompatProvider
+from cachellm.settings import Settings
+
+log = structlog.get_logger(__name__)
+
+# Model families we can attribute to a vendor with confidence. A bare
+# "gpt-4o-mini" from a drop-in client means OpenAI, whatever the configured
+# default provider happens to be.
+OPENAI_MODEL_PREFIXES = (
+    "gpt-",
+    "o1",
+    "o3",
+    "o4",
+    "chatgpt-",
+    "text-embedding-",
+    "davinci",
+    "babbage",
+)
+
+# Shown by GET /v1/models so a client can discover what this proxy will route.
+SUGGESTED_MODELS = [
+    "bedrock/us.amazon.nova-micro-v1:0",
+    "bedrock/us.amazon.nova-lite-v1:0",
+    "bedrock/us.amazon.nova-pro-v1:0",
+    "bedrock/anthropic.claude-3-haiku-20240307-v1:0",
+    "openai/gpt-4o-mini",
+    "fake/echo",
+]
+
+
+class ProviderRegistry:
+    def __init__(self, settings: Settings, providers: dict[str, Provider] | None = None) -> None:
+        self._settings = settings
+        self._providers: dict[str, Provider] = providers or {}
+
+    def _get_or_create(self, name: str) -> Provider:
+        if name not in self._providers:
+            if name == "bedrock":
+                self._providers[name] = BedrockProvider(self._settings)
+            elif name == "openai":
+                self._providers[name] = OpenAICompatProvider(self._settings)
+            elif name == "fake":
+                self._providers[name] = FakeProvider(latency_ms=self._settings.fake_latency_ms)
+            else:  # pragma: no cover - guarded by resolve()
+                raise ValueError(f"unknown provider {name!r}")
+        return self._providers[name]
+
+    def provider_name_for(self, model: str) -> str:
+        model = model.strip()
+        if "/" in model:
+            prefix = model.split("/", 1)[0].lower()
+            if prefix in ("bedrock", "openai", "fake"):
+                return prefix
+        if model.startswith("fake-"):
+            return "fake"
+        if looks_like_bedrock_model(model):
+            return "bedrock"
+        if model.lower().startswith(OPENAI_MODEL_PREFIXES):
+            return "openai"
+        return self._settings.default_provider
+
+    def resolve(self, model: str) -> tuple[Provider, str]:
+        name = self.provider_name_for(model)
+        return self._get_or_create(name), name
+
+    def register(self, name: str, provider: Provider) -> None:
+        self._providers[name] = provider
+
+    def models(self) -> list[str]:
+        return list(SUGGESTED_MODELS)
+
+    async def close(self) -> None:
+        for provider in self._providers.values():
+            await provider.close()
+        self._providers.clear()

cachellm/settings.py

@@ -0,0 +1,230 @@
+"""Typed configuration for CacheLLM.
+
+Every value can be set through an environment variable prefixed with
+``CACHELLM_`` (for example ``CACHELLM_THRESHOLD_DEFAULT=0.94``) or through a
+``.env`` file sitting next to the process working directory. Defaults are
+deliberately conservative: it is far worse to serve a confidently wrong cached
+answer than to miss the cache and pay for a fresh call.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Literal
+
+from pydantic import field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+# Embedding dimensions for the models we ship presets for. Used to catch the
+# classic "changed the model but not the index" mistake at startup.
+KNOWN_EMBEDDING_DIMS: dict[str, int] = {
+    "BAAI/bge-small-en-v1.5": 384,
+    "BAAI/bge-base-en-v1.5": 768,
+    "sentence-transformers/all-MiniLM-L6-v2": 384,
+    "snowflake/snowflake-arctic-embed-s": 384,
+    "jinaai/jina-embeddings-v2-small-en": 512,
+    "thenlper/gte-base": 768,
+    "hash-embedder": 384,
+}
+
+# Similarity thresholds DO NOT transfer between embedding models. Measured on
+# the project's own labelled corpus (see `bench/compare_models.py` and
+# docs/evaluation.md), these are the highest thresholds at which each model
+# served zero hard negatives: pairs one word apart with opposite meaning, like
+# "undo the last git commit" versus "undo the last git merge".
+#
+# The spread is the point. bge-small needs 0.96 where MiniLM needs 0.89, so a
+# threshold copied from a blog post is meaningless without naming the model it
+# was measured on.
+CALIBRATED_THRESHOLDS: dict[str, float] = {
+    "sentence-transformers/all-MiniLM-L6-v2": 0.89,
+    "BAAI/bge-small-en-v1.5": 0.96,
+    "BAAI/bge-base-en-v1.5": 0.94,
+    "snowflake/snowflake-arctic-embed-s": 0.98,
+    "jinaai/jina-embeddings-v2-small-en": 0.96,
+    "thenlper/gte-base": 0.96,
+}
+
+Category = Literal[
+    "factual",
+    "classification",
+    "creative",
+    "volatile",
+    "conversational",
+    "default",
+]
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        env_prefix="CACHELLM_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    # ------------------------------------------------------------------ server
+    host: str = "127.0.0.1"
+    port: int = 8080
+    log_level: str = "INFO"
+    log_json: bool = True
+    #: Prompt and completion text is NEVER written to logs unless this is on.
+    log_prompts: bool = False
+
+    # -------------------------------------------------------------------- auth
+    #: Comma-separated client keys. Empty means auth is disabled (local dev only).
+    api_keys: str = ""
+    require_auth_for_admin: bool = True
+
+    # ------------------------------------------------------------------- redis
+    redis_url: str = "redis://localhost:6379/0"
+    index_name: str = "cachellm_idx"
+    entry_prefix: str = "cachellm:e:"
+    exact_prefix: str = "cachellm:x:"
+    stats_prefix: str = "cachellm:s:"
+
+    # -------------------------------------------------------------- embeddings
+    embedding_backend: Literal["fastembed", "hash"] = "fastembed"
+    #: MiniLM is the default because it gave four times the safe recall of
+    #: bge-small on the evaluation corpus, at a third of the download size.
+    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
+    embedding_dim: int = 384
+    embedding_cache_size: int = 2048
+
+    # ---------------------------------------------------------------- caching
+    enabled: bool = True
+    #: Observe-only. Looks up, records what it *would* have served, then still
+    #: calls the provider. The safe way to roll out onto real traffic.
+    shadow_mode: bool = False
+    cache_streaming: bool = True
+    #: Store responses that stopped because they hit max_tokens. Off by default:
+    #: a truncated answer served from cache forever is a silent quality bug. Turn
+    #: it on if you cap max_tokens deliberately and want those answers cached.
+    cache_truncated: bool = False
+    top_k: int = 3
+    #: A miss within this margin of the threshold is recorded as a near miss.
+    near_miss_margin: float = 0.06
+    near_miss_log_size: int = 500
+
+    #: 0.0 means "use the calibrated value for the configured embedding model".
+    #: Set any of these explicitly to override the calibration.
+    threshold_default: float = 0.0
+    threshold_factual: float = 0.0
+    threshold_classification: float = 0.0
+    threshold_creative: float = 0.0
+    threshold_volatile: float = 0.0
+    threshold_conversational: float = 0.0
+
+    #: Offsets applied to the calibrated threshold, per category. Classification
+    #: has a small, constrained answer space so it tolerates looser matching;
+    #: creative generation tolerates almost none.
+    category_threshold_offsets: dict[str, float] = {
+        "classification": -0.02,
+        "factual": 0.0,
+        "conversational": +0.03,
+        "volatile": +0.03,
+        "creative": +0.06,
+        "default": 0.0,
+    }
+
+    ttl_default: int = 86_400  # 1 day
+    ttl_factual: int = 604_800  # 7 days
+    ttl_classification: int = 604_800
+    ttl_creative: int = 3_600
+    ttl_volatile: int = 900  # 15 minutes
+    ttl_conversational: int = 3_600
+
+    # ------------------------------------------------------- cacheability rules
+    max_cacheable_temperature: float = 0.3
+    max_prompt_chars: int = 8_000
+    cache_multi_turn: bool = False
+    cache_json_mode: bool = False
+    cache_tool_calls: bool = False
+    #: Skip caching prompts that look like they carry personal data.
+    pii_guard: bool = True
+    #: Strip filler words before embedding. Raises hit rate, slightly raises risk.
+    strip_filler_words: bool = False
+
+    # --------------------------------------------------------------- providers
+    default_provider: Literal["bedrock", "openai", "fake"] = "bedrock"
+    aws_region: str = "us-east-1"
+    aws_profile: str = ""
+    #: Simulated upstream latency for the `fake` provider, so a reproducible
+    #: benchmark can show a realistic hit-versus-miss gap without spending money.
+    fake_latency_ms: float = 0.0
+    openai_base_url: str = "https://api.openai.com/v1"
+    openai_api_key: str = ""
+    request_timeout: float = 120.0
+    provider_max_retries: int = 2
+
+    # ----------------------------------------------------------- observability
+    metrics_enabled: bool = True
+    tracing_enabled: bool = False
+    otlp_endpoint: str = ""
+    service_name: str = "cachellm"
+
+    @field_validator(
+        "threshold_default",
+        "threshold_factual",
+        "threshold_classification",
+        "threshold_creative",
+        "threshold_volatile",
+        "threshold_conversational",
+    )
+    @classmethod
+    def _valid_threshold(cls, v: float) -> float:
+        if not 0.0 <= v <= 1.0:
+            raise ValueError("similarity thresholds must be between 0 and 1")
+        return v
+
+    # ------------------------------------------------------------------ helpers
+    @property
+    def client_keys(self) -> set[str]:
+        return {k.strip() for k in self.api_keys.split(",") if k.strip()}
+
+    @property
+    def auth_enabled(self) -> bool:
+        return bool(self.client_keys)
+
+    @property
+    def calibrated_threshold(self) -> float:
+        """The safe threshold for the configured embedding model.
+
+        0.92 is the fallback for an uncalibrated model. It is a guess, and the
+        proxy says so at startup, because the measured safe points across six
+        models span 0.89 to 0.98.
+        """
+        if self.embedding_backend == "hash":
+            return 0.92
+        return CALIBRATED_THRESHOLDS.get(self.embedding_model, 0.92)
+
+    @property
+    def is_calibrated(self) -> bool:
+        return self.embedding_backend != "hash" and self.embedding_model in CALIBRATED_THRESHOLDS
+
+    def threshold_for(self, category: str) -> float:
+        explicit = float(getattr(self, f"threshold_{category}", 0.0) or 0.0)
+        if explicit > 0.0:
+            return explicit
+        base = self.threshold_default if self.threshold_default > 0.0 else self.calibrated_threshold
+        offset = self.category_threshold_offsets.get(category, 0.0)
+        # Cap below 1.0: an exact-similarity requirement would make the semantic
+        # tier dead weight, and the exact tier already covers literal repeats.
+        return round(min(0.995, max(0.0, base + offset)), 4)
+
+    def ttl_for(self, category: str) -> int:
+        return int(getattr(self, f"ttl_{category}", self.ttl_default))
+
+    def expected_dim(self) -> int:
+        return KNOWN_EMBEDDING_DIMS.get(self.embedding_model, self.embedding_dim)
+
+
+@lru_cache(maxsize=1)
+def get_settings() -> Settings:
+    return Settings()
+
+
+def reset_settings_cache() -> None:
+    """Test helper: drop the memoised Settings instance."""
+    get_settings.cache_clear()