cachellm-proxy 0.1.0__py3-none-any.whl

cachellm/observability/tracing.py

@@ -0,0 +1,107 @@
+"""Optional OpenTelemetry tracing.
+
+OTel is the spine, not the destination. Instrument once and the same spans can
+go to Langfuse (which reads them as LLM generations, with cost and prompts),
+Grafana Tempo, Jaeger, or anything else that accepts OTLP. Nothing here is
+required for the proxy to run: with tracing disabled every helper degrades to a
+no-op context manager.
+
+Span attributes follow the OpenTelemetry GenAI semantic conventions where they
+exist (``gen_ai.*``), with CacheLLM's own facts under ``cachellm.*``.
+"""
+
+from __future__ import annotations
+
+import base64
+import contextlib
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import structlog
+
+from cachellm.settings import Settings
+
+log = structlog.get_logger(__name__)
+
+_tracer: Any = None
+_enabled = False
+
+
+def langfuse_otlp_headers(public_key: str, secret_key: str) -> str:
+    """Langfuse authenticates OTLP with HTTP basic auth over the key pair."""
+    token = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode()
+    return f"Authorization=Basic {token}"
+
+
+def setup_tracing(settings: Settings, app: Any = None) -> bool:
+    """Configure the global tracer. Returns True when tracing is live."""
+    global _tracer, _enabled
+    if not settings.tracing_enabled:
+        return False
+    try:
+        from opentelemetry import trace
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+        from opentelemetry.sdk.resources import Resource
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+    except ImportError:
+        log.warning(
+            "tracing_unavailable",
+            hint="install the observability extra: uv sync --extra observability",
+        )
+        return False
+
+    endpoint = settings.otlp_endpoint or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "")
+    if not endpoint:
+        log.warning("tracing_disabled_no_endpoint")
+        return False
+
+    provider = TracerProvider(resource=Resource.create({"service.name": settings.service_name}))
+    provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(endpoint=endpoint)))
+    trace.set_tracer_provider(provider)
+    _tracer = trace.get_tracer("cachellm")
+    _enabled = True
+
+    for module, instrument in (
+        ("opentelemetry.instrumentation.redis", "RedisInstrumentor"),
+        ("opentelemetry.instrumentation.httpx", "HTTPXClientInstrumentor"),
+        ("opentelemetry.instrumentation.botocore", "BotocoreInstrumentor"),
+    ):
+        with contextlib.suppress(Exception):
+            mod = __import__(module, fromlist=[instrument])
+            getattr(mod, instrument)().instrument()
+
+    if app is not None:
+        with contextlib.suppress(Exception):
+            from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+            FastAPIInstrumentor.instrument_app(app, excluded_urls="healthz,readyz,metrics")
+
+    log.info("tracing_enabled", endpoint=endpoint, service=settings.service_name)
+    return True
+
+
+@contextlib.contextmanager
+def span(name: str, **attributes: Any) -> Iterator[Any]:
+    """Start a span, or do nothing at all when tracing is off."""
+    if not _enabled or _tracer is None:
+        yield None
+        return
+    with _tracer.start_as_current_span(name) as current:
+        for key, value in attributes.items():
+            if value is not None:
+                current.set_attribute(key, value)
+        yield current
+
+
+def set_attributes(current: Any, **attributes: Any) -> None:
+    if current is None:
+        return
+    for key, value in attributes.items():
+        if value is not None:
+            current.set_attribute(key, value)
+
+
+def is_enabled() -> bool:
+    return _enabled

cachellm/pricing.py

@@ -0,0 +1,108 @@
+"""Token prices used to compute "money saved" for every cache hit.
+
+Prices are USD per one million tokens and are *defaults*: providers change them,
+so treat this table as a starting point and override it with
+``CACHELLM_PRICING_FILE`` pointing at a JSON file of the same shape.
+
+The saving reported by CacheLLM is a modelled number, not a bill. It is
+"what these tokens would have cost at list price if we had called the provider",
+which is exactly the number a team wants when deciding whether to deploy this.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+
+__all__ = ["PRICES", "ModelPrice", "estimate_cost", "price_for"]
+
+
+@dataclass(frozen=True)
+class ModelPrice:
+    """USD per 1M tokens."""
+
+    input_per_m: float
+    output_per_m: float
+
+    def cost(self, input_tokens: int, output_tokens: int) -> float:
+        return (input_tokens * self.input_per_m + output_tokens * self.output_per_m) / 1_000_000
+
+
+# Defaults verified against public list prices; override per deployment.
+PRICES: dict[str, ModelPrice] = {
+    # --- AWS Bedrock: Amazon Nova ---
+    "amazon.nova-micro-v1:0": ModelPrice(0.035, 0.14),
+    "amazon.nova-lite-v1:0": ModelPrice(0.06, 0.24),
+    "amazon.nova-pro-v1:0": ModelPrice(0.80, 3.20),
+    # --- AWS Bedrock: Anthropic ---
+    "anthropic.claude-3-haiku-20240307-v1:0": ModelPrice(0.25, 1.25),
+    "anthropic.claude-3-5-haiku-20241022-v1:0": ModelPrice(0.80, 4.00),
+    "anthropic.claude-3-sonnet-20240229-v1:0": ModelPrice(3.00, 15.00),
+    "anthropic.claude-3-5-sonnet-20240620-v1:0": ModelPrice(3.00, 15.00),
+    # --- AWS Bedrock: Meta / Mistral ---
+    "meta.llama3-8b-instruct-v1:0": ModelPrice(0.30, 0.60),
+    "meta.llama3-70b-instruct-v1:0": ModelPrice(2.65, 3.50),
+    "mistral.mistral-7b-instruct-v0:2": ModelPrice(0.15, 0.20),
+    "mistral.mixtral-8x7b-instruct-v0:1": ModelPrice(0.45, 0.70),
+    "mistral.mistral-large-2402-v1:0": ModelPrice(4.00, 12.00),
+    # --- OpenAI ---
+    "gpt-4o-mini": ModelPrice(0.15, 0.60),
+    "gpt-4o": ModelPrice(2.50, 10.00),
+    "gpt-4.1-mini": ModelPrice(0.40, 1.60),
+    # --- embeddings (input only) ---
+    "amazon.titan-embed-text-v2:0": ModelPrice(0.02, 0.0),
+    "text-embedding-3-small": ModelPrice(0.02, 0.0),
+    "text-embedding-3-large": ModelPrice(0.13, 0.0),
+    # --- test double ---
+    "fake-echo": ModelPrice(0.50, 1.50),
+}
+
+_FALLBACK = ModelPrice(0.15, 0.60)
+_loaded_override = False
+
+
+def _load_override() -> None:
+    global _loaded_override
+    if _loaded_override:
+        return
+    _loaded_override = True
+    path = os.getenv("CACHELLM_PRICING_FILE", "").strip()
+    if not path or not os.path.exists(path):
+        return
+    with open(path, encoding="utf-8") as fh:
+        raw = json.load(fh)
+    for model, entry in raw.items():
+        PRICES[model] = ModelPrice(float(entry["input_per_m"]), float(entry["output_per_m"]))
+
+
+def normalise_model_id(model: str) -> str:
+    """Strip routing and inference-profile decorations from a model id.
+
+    ``bedrock/us.amazon.nova-lite-v1:0`` and ``amazon.nova-lite-v1:0`` are the
+    same billable model, and cross-region profiles just add a region prefix.
+    """
+    m = model.strip()
+    if "/" in m:
+        m = m.split("/", 1)[1]
+    for region_prefix in ("us.", "eu.", "apac.", "ap."):
+        if m.startswith(region_prefix):
+            m = m[len(region_prefix) :]
+            break
+    return m
+
+
+def price_for(model: str) -> ModelPrice:
+    _load_override()
+    key = normalise_model_id(model)
+    if key in PRICES:
+        return PRICES[key]
+    # Prefix match so undated model revisions still price sensibly.
+    for known, price in PRICES.items():
+        if key.startswith(known.split("-2024")[0].split("-v1:")[0]):
+            return price
+    return _FALLBACK
+
+
+def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
+    return price_for(model).cost(input_tokens, output_tokens)

cachellm/providers/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from cachellm.providers.base import Provider, ProviderResult, StreamEvent
+from cachellm.providers.fake import FakeProvider
+from cachellm.providers.registry import ProviderRegistry
+
+__all__ = ["FakeProvider", "Provider", "ProviderRegistry", "ProviderResult", "StreamEvent"]

cachellm/providers/base.py

@@ -0,0 +1,84 @@
+"""Provider interface.
+
+Everything the cache needs from an upstream model lives behind two methods.
+Keeping this surface tiny is what makes the cache provider-agnostic: a cached
+Bedrock answer never gets served to an OpenAI request (provider is part of the
+namespace), but the caching logic itself is written once.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from typing import Any
+
+from cachellm.models import ChatCompletionRequest
+
+# Upstream "finished cleanly" markers, normalised to the OpenAI vocabulary.
+FINISH_REASON_MAP = {
+    "end_turn": "stop",
+    "stop_sequence": "stop",
+    "stop": "stop",
+    "max_tokens": "length",
+    "length": "length",
+    "content_filtered": "content_filter",
+    "content_filter": "content_filter",
+    "tool_use": "tool_calls",
+    "tool_calls": "tool_calls",
+    "guardrail_intervened": "content_filter",
+}
+
+
+def normalise_finish_reason(raw: str | None) -> str:
+    if not raw:
+        return "stop"
+    return FINISH_REASON_MAP.get(str(raw).lower(), str(raw).lower())
+
+
+@dataclass
+class ProviderResult:
+    text: str
+    model: str
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    finish_reason: str = "stop"
+    tool_calls: list[dict[str, Any]] | None = None
+    raw: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class StreamEvent:
+    """One step of a streamed completion."""
+
+    delta: str = ""
+    finish_reason: str | None = None
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+class Provider(abc.ABC):
+    name: str
+
+    @abc.abstractmethod
+    async def complete(self, request: ChatCompletionRequest) -> ProviderResult: ...
+
+    @abc.abstractmethod
+    def stream(self, request: ChatCompletionRequest) -> AsyncIterator[StreamEvent]: ...
+
+    def resolve_model(self, model: str) -> str:
+        """Strip the routing prefix, if any, to get the upstream model id."""
+        return model.split("/", 1)[1] if "/" in model else model
+
+    async def close(self) -> None:  # pragma: no cover - default no-op
+        return None
+
+    @staticmethod
+    def approx_tokens(text: str) -> int:
+        """Rough token count for providers that do not report usage.
+
+        Four characters per token is the widely used English approximation. It
+        is only used when the upstream is silent, and it is flagged as an
+        estimate wherever it reaches a number the user sees.
+        """
+        return max(1, len(text) // 4)

cachellm/providers/bedrock.py

@@ -0,0 +1,238 @@
+"""AWS Bedrock provider using the Converse API.
+
+Converse is the reason this project needs no Anthropic or Meta key: one request
+shape reaches Nova, Claude, Llama and Mistral alike. The interesting work here
+is the translation, because Bedrock's contract differs from OpenAI's in three
+ways that bite in production:
+
+* system prompts are a separate top-level field, not a message;
+* messages must strictly alternate user/assistant and must start with user;
+* boto3 is synchronous, so every call is pushed to a worker thread and the
+  streaming iterator is bridged onto the event loop by hand.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import AsyncIterator
+from typing import Any
+
+import structlog
+
+from cachellm.errors import UpstreamError
+from cachellm.models import ChatCompletionRequest, ChatMessage
+from cachellm.providers.base import (
+    Provider,
+    ProviderResult,
+    StreamEvent,
+    normalise_finish_reason,
+)
+from cachellm.settings import Settings
+
+log = structlog.get_logger(__name__)
+
+BEDROCK_VENDOR_PREFIXES = (
+    "amazon.",
+    "anthropic.",
+    "meta.",
+    "mistral.",
+    "cohere.",
+    "ai21.",
+    "deepseek.",
+    "qwen.",
+    "openai.",
+    "google.",
+    "moonshot",
+    "minimax.",
+    "nvidia.",
+    "zai.",
+    "writer.",
+)
+
+
+def _explain(exc: Exception) -> str:
+    """Turn boto's less obvious failures into something actionable.
+
+    The credential one is worth special-casing: the AWS CLI's `aws login` flow
+    stores credentials through a provider that needs the optional CRT extra,
+    and boto's own message names the package without naming the install.
+    """
+    text = str(exc)
+    if "botocore[crt]" in text or "Missing Dependency" in text:
+        return (
+            "Bedrock call failed: your AWS credentials come from `aws login`, which needs "
+            "the optional CRT extra. Install it with `uv sync --extra aws` "
+            "(or `pip install 'botocore[crt]'`), then restart the proxy. "
+            "Static keys and SSO profiles do not need this."
+        )
+    if "ExpiredToken" in text or "security token included in the request is expired" in text:
+        return "Bedrock call failed: AWS credentials have expired. Run `aws login` again."
+    if "AccessDeniedException" in text:
+        return (
+            f"Bedrock call failed: access denied. Check the model is enabled in this region "
+            f"and the identity has bedrock:InvokeModel. Original: {text[:200]}"
+        )
+    if "ThrottlingException" in text or "TooManyRequests" in text:
+        return f"Bedrock call failed: throttled by AWS. Lower concurrency and retry. {text[:160]}"
+    if "ValidationException" in text:
+        return f"Bedrock rejected the request: {text[:280]}"
+    return f"Bedrock call failed: {text[:300]}"
+
+
+def looks_like_bedrock_model(model: str) -> bool:
+    candidate = model.split("/", 1)[1] if model.startswith("bedrock/") else model
+    for region in ("us.", "eu.", "apac.", "ap."):
+        if candidate.startswith(region):
+            candidate = candidate[len(region) :]
+            break
+    return candidate.startswith(BEDROCK_VENDOR_PREFIXES)
+
+
+def to_converse_messages(messages: list[ChatMessage]) -> list[dict[str, Any]]:
+    """Convert OpenAI messages into Converse messages, merging same-role runs."""
+    converted: list[dict[str, Any]] = []
+    for message in messages:
+        if message.role in ("system", "developer"):
+            continue
+        role = "assistant" if message.role == "assistant" else "user"
+        text = message.as_text()
+        if not text.strip():
+            continue
+        if converted and converted[-1]["role"] == role:
+            converted[-1]["content"].append({"text": text})
+        else:
+            converted.append({"role": role, "content": [{"text": text}]})
+    # Converse rejects a conversation that opens with the assistant.
+    while converted and converted[0]["role"] != "user":
+        converted.pop(0)
+    return converted
+
+
+def build_converse_kwargs(request: ChatCompletionRequest, model_id: str) -> dict[str, Any]:
+    inference: dict[str, Any] = {}
+    if request.effective_max_tokens:
+        inference["maxTokens"] = int(request.effective_max_tokens)
+    if request.temperature is not None:
+        inference["temperature"] = float(request.temperature)
+    if request.top_p is not None:
+        inference["topP"] = float(request.top_p)
+    if request.stop:
+        stops = [request.stop] if isinstance(request.stop, str) else list(request.stop)
+        inference["stopSequences"] = stops[:4]
+
+    kwargs: dict[str, Any] = {
+        "modelId": model_id,
+        "messages": to_converse_messages(request.messages),
+    }
+    system_text = request.system_text()
+    if system_text:
+        kwargs["system"] = [{"text": system_text}]
+    if inference:
+        kwargs["inferenceConfig"] = inference
+    return kwargs
+
+
+class BedrockProvider(Provider):
+    name = "bedrock"
+
+    def __init__(self, settings: Settings) -> None:
+        self._settings = settings
+        self._client = None
+        self._lock = threading.Lock()
+
+    def _get_client(self):
+        if self._client is None:
+            with self._lock:
+                if self._client is None:
+                    import boto3
+                    from botocore.config import Config
+
+                    session_kwargs: dict[str, Any] = {"region_name": self._settings.aws_region}
+                    if self._settings.aws_profile:
+                        session_kwargs["profile_name"] = self._settings.aws_profile
+                    session = boto3.Session(**session_kwargs)
+                    self._client = session.client(
+                        "bedrock-runtime",
+                        config=Config(
+                            retries={
+                                "max_attempts": self._settings.provider_max_retries,
+                                "mode": "standard",
+                            },
+                            read_timeout=int(self._settings.request_timeout),
+                            connect_timeout=10,
+                        ),
+                    )
+        return self._client
+
+    # ---------------------------------------------------------------- complete
+    async def complete(self, request: ChatCompletionRequest) -> ProviderResult:
+        model_id = self.resolve_model(request.model)
+        kwargs = build_converse_kwargs(request, model_id)
+        try:
+            response = await asyncio.to_thread(lambda: self._get_client().converse(**kwargs))
+        except Exception as exc:  # boto raises many distinct client errors
+            log.warning("bedrock_error", model=model_id, error=str(exc)[:300])
+            raise UpstreamError(_explain(exc)) from exc
+
+        message = response.get("output", {}).get("message", {})
+        text = "".join(
+            block.get("text", "") for block in message.get("content", []) if "text" in block
+        )
+        usage = response.get("usage", {})
+        return ProviderResult(
+            text=text,
+            model=model_id,
+            prompt_tokens=int(usage.get("inputTokens", 0) or 0),
+            completion_tokens=int(usage.get("outputTokens", 0) or 0),
+            finish_reason=normalise_finish_reason(response.get("stopReason")),
+            raw={"metrics": response.get("metrics", {})},
+        )
+
+    # ------------------------------------------------------------------ stream
+    async def stream(self, request: ChatCompletionRequest) -> AsyncIterator[StreamEvent]:
+        model_id = self.resolve_model(request.model)
+        kwargs = build_converse_kwargs(request, model_id)
+        loop = asyncio.get_running_loop()
+        queue: asyncio.Queue[StreamEvent | Exception | None] = asyncio.Queue()
+
+        def pump() -> None:
+            """Consume the synchronous boto EventStream on a worker thread."""
+            try:
+                response = self._get_client().converse_stream(**kwargs)
+                for event in response.get("stream", []):
+                    if "contentBlockDelta" in event:
+                        delta = event["contentBlockDelta"].get("delta", {}).get("text", "")
+                        if delta:
+                            loop.call_soon_threadsafe(queue.put_nowait, StreamEvent(delta=delta))
+                    elif "messageStop" in event:
+                        reason = normalise_finish_reason(event["messageStop"].get("stopReason"))
+                        loop.call_soon_threadsafe(
+                            queue.put_nowait, StreamEvent(finish_reason=reason)
+                        )
+                    elif "metadata" in event:
+                        usage = event["metadata"].get("usage", {})
+                        loop.call_soon_threadsafe(
+                            queue.put_nowait,
+                            StreamEvent(
+                                prompt_tokens=int(usage.get("inputTokens", 0) or 0),
+                                completion_tokens=int(usage.get("outputTokens", 0) or 0),
+                            ),
+                        )
+            except Exception as exc:  # surfaced to the consumer below
+                loop.call_soon_threadsafe(queue.put_nowait, exc)
+            finally:
+                loop.call_soon_threadsafe(queue.put_nowait, None)
+
+        task = asyncio.create_task(asyncio.to_thread(pump))
+        try:
+            while True:
+                item = await queue.get()
+                if item is None:
+                    break
+                if isinstance(item, Exception):
+                    log.warning("bedrock_stream_error", model=model_id, error=str(item)[:300])
+                    raise UpstreamError(_explain(item)) from item
+                yield item
+        finally:
+            task.cancel()

cachellm/providers/fake.py

@@ -0,0 +1,56 @@
+"""Deterministic in-process provider.
+
+Every test, the whole CI pipeline and the reproducible benchmark run against
+this. It costs nothing, never rate-limits, and returns byte-identical answers
+on every machine, which is what makes the published hit-rate numbers something
+a reader can re-run rather than take on trust.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from collections.abc import AsyncIterator
+
+from cachellm.models import ChatCompletionRequest
+from cachellm.providers.base import Provider, ProviderResult, StreamEvent
+
+
+class FakeProvider(Provider):
+    name = "fake"
+
+    def __init__(self, latency_ms: float = 0.0, answers: dict[str, str] | None = None) -> None:
+        self.latency_ms = latency_ms
+        self.answers = answers or {}
+        self.calls = 0
+
+    def _answer(self, prompt: str) -> str:
+        if prompt in self.answers:
+            return self.answers[prompt]
+        digest = hashlib.blake2b(prompt.strip().lower().encode(), digest_size=6).hexdigest()
+        return f"[fake] answer to {prompt.strip()[:80]!r} (id {digest})"
+
+    async def complete(self, request: ChatCompletionRequest) -> ProviderResult:
+        self.calls += 1
+        if self.latency_ms:
+            await asyncio.sleep(self.latency_ms / 1000.0)
+        prompt = request.last_user_text()
+        text = self._answer(prompt)
+        return ProviderResult(
+            text=text,
+            model=self.resolve_model(request.model),
+            prompt_tokens=self.approx_tokens(prompt),
+            completion_tokens=self.approx_tokens(text),
+            finish_reason="stop",
+        )
+
+    async def stream(self, request: ChatCompletionRequest) -> AsyncIterator[StreamEvent]:
+        result = await self.complete(request)
+        words = result.text.split(" ")
+        for i, word in enumerate(words):
+            yield StreamEvent(delta=word if i == 0 else f" {word}")
+        yield StreamEvent(
+            finish_reason="stop",
+            prompt_tokens=result.prompt_tokens,
+            completion_tokens=result.completion_tokens,
+        )