tierwise 0.8.0__py3-none-any.whl

tierwise/__init__.py

@@ -0,0 +1,3 @@
+"""tierwise: cost-aware, self-hosted LLM gateway with a learned router."""
+
+__version__ = "0.1.0"

tierwise/__main__.py

@@ -0,0 +1,101 @@
+"""Run the gateway (`tierwise`, `python -m tierwise`), manage API keys (`tierwise keys`) or
+download the learned router's encoder (`tierwise fetch-model`)."""
+
+import argparse
+import asyncio
+import sys
+import time
+from pathlib import Path
+
+import uvicorn
+
+from tierwise import fetch
+from tierwise.api.app import create_app
+from tierwise.budgets import month_of
+from tierwise.config import CACHE_ROUTER_DIR, Settings
+from tierwise.keys import generate_key, hash_key
+from tierwise.store import RequestLog
+
+
+def _rpm(value: str) -> int:
+    rpm = int(value)
+    if rpm < 1:
+        raise argparse.ArgumentTypeError("must be at least 1 (omit --rpm for no limit)")
+    return rpm
+
+
+def _budget(value: str) -> float:
+    budget = float(value)
+    if budget < 0:
+        raise argparse.ArgumentTypeError("must be zero or more (omit --budget for no limit)")
+    return budget
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="tierwise")
+    commands = parser.add_subparsers(dest="command")
+    commands.add_parser("serve", help="run the gateway (the default)")
+    keys = commands.add_parser("keys", help="manage gateway API keys").add_subparsers(
+        dest="action", required=True
+    )
+    create = keys.add_parser("create", help="create a key (printed once)")
+    create.add_argument("--name", required=True)
+    create.add_argument("--budget", type=_budget, help="monthly budget in USD")
+    create.add_argument("--rpm", type=_rpm, help="requests per minute")
+    keys.add_parser("list", help="list keys (no secrets)")
+    revoke = keys.add_parser("revoke", help="revoke a key by id")
+    revoke.add_argument("id")
+    fetch_model = commands.add_parser(
+        "fetch-model", help="download the learned router's encoder (90 MB, sha256-checked)"
+    )
+    fetch_model.add_argument("--out", type=Path, default=CACHE_ROUTER_DIR)
+    return parser
+
+
+async def _keys(args: argparse.Namespace, db_url: str) -> int:
+    log = RequestLog(db_url)
+    await log.start()
+    try:
+        if args.action == "create":
+            key = generate_key()
+            row = await log.create_key(args.name, hash_key(key), key[:10], args.budget, args.rpm)
+            print(f"id: {row.id}\nkey: {key}\nStore the key now: it is not shown again.")
+            return 0
+        if args.action == "revoke":
+            if await log.revoke_key(args.id):
+                print(f"revoked {args.id}")
+                return 0
+            print(f"no active key with id {args.id}", file=sys.stderr)
+            return 1
+        spend = await log.load_month_spend(month_of(time.time()))
+        print(f"{'id':<14}{'name':<16}{'prefix':<12}{'budget':>10}{'spent':>10}{'rpm':>6}  status")
+        for row in await log.list_keys():
+            budget = "-" if row.monthly_budget_usd is None else f"{row.monthly_budget_usd:.2f}"
+            rpm = "-" if row.rpm is None else str(row.rpm)
+            status = "active" if row.revoked_ts is None else "revoked"
+            print(f"{row.id:<14}{row.name:<16}{row.prefix:<12}{budget:>10}"
+                  f"{spend.get(row.id, 0.0):>10.2f}{rpm:>6}  {status}")  # fmt: skip
+        return 0
+    finally:
+        await log.stop()
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _parser().parse_args(argv)
+    if args.command == "fetch-model":
+        try:
+            path = fetch.fetch_model(args.out)
+        except fetch.FetchError as exc:
+            print(f"fetch-model failed: {exc}", file=sys.stderr)
+            return 1
+        print(f"wrote {path}")
+        return 0
+    settings = Settings()
+    if args.command == "keys":
+        return asyncio.run(_keys(args, settings.db_url))
+    uvicorn.run(create_app(settings), host=settings.host, port=settings.port)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tierwise/affinity.py

@@ -0,0 +1,137 @@
+"""Conversation memory: which model served which prompt prefix, and when.
+
+Providers cache the start of a prompt, so the model that served a conversation's earlier
+turns can be cheaper for the next turn than a "cheaper" model whose cache is cold.
+"""
+
+import hashlib
+import json
+import time
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from tierwise.schema import ChatRequest
+from tierwise.tokens import estimate_text, text_of
+
+
+def _items(req: ChatRequest) -> list[tuple[str, int]]:
+    """The prompt as cache-able items, in the order providers see them: tools, then messages."""
+    items: list[tuple[str, int]] = []
+    if req.tools:
+        tools = [t.model_dump(exclude_none=True) for t in req.tools]
+        items.append((_canonical({"tools": tools}), estimate_text(json.dumps(tools))))
+    for m in req.messages:
+        items.append(
+            (_canonical(m.model_dump(exclude_none=True)), estimate_text(text_of(m.content)))
+        )
+    return items
+
+
+def _canonical(value: Any) -> str:
+    return json.dumps(value, sort_keys=True, separators=(",", ":"))
+
+
+def prefix_boundaries(req: ChatRequest) -> list[tuple[str, int]]:
+    """(rolling hash, cumulative estimated tokens) after each item of the prompt.
+
+    Providers cache the *beginning* of a prompt, so turn N+1 of a conversation can reuse the
+    cache written for turn N: the whole turn-N prompt is one of turn N+1's boundaries.
+    The last entry covers the full prompt.
+    """
+    boundaries: list[tuple[str, int]] = []
+    digest, tokens = b"", 0
+    for canonical, item_tokens in _items(req):
+        digest = hashlib.sha256(digest + canonical.encode()).digest()
+        tokens += item_tokens
+        boundaries.append((digest.hex(), tokens))
+    return boundaries
+
+
+def split_prefix(req: ChatRequest) -> tuple[str | None, int]:
+    """Hash and size of everything before the last user message (the conversation so far)."""
+    last_user = max((i for i, m in enumerate(req.messages) if m.role == "user"), default=None)
+    count = (1 if req.tools else 0) + (len(req.messages) if last_user is None else last_user)
+    if count == 0:
+        return None, 0
+    return prefix_boundaries(req)[count - 1]
+
+
+class AffinityTable:
+    """In-memory LRU of prefix hash -> {(model id, effort sent): when it last served it}."""
+
+    def __init__(
+        self, max_entries: int = 10_000, clock: Callable[[], float] = time.monotonic
+    ) -> None:
+        self._entries: OrderedDict[str, dict[tuple[str, str | None], float]] = OrderedDict()
+        self._max_entries = max_entries
+        self._clock = clock
+
+    def record(self, prefix_hash: str, model_id: str, effort: str | None = None) -> None:
+        served = self._entries.pop(prefix_hash, {})
+        served[(model_id, effort)] = self._clock()
+        self._entries[prefix_hash] = served
+        while len(self._entries) > self._max_entries:
+            self._entries.popitem(last=False)
+
+    def warm_models(self, prefix_hash: str | None) -> dict[str, float]:
+        """Model id -> most recent time it served this prefix, at any effort."""
+        latest: dict[str, float] = {}
+        for (model_id, _), served_at in self._entries.get(prefix_hash or "", {}).items():
+            latest[model_id] = max(served_at, latest.get(model_id, served_at))
+        return latest
+
+    def is_warm(
+        self,
+        prefix_hash: str | None,
+        model_id: str,
+        ttl: int | None,
+        effort: str | None = None,
+        effort_matters: bool = False,
+    ) -> bool:
+        """Whether ``model_id`` still has this prefix cached.
+
+        Where effort is part of the provider's cache key, only the same effort counts.
+        """
+        if ttl is None:
+            return False
+        now = self._clock()
+        return any(
+            mid == model_id and (not effort_matters or used == effort) and now - at <= ttl
+            for (mid, used), at in self._entries.get(prefix_hash or "", {}).items()
+        )
+
+    def longest_warm(
+        self,
+        boundaries: Sequence[tuple[str, int]],
+        model_id: str,
+        ttl: int | None,
+        effort: str | None = None,
+        effort_matters: bool = False,
+    ) -> int:
+        """Estimated tokens of the longest boundary ``model_id`` has warm in its cache."""
+        for prefix_hash, tokens in reversed(boundaries):
+            if self.is_warm(prefix_hash, model_id, ttl, effort, effort_matters):
+                return tokens
+        return 0
+
+    def chain_incumbent(
+        self, boundaries: Sequence[tuple[str, int]], ttl_for: Callable[[str], int | None]
+    ) -> str | None:
+        """The warm model for the longest boundary that has one."""
+        for prefix_hash, _ in reversed(boundaries):
+            if (model_id := self.incumbent(prefix_hash, ttl_for)) is not None:
+                return model_id
+        return None
+
+    def incumbent(
+        self, prefix_hash: str | None, ttl_for: Callable[[str], int | None]
+    ) -> str | None:
+        """The most recent model whose cache for this prefix should still be warm."""
+        now = self._clock()
+        warm = [
+            (served_at, model_id)
+            for model_id, served_at in self.warm_models(prefix_hash).items()
+            if (ttl := ttl_for(model_id)) is not None and now - served_at <= ttl
+        ]
+        return max(warm)[1] if warm else None

tierwise/api/app.py

@@ -0,0 +1,127 @@
+"""Application factory: registry, provider clients, learned router and the request log."""
+
+import asyncio
+import contextlib
+import time
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+import httpx
+from fastapi import FastAPI
+
+from tierwise import __version__
+from tierwise.affinity import AffinityTable
+from tierwise.api.errors import install_error_handlers
+from tierwise.api.routes import router
+from tierwise.budgets import Budgets, RateLimiter, month_of
+from tierwise.cache import SemanticCache
+from tierwise.config import Settings, registry_file
+from tierwise.keys import KeyStore
+from tierwise.metrics import Metrics
+from tierwise.providers.anthropic import AnthropicProvider
+from tierwise.providers.base import Provider
+from tierwise.providers.gemini import GeminiProvider
+from tierwise.providers.ollama import OllamaProvider
+from tierwise.providers.openai import OpenAIProvider
+from tierwise.registry import Registry, load_registry
+from tierwise.reliability.breaker import Breakers
+from tierwise.reliability.retry import RetryPolicy
+from tierwise.router.learned import load_learned
+from tierwise.store import RequestLog
+from tierwise.telemetry import configure_tracing
+
+LIMITS = httpx.Limits(max_keepalive_connections=100, max_connections=250, keepalive_expiry=30)
+TIMEOUT = httpx.Timeout(connect=5.0, read=120.0, write=30.0, pool=2.0)
+
+
+def create_app(
+    settings: Settings | None = None,
+    transport: httpx.AsyncBaseTransport | None = None,
+    registry: Registry | None = None,
+) -> FastAPI:
+    config = settings or Settings()
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+        tracing = configure_tracing(config.otel_exporter_otlp_endpoint)
+        # Fail fast: an invalid registry stops startup instead of breaking requests later.
+        app.state.registry = registry or load_registry(registry_file(config))
+
+        # One pooled client per provider for the app's whole life: no per-request handshakes.
+        clients: dict[str, httpx.AsyncClient] = {}
+        providers: dict[str, Provider] = {}
+
+        def client_for(name: str) -> httpx.AsyncClient:
+            clients[name] = httpx.AsyncClient(limits=LIMITS, timeout=TIMEOUT, transport=transport)
+            return clients[name]
+
+        if config.anthropic_api_key:
+            providers["anthropic"] = AnthropicProvider(
+                client_for("anthropic"), config.anthropic_api_key
+            )
+        if config.openai_api_key:
+            providers["openai"] = OpenAIProvider(
+                client_for("openai"), config.openai_api_key, config.openai_base_url
+            )
+        if config.gemini_api_key:
+            providers["gemini"] = GeminiProvider(client_for("gemini"), config.gemini_api_key)
+        if config.ollama_base_url:
+            providers["ollama"] = OllamaProvider(client_for("ollama"), config.ollama_base_url)
+        app.state.clients = clients
+        app.state.providers = providers
+        app.state.affinity = AffinityTable()
+        app.state.breakers = Breakers(
+            config.breaker_failures, config.breaker_window_s, config.breaker_cooldown_s
+        )
+        app.state.retry = RetryPolicy(config.retries, config.retry_base_s, config.retry_max_s)
+        app.state.stream_hold = (config.stream_hold_s, config.stream_hold_bytes)
+
+        # Loading and warming the ONNX session blocks; keep it off the event loop.
+        learned = await asyncio.to_thread(load_learned, config)
+        app.state.learned = learned
+        app.state.cache = (
+            SemanticCache(
+                learned.embed_text,
+                config.cache_similarity,
+                config.cache_ttl_s,
+                config.cache_max_entries,
+            )
+            if learned is not None and config.cache_enabled
+            else None
+        )
+        app.state.cache_max_temperature = config.cache_max_temperature
+
+        log = RequestLog(config.db_url)
+        await log.start()
+        app.state.log = log
+        keys = KeyStore()
+        keys.replace(await log.list_keys())
+        app.state.keys = keys
+        app.state.require_api_key = config.require_api_key
+        app.state.limiter = RateLimiter()
+        app.state.metrics = Metrics(
+            app.state.breakers, log, [m.id for m in app.state.registry.models]
+        )
+        budgets = Budgets(log)
+        month = month_of(time.time())
+        budgets.load(month, await log.load_month_spend(month))
+        app.state.budgets = budgets
+        refresh = asyncio.create_task(keys.refresh_forever(log, config.keys_refresh_s))
+        try:
+            yield
+        finally:
+            refresh.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await refresh
+            for client in clients.values():
+                await client.aclose()
+            await log.stop()
+            if learned is not None:
+                learned.close()
+            if tracing is not None:  # last: a slow collector must not delay the debit flush
+                tracing.shutdown()
+
+    app = FastAPI(title="tierwise", version=__version__, lifespan=lifespan)
+    install_error_handlers(app)
+    app.include_router(router)
+    return app

tierwise/api/auth.py

@@ -0,0 +1,22 @@
+"""Bearer-token authentication against the gateway's API keys."""
+
+from tierwise.api.errors import InvalidApiKey
+from tierwise.keys import KeyStore
+from tierwise.store import KeyRow
+
+
+def authenticate(header: str | None, store: KeyStore, required: bool) -> KeyRow | None:
+    """The key named by ``Authorization: Bearer <key>``.
+
+    When keys are required, a missing, malformed, unknown or revoked key is a 401. When they
+    aren't, a valid key is still resolved so its spend is tracked, and anything else is
+    anonymous (None).
+    """
+    parts = (header or "").split()
+    if len(parts) == 2 and parts[0].lower() == "bearer":
+        row = store.find(parts[1])
+        if row is not None:
+            return row
+    if required:
+        raise InvalidApiKey()
+    return None

tierwise/api/errors.py

@@ -0,0 +1,84 @@
+"""OpenAI-shaped error responses for everything the gateway can raise."""
+
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+
+from tierwise.providers.base import AuthError, FatalError, ProviderError, RetryableError
+
+
+class GatewayError(Exception):
+    def __init__(
+        self, status: int, code: str, message: str, type_: str = "invalid_request_error"
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.code = code
+        self.message = message
+        self.type = type_
+
+
+class InvalidApiKey(GatewayError):
+    def __init__(self) -> None:
+        super().__init__(401, "invalid_api_key", "Invalid or missing API key.")
+
+
+class RateLimited(GatewayError):
+    def __init__(self, retry_after: int) -> None:
+        message = f"Rate limit exceeded for this API key; retry in {retry_after} s."
+        super().__init__(429, "rate_limited", message, "rate_limit_error")
+        self.retry_after = retry_after
+
+
+class BudgetExceeded(GatewayError):
+    def __init__(self, message: str) -> None:
+        super().__init__(429, "budget_exceeded", message, "insufficient_quota")
+
+
+class AllCircuitsOpen(GatewayError):
+    """Every candidate's circuit breaker is open, so nothing was attempted."""
+
+    def __init__(self, skipped: list[str]) -> None:
+        names = ", ".join(skipped) or "none"
+        message = f"All candidate models are temporarily unavailable (circuit open: {names})."
+        super().__init__(503, "upstream_unavailable", message, "api_error")
+
+
+def error_response(status: int, code: str, message: str, type_: str) -> JSONResponse:
+    return JSONResponse(
+        status_code=status, content={"error": {"message": message, "type": type_, "code": code}}
+    )
+
+
+def gateway_error_from(exc: ProviderError) -> GatewayError:
+    if isinstance(exc, AuthError):
+        # Never echo the upstream body: it can contain (part of) our key.
+        message = f"Upstream provider '{exc.provider}' rejected the gateway's credentials."
+        return GatewayError(502, "upstream_auth_error", message, "api_error")
+    if isinstance(exc, FatalError):
+        return GatewayError(exc.status or 400, "invalid_request_error", exc.message)
+    if isinstance(exc, RetryableError):
+        message = f"Upstream provider '{exc.provider}' is unavailable: {exc.message}"
+        return GatewayError(503, "upstream_unavailable", message, "api_error")
+    return GatewayError(502, "upstream_error", exc.message, "api_error")
+
+
+def install_error_handlers(app: FastAPI) -> None:
+    @app.exception_handler(GatewayError)
+    async def _gateway(_: Request, exc: GatewayError) -> JSONResponse:
+        response = error_response(exc.status, exc.code, exc.message, exc.type)
+        if isinstance(exc, RateLimited):
+            response.headers["Retry-After"] = str(exc.retry_after)
+        return response
+
+    @app.exception_handler(ProviderError)
+    async def _provider(_: Request, exc: ProviderError) -> JSONResponse:
+        err = gateway_error_from(exc)
+        return error_response(err.status, err.code, err.message, err.type)
+
+    @app.exception_handler(RequestValidationError)
+    async def _validation(_: Request, exc: RequestValidationError) -> JSONResponse:
+        problems = "; ".join(
+            f"{'.'.join(str(p) for p in e['loc'])}: {e['msg']}" for e in exc.errors()
+        )
+        return error_response(400, "invalid_request_error", problems, "invalid_request_error")