tokenbiryani 0.2.0__py3-none-any.whl

tokenbiryani/core/router.py

@@ -0,0 +1,323 @@
+"""Routing: filter, score, pick — and record why every loser lost.
+
+The decision object is deliberately verbose. An operator who cannot see why a request
+went where it went has no reason to trust the router, so every candidate carries its
+weighted terms and a verdict, and that travels straight through to the inspector.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterable, List, Optional, Protocol, Sequence
+
+from ..config import RoutingWeights
+from .account import AccountRuntime
+from .limits import TokenEstimate
+
+logger = logging.getLogger("tokenbiryani")
+
+
+#: Entry point group third-party packages publish strategies under.
+PLUGIN_GROUP = "tokenbiryani.strategies"
+
+
+class Scorer(Protocol):
+    """A strategy that computes its own terms instead of reweighting the built-ins.
+
+    Terms are summed to produce the score, and each one is shown per candidate in the
+    request inspector — so name them for a reader, not for the arithmetic.
+    """
+
+    def terms(
+        self,
+        account: AccountRuntime,
+        ctx: RoutingContext,
+        weights: RoutingWeights,
+        index: int,
+        count: int,
+    ) -> Dict[str, float]:
+        ...  # pragma: no cover - protocol
+
+
+@dataclass
+class StrategySpec:
+    """Per-strategy multipliers over the operator's base weights."""
+
+    affinity: float = 1.0
+    headroom: float = 1.0
+    priority: float = 1.0
+    load: float = 1.0
+    errors: float = 1.0
+    cost: float = 1.0
+    rotate: bool = False
+
+
+STRATEGIES: Dict[str, StrategySpec] = {
+    # Affinity first, then headroom. Keeps prompt caches warm; the default.
+    "sticky_headroom": StrategySpec(),
+    # Pure most-available. Correct for stateless batch traffic, wasteful for chat.
+    "headroom": StrategySpec(affinity=0.0),
+    # Drain the cheap accounts, spill upward.
+    "cost_tiered": StrategySpec(affinity=0.5, headroom=0.25, cost=8.0),
+    # Strict ordered failover: primary, then backup.
+    "priority": StrategySpec(affinity=0.25, headroom=0.25, priority=8.0),
+    # Baseline. Shipped so the default can be benchmarked against it.
+    "least_loaded": StrategySpec(affinity=0.0, headroom=0.2, load=8.0),
+    # Baseline. Ignores every signal on purpose.
+    "round_robin": StrategySpec(
+        affinity=0.0, headroom=0.0, priority=0.0, load=0.0, errors=0.0, cost=0.0, rotate=True
+    ),
+}
+
+
+#: Strategies registered at runtime, by this process or by a plugin.
+CUSTOM_STRATEGIES: Dict[str, Scorer] = {}
+
+_plugins_loaded = False
+
+
+def register_strategy(name: str, strategy) -> None:
+    """Register a strategy under `name`.
+
+    Accepts a StrategySpec (reweights the built-in terms) or any object with a
+    `terms(...)` method (computes its own).
+    """
+    if isinstance(strategy, StrategySpec):
+        STRATEGIES[name] = strategy
+        CUSTOM_STRATEGIES.pop(name, None)
+        return
+    if not hasattr(strategy, "terms"):
+        raise TypeError(
+            f"strategy {name!r} must be a StrategySpec or expose a terms() method"
+        )
+    CUSTOM_STRATEGIES[name] = strategy
+    STRATEGIES.pop(name, None)
+
+
+def load_plugin_strategies(force: bool = False) -> List[str]:
+    """Import strategies published under the `tokenbiryani.strategies` entry point.
+
+    A plugin that fails to import is logged and skipped: a broken third-party
+    package must not stop the gateway from starting on its built-in strategies.
+    """
+    global _plugins_loaded
+    if _plugins_loaded and not force:
+        return []
+    _plugins_loaded = True
+    loaded: List[str] = []
+    for entry in _entry_points(PLUGIN_GROUP):
+        try:
+            resolved = entry.load()
+            if callable(resolved) and not isinstance(resolved, StrategySpec):
+                candidate = resolved() if not hasattr(resolved, "terms") else resolved
+            else:
+                candidate = resolved
+            register_strategy(entry.name, candidate)
+            loaded.append(entry.name)
+        except Exception as exc:  # noqa: BLE001 - a bad plugin must not be fatal
+            logger.warning("could not load routing strategy %r: %s", entry.name, exc)
+    return loaded
+
+
+def _entry_points(group: str):
+    """importlib.metadata changed shape in 3.10; support both."""
+    try:
+        from importlib.metadata import entry_points
+    except ImportError:  # pragma: no cover - Python < 3.8 only
+        return []
+    found = entry_points()
+    selector = getattr(found, "select", None)
+    if selector is not None:
+        return list(selector(group=group))
+    return list(found.get(group, []))
+
+
+def available_strategies() -> List[str]:
+    load_plugin_strategies()
+    return sorted(set(STRATEGIES) | set(CUSTOM_STRATEGIES))
+
+
+@dataclass
+class RoutingContext:
+    model: str
+    estimate: TokenEstimate
+    session_key: str
+    now: float
+    sticky_owner: Optional[str] = None
+    pool: Sequence[str] = ()
+    exclude: Sequence[str] = ()
+
+
+@dataclass
+class ScoredCandidate:
+    account_id: str
+    verdict: str
+    score: Optional[float] = None
+    terms: Dict[str, float] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "account_id": self.account_id,
+            "verdict": self.verdict,
+            "score": None if self.score is None else round(self.score, 4),
+            "terms": {k: round(v, 4) for k, v in self.terms.items()},
+        }
+
+
+@dataclass
+class Decision:
+    chosen: Optional[AccountRuntime]
+    candidates: List[ScoredCandidate]
+    strategy: str
+    affinity_honored: bool = False
+    affinity_broken: bool = False
+    #: True when nothing is eligible *now* but something will be after a reset
+    retryable_later: bool = False
+    soonest_available: Optional[float] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "strategy": self.strategy,
+            "chosen": self.chosen.id if self.chosen else None,
+            "affinity_honored": self.affinity_honored,
+            "affinity_broken": self.affinity_broken,
+            "soonest_available": (
+                round(self.soonest_available, 1) if self.soonest_available is not None else None
+            ),
+            "candidates": [c.to_dict() for c in self.candidates],
+        }
+
+
+def _jitter(session_key: str, account_id: str) -> float:
+    """Deterministic tie-break: stable under replay, spread across sessions."""
+    digest = hashlib.sha256((session_key + "|" + account_id).encode("utf-8")).digest()
+    return int.from_bytes(digest[:4], "big") / float(1 << 32) * 1e-6
+
+
+class Router:
+    def __init__(self, strategy: str, weights: RoutingWeights) -> None:
+        load_plugin_strategies()
+        if strategy not in STRATEGIES and strategy not in CUSTOM_STRATEGIES:
+            raise ValueError(
+                "unknown routing strategy {!r}; known: {}".format(
+                    strategy, ", ".join(available_strategies())
+                )
+            )
+        self.strategy = strategy
+        self.scorer: Optional[Scorer] = CUSTOM_STRATEGIES.get(strategy)
+        self.spec = STRATEGIES.get(strategy, StrategySpec())
+        self.weights = weights
+        self._cursor = 0
+
+    def select(
+        self, accounts: Iterable[AccountRuntime], ctx: RoutingContext
+    ) -> Decision:
+        accounts = list(accounts)
+        excluded = set(ctx.exclude)
+        pool = set(ctx.pool) if ctx.pool else None
+
+        eligible: List[AccountRuntime] = []
+        candidates: List[ScoredCandidate] = []
+        soonest: Optional[float] = None
+
+        for account in accounts:
+            if pool is not None and account.id not in pool:
+                candidates.append(
+                    ScoredCandidate(account.id, "filtered — outside this key's pool")
+                )
+                continue
+            if account.id in excluded:
+                candidates.append(
+                    ScoredCandidate(account.id, "filtered — already attempted")
+                )
+                continue
+            if account.mirror.exceeds_capacity(ctx.estimate):
+                candidates.append(
+                    ScoredCandidate(account.id, "filtered — request exceeds account limit")
+                )
+                continue
+            reason = account.eligible(ctx.model, ctx.estimate, ctx.now)
+            if reason is not None:
+                candidates.append(ScoredCandidate(account.id, "filtered — " + reason))
+                wait = account.cooling_for(ctx.now)
+                if wait is None and "headroom" in reason:
+                    wait = account.mirror.next_reset(ctx.now)
+                if wait is not None:
+                    soonest = wait if soonest is None else min(soonest, wait)
+                continue
+            eligible.append(account)
+
+        if not eligible:
+            return Decision(
+                None,
+                candidates,
+                self.strategy,
+                retryable_later=soonest is not None,
+                soonest_available=soonest,
+            )
+
+        tiers = [a.config.cost_tier for a in eligible]
+        low, high = min(tiers), max(tiers)
+        spread = (high - low) or 1.0
+
+        scored: List[Any] = []
+        for index, account in enumerate(eligible):
+            terms = self._terms(account, ctx, index, len(eligible), low, spread)
+            total = sum(terms.values()) + _jitter(ctx.session_key, account.id)
+            scored.append((total, account, terms))
+
+        scored.sort(key=lambda item: item[0], reverse=True)
+        best_score, best, best_terms = scored[0]
+
+        for total, account, terms in scored:
+            candidates.append(
+                ScoredCandidate(
+                    account.id,
+                    "chosen" if account is best else "eligible",
+                    score=total,
+                    terms=terms,
+                )
+            )
+
+        if self.spec.rotate:
+            self._cursor = (self._cursor + 1) % max(1, len(eligible))
+
+        honored = bool(ctx.sticky_owner) and ctx.sticky_owner == best.id
+        broken = bool(ctx.sticky_owner) and ctx.sticky_owner != best.id
+        return Decision(
+            best,
+            candidates,
+            self.strategy,
+            affinity_honored=honored,
+            affinity_broken=broken,
+        )
+
+    def _terms(
+        self,
+        account: AccountRuntime,
+        ctx: RoutingContext,
+        index: int,
+        count: int,
+        low_tier: float,
+        tier_spread: float,
+    ) -> Dict[str, float]:
+        if self.scorer is not None:
+            return self.scorer.terms(account, ctx, self.weights, index, count)
+        w, s = self.weights, self.spec
+        if s.rotate:
+            # Distance forward from the cursor; the account at the cursor scores highest.
+            offset = (index - self._cursor) % max(1, count)
+            return {"rotation": -float(offset)}
+
+        is_owner = 1.0 if ctx.sticky_owner == account.id else 0.0
+        normalized_cost = (account.config.cost_tier - low_tier) / tier_spread
+        return {
+            "affinity": w.affinity * s.affinity * is_owner,
+            "headroom": w.headroom * s.headroom * account.mirror.headroom(ctx.now),
+            "priority": w.priority * s.priority * account.config.priority,
+            "load": -w.load * s.load * account.load(),
+            "errors": -w.errors * s.errors * account.error_rate(),
+            "cost": -w.cost * s.cost * normalized_cost,
+        }

tokenbiryani/core/secrets.py

@@ -0,0 +1,114 @@
+"""Encryption for credentials the gateway stores itself.
+
+Config-file accounts keep their credentials in the operator's own file, protected by
+file permissions. Accounts added through the API are different: the gateway writes
+them to a database that may be a shared Redis or a file that ends up in a backup. So
+those are encrypted before they leave the process, and the key lives somewhere the
+database does not.
+
+Key resolution, in order:
+  1. ``TOKENBIRYANI_SECRET_KEY`` — a urlsafe-base64 Fernet key. Right for containers.
+  2. A key file beside the store, created 0600 on first use. Right for a laptop.
+
+Losing the key means the stored credentials are unrecoverable, which is the point.
+"""
+
+from __future__ import annotations
+
+import os
+import stat
+from typing import Optional
+
+ENV_VAR = "TOKENBIRYANI_SECRET_KEY"
+DEFAULT_KEY_FILE = "tokenbiryani.key"
+
+#: Ciphertext is prefixed so a future scheme can be told apart from this one.
+PREFIX = "v1:"
+
+
+class SecretError(RuntimeError):
+    """Raised when a secret cannot be encrypted or read back."""
+
+
+def _fernet(key: bytes):
+    try:
+        from cryptography.fernet import Fernet
+    except ImportError as exc:  # pragma: no cover - error path only
+        raise SecretError(
+            "storing account credentials needs the cryptography package, which is a "
+            "base dependency — reinstall with: pip install --upgrade tokenbiryani"
+        ) from exc
+    return Fernet(key)
+
+
+def generate_key() -> bytes:
+    # Guarded like _fernet: this is reached on the first attempt to store an
+    # account, and a bare ModuleNotFoundError there tells the operator nothing
+    # about which extra is missing.
+    try:
+        from cryptography.fernet import Fernet
+    except ImportError as exc:  # pragma: no cover - error path only
+        raise SecretError(
+            "storing account credentials needs the cryptography package, which is a "
+            "base dependency — reinstall with: pip install --upgrade tokenbiryani"
+        ) from exc
+
+    return Fernet.generate_key()
+
+
+def load_key(key_path: Optional[str] = None) -> bytes:
+    """Find the key, or create one. Never returns a key from a world-readable file."""
+    from_env = os.environ.get(ENV_VAR, "").strip()
+    if from_env:
+        return from_env.encode("utf-8")
+
+    path = os.path.abspath(os.path.expanduser(key_path or DEFAULT_KEY_FILE))
+    if os.path.exists(path):
+        mode = stat.S_IMODE(os.stat(path).st_mode)
+        if mode & 0o077:
+            raise SecretError(
+                f"{path} is readable by other users (mode {mode:o}). "
+                "Run: chmod 600 " + path
+            )
+        with open(path, "rb") as handle:
+            return handle.read().strip()
+
+    directory = os.path.dirname(path)
+    if directory and not os.path.isdir(directory):
+        os.makedirs(directory, exist_ok=True)
+    key = generate_key()
+    # Create with the right mode from the start; do not write then chmod.
+    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
+    with os.fdopen(descriptor, "wb") as handle:
+        handle.write(key)
+    return key
+
+
+class SecretBox:
+    """Encrypts and decrypts one kind of thing: a credential going into the store."""
+
+    def __init__(self, key: Optional[bytes] = None, key_path: Optional[str] = None) -> None:
+        self._key = key or load_key(key_path)
+        self._cipher = _fernet(self._key)
+
+    def encrypt(self, plaintext: str) -> str:
+        if not plaintext:
+            return ""
+        return PREFIX + self._cipher.encrypt(plaintext.encode("utf-8")).decode("ascii")
+
+    def decrypt(self, ciphertext: str) -> str:
+        if not ciphertext:
+            return ""
+        if not ciphertext.startswith(PREFIX):
+            # A value written before encryption existed, or hand-edited into the
+            # store. Accept it rather than lock the operator out of their own pool.
+            return ciphertext
+        from cryptography.fernet import InvalidToken
+
+        try:
+            return self._cipher.decrypt(ciphertext[len(PREFIX):].encode("ascii")).decode("utf-8")
+        except (InvalidToken, ValueError) as exc:
+            raise SecretError(
+                "could not decrypt a stored credential — the secret key has changed. "
+                "Restore the original key, or delete and re-add the account."
+            ) from exc

tokenbiryani/core/session.py

@@ -0,0 +1,117 @@
+"""Session affinity.
+
+Anthropic's prompt cache is scoped per credential. A conversation that hops accounts
+pays full price for a prefix it already cached somewhere else, so the router needs a
+stable key for "this conversation" without the client having to supply one.
+
+The fingerprint deliberately covers only the *stable* head of a request — system
+prompt, tool definitions, and the opening messages — so it stays constant as a
+conversation grows.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any, List, Mapping, Optional
+
+SESSION_HEADER = "x-tokenbiryani-session"
+
+#: How many leading messages take part in the fingerprint. Enough to distinguish two
+#: conversations that share a system prompt; few enough to survive history growth.
+LEADING_MESSAGES = 2
+
+#: Per-message character budget, so a long opening message can't dominate the hash.
+MESSAGE_BUDGET = 512
+
+
+def _flatten(content: Any, budget: int = MESSAGE_BUDGET) -> str:
+    """Content is a string or a list of typed blocks. Reduce it to stable text."""
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content[:budget]
+    if isinstance(content, Mapping):
+        kind = content.get("type", "")
+        if kind == "text":
+            return "text:" + str(content.get("text", ""))[:budget]
+        if kind == "tool_use":
+            return "tool_use:" + str(content.get("name", ""))
+        if kind == "tool_result":
+            return "tool_result:" + str(content.get("tool_use_id", ""))
+        return str(kind)
+    if isinstance(content, list):
+        parts: List[str] = []
+        remaining = budget
+        for block in content:
+            if remaining <= 0:
+                break
+            piece = _flatten(block, remaining)
+            remaining -= len(piece)
+            parts.append(piece)
+        return "|".join(parts)
+    return str(content)[:budget]
+
+
+def fingerprint(body: Mapping[str, Any]) -> str:
+    """Hash the stable prefix of a Messages request."""
+    parts: List[str] = ["model=" + str(body.get("model", ""))]
+    parts.append("system=" + _flatten(body.get("system"), budget=4096))
+
+    tools = body.get("tools")
+    if isinstance(tools, list):
+        names = []
+        for tool in tools:
+            if isinstance(tool, Mapping):
+                names.append(str(tool.get("name", "")))
+        parts.append("tools=" + ",".join(sorted(names)))
+
+    messages = body.get("messages")
+    if isinstance(messages, list):
+        for message in messages[:LEADING_MESSAGES]:
+            if isinstance(message, Mapping):
+                parts.append(
+                    "{}:{}".format(message.get("role", ""), _flatten(message.get("content")))
+                )
+
+    digest = hashlib.sha256("\x1f".join(parts).encode("utf-8", "replace")).hexdigest()
+    return digest[:32]
+
+
+def session_key(
+    body: Mapping[str, Any],
+    headers: Optional[Mapping[str, str]] = None,
+    scope: str = "",
+) -> str:
+    """An explicit header always wins; otherwise fall back to the fingerprint.
+
+    ``scope`` namespaces the result to one tenant. Without it the key is global, and
+    two tenants collide in two ways that both matter: either can steer the other's
+    affinity by choosing the same ``X-TokenBiryani-Session`` value, and two callers
+    running the same agent share a fingerprint whenever their system prompt, tool
+    names and opening messages match — which is the normal case for one popular
+    client, not a contrived one. Pass the virtual key's name and neither happens.
+    """
+    if headers:
+        for name, value in headers.items():
+            if name.lower() == SESSION_HEADER and value.strip():
+                return _scoped(scope, "hdr:" + value.strip()[:128])
+    return _scoped(scope, "fp:" + fingerprint(body))
+
+
+def _scoped(scope: str, key: str) -> str:
+    """Prefix a session key with its tenant, when there is one."""
+    if not scope:
+        return key
+    return "k:" + str(scope)[:64] + "|" + key
+
+
+def cache_prefix_size(body: Mapping[str, Any]) -> int:
+    """Rough size of the cacheable prefix — how much is at stake in a cache break."""
+    try:
+        return len(json.dumps(
+            {k: body.get(k) for k in ("system", "tools") if body.get(k) is not None},
+            ensure_ascii=False,
+        ))
+    except (TypeError, ValueError):
+        return 0

tokenbiryani/dashboard/__init__.py

@@ -0,0 +1,56 @@
+"""The operator console: a server-rendered page and its stylesheet, no build step.
+
+It ships inside the package deliberately. Adding a Node toolchain to a `pipx install`
+to get a dashboard would be a bad trade for a tool whose whole appeal is that it runs
+from one command.
+
+The CSS is a separate file rather than a `<style>` block because it is a design system
+— tokens, primitives, components — that is read and edited on its own terms. Splitting
+it costs one cached request and no build step at all.
+
+Both are re-read when they change on disk, so editing the console and refreshing the
+browser is the whole loop — no restart, and no build.
+"""
+
+from __future__ import annotations
+
+import os
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+CONSOLE_PATH = os.path.join(_HERE, "console.html")
+CONSOLE_CSS_PATH = os.path.join(_HERE, "console.css")
+
+#: path -> (mtime, contents)
+_cache: dict = {}
+
+
+def _read(path: str) -> str:
+    """Read once, and again whenever the file changes underneath us.
+
+    Caching outright was wrong. In a container with the source bind-mounted, and
+    under `serve --reload` (which watches Python files, not these), editing the
+    console meant refreshing the browser and seeing the old page until something
+    else happened to restart the process.
+
+    A stat per request is nothing next to proxying a model call, and it makes
+    "edit the console, refresh" true rather than nearly true.
+    """
+    try:
+        mtime = os.path.getmtime(path)
+    except OSError:
+        mtime = None
+    cached = _cache.get(path)
+    if cached is not None and cached[0] == mtime:
+        return cached[1]
+    with open(path, encoding="utf-8") as handle:
+        contents = handle.read()
+    _cache[path] = (mtime, contents)
+    return contents
+
+
+def console_html() -> str:
+    return _read(CONSOLE_PATH)
+
+
+def console_css() -> str:
+    return _read(CONSOLE_CSS_PATH)