tokenbiryani 0.2.0__py3-none-any.whl

tokenbiryani/core/diagnostics.py

@@ -0,0 +1,131 @@
+"""What the rate-limit headers on a real response actually say.
+
+This is the one check no mock-backed test can make: every test in this project runs
+against the suite's fake upstream, which encodes an assumption about how
+the upstream spells its headers. If the real API spells one differently, that window
+stays empty, the account reads as full, and the router silently degrades to
+round-robin — shredding the prompt cache while every screen looks healthy.
+
+It lives here, rather than inside `cmd_doctor`, because two callers need the same
+answer from the same list: `tokenbiryani doctor` in the terminal, and the console's
+verify step. Two copies of the expected-header list is exactly the drift the check
+exists to catch.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, List, Mapping
+
+from .limits import UNIFIED_WINDOWS, LimitMirror
+
+#: The three windows the limit mirror models, in the spelling the upstream uses.
+LIMIT_WINDOWS = ["requests", "input-tokens", "output-tokens"]
+
+#: ...and the three facts it needs about each of them.
+HEADER_TEMPLATES = [
+    "anthropic-ratelimit-{}-limit",
+    "anthropic-ratelimit-{}-remaining",
+    "anthropic-ratelimit-{}-reset",
+]
+
+#: A subscription session answers a different question. It never sends the triples
+#: above; it sends how much of each rolling window is spent. An account reporting
+#: these is not broken, so checking it against the API-key list would report nine
+#: faults where there are none — which is the exact false alarm this module exists
+#: to prevent, pointed the other way.
+UNIFIED_TEMPLATES = [
+    "anthropic-ratelimit-unified-{}-status",
+    "anthropic-ratelimit-unified-{}-utilization",
+    "anthropic-ratelimit-unified-{}-reset",
+]
+
+
+def expected_headers(unified: bool = False) -> List[str]:
+    """Every header the router depends on, in the order a report should show them."""
+    if unified:
+        return [
+            template.format(window)
+            for window in UNIFIED_WINDOWS
+            for template in UNIFIED_TEMPLATES
+        ]
+    return [
+        template.format(window)
+        for window in LIMIT_WINDOWS
+        for template in HEADER_TEMPLATES
+    ]
+
+
+def looks_unified(lowered: Mapping[str, str]) -> bool:
+    """True when the response is a subscription session's, not an API key's."""
+    return any(
+        str(name).startswith("anthropic-ratelimit-unified-") for name in lowered
+    )
+
+
+def inspect_headers(headers: Mapping[str, str], now: float = 0.0) -> Dict[str, Any]:
+    """Compare what arrived against what the mirror looks for, and parse it.
+
+    Returns the whole picture rather than a verdict, because the useful thing to show
+    a human is the mismatch itself: which header was expected, what came instead, and
+    what the router would therefore believe.
+    """
+    now = now or time.time()
+    lowered = {str(k).lower(): str(v) for k, v in dict(headers).items()}
+    unified = looks_unified(lowered)
+
+    checked = []
+    missing = []
+    for name in expected_headers(unified):
+        present = name in lowered
+        if not present:
+            missing.append(name)
+        checked.append({"header": name, "present": present, "value": lowered.get(name)})
+
+    # The real proof is not that the keys exist but that the mirror the router reads
+    # parsed something out of them.
+    mirror = LimitMirror(observable=not unified)
+    mirror.update_from_headers(dict(headers), now)
+    snapshot = mirror.snapshot(now)
+    if unified:
+        parsed = {
+            f"unified_{name}": snapshot["unified"][name]
+            for name in snapshot["unified"]
+        }
+    else:
+        parsed = {
+            label: {
+                "limit": snapshot[label]["limit"],
+                "remaining": snapshot[label]["remaining"],
+                "reset_in": snapshot[label]["reset_in"],
+            }
+            for label in ("requests", "input_tokens", "output_tokens")
+        }
+
+    if not missing:
+        consequence = ""
+    elif unified:
+        consequence = (
+            f"{len(missing)} of {len(expected_headers(True))} unified headers are "
+            "missing or spelled differently. This account routes on assumed headroom "
+            "instead of its real utilisation, so a nearly-exhausted subscription "
+            "still reads as half full."
+        )
+    else:
+        consequence = (
+            f"{len(missing)} of {len(expected_headers())} headers are missing or "
+            "spelled differently. Those windows stay empty, so this account reads "
+            "as full and routing degrades to round-robin — which shreds the prompt "
+            "cache while every meter looks healthy."
+        )
+
+    return {
+        "ok": not missing,
+        "family": "unified" if unified else "limits",
+        "headroom": snapshot["headroom"],
+        "checked": checked,
+        "missing": missing,
+        "returned": {k: v for k, v in sorted(lowered.items()) if k.startswith("anthropic-")},
+        "parsed": parsed,
+        "consequence": consequence,
+    }

tokenbiryani/core/estimator.py

@@ -0,0 +1,180 @@
+"""Predicting what a request will actually return, instead of believing its ceiling.
+
+`max_tokens` is the only honest upper bound available before a request runs, which
+is why the lease started out using it. It is also wildly wrong in the case this
+gateway exists to serve: an agent client sends `max_tokens: 32000` and returns a few
+hundred tokens.
+
+The lease is released afterwards, so no quota is *spent* on the difference. The cost
+is paid during the request instead. An account whose output window is leased at 32k
+reads as full, an account that reads as full is filtered out of routing, and a
+conversation whose owner is filtered out gets re-homed onto a credential that has
+never seen its prefix. That is how an estimation problem turns into a cache break,
+which is the most expensive thing that can happen here.
+
+Measured against a mirrored output window, believing the ceiling costs most of the
+pool's concurrency:
+
+    output limit 16,000/window, max_tokens  8,192 ->  1 concurrent request
+    output limit 64,000/window, max_tokens 32,000 ->  2 concurrent requests
+
+So: keep a rolling sample of what each model really returns, lease a high quantile
+of it, and never lease more than the caller's own ceiling — the prediction can only
+ever reserve *less* than today. Fall back to the ceiling until there is enough
+evidence to beat it, and again if the prediction turns out to be beaten too often.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from typing import Any, Deque, Dict
+
+#: Estimating from fewer samples than this is guessing with extra steps.
+DEFAULT_MIN_SAMPLES = 20
+
+#: How many recent completions per model take part. Long enough to be stable across
+#: a mix of short and long turns, short enough to follow a change in how a client
+#: is being used.
+DEFAULT_WINDOW = 200
+
+#: Never predict below this. A model that has only ever answered in one word must
+#: not leave a two-word answer under-reserved.
+DEFAULT_FLOOR = 256
+
+#: A p95 predictor is beaten about 5% of the time by construction. This is the
+#: safety valve for a distribution that has changed shape, not for that.
+DEFAULT_MAX_UNDERSHOOT = 0.20
+
+MODE_ADAPTIVE = "adaptive"
+MODE_CEILING = "max_tokens"
+
+
+class OutputEstimator:
+    """Per-model output-size prediction, bounded by the caller's own `max_tokens`.
+
+    Shared across the pool rather than held per account: what a model returns is a
+    property of the model and the traffic, not of the credential it was billed to,
+    and splitting the samples per account would just make each one slower to learn
+    the same number.
+    """
+
+    def __init__(
+        self,
+        mode: str = MODE_ADAPTIVE,
+        quantile: float = 0.95,
+        min_samples: int = DEFAULT_MIN_SAMPLES,
+        window: int = DEFAULT_WINDOW,
+        floor: int = DEFAULT_FLOOR,
+        max_undershoot: float = DEFAULT_MAX_UNDERSHOOT,
+    ) -> None:
+        self.mode = mode
+        self.quantile = min(1.0, max(0.5, float(quantile)))
+        self.min_samples = max(1, int(min_samples))
+        self.window = max(self.min_samples, int(window))
+        self.floor = max(1, int(floor))
+        self.max_undershoot = max(0.0, min(1.0, float(max_undershoot)))
+        self._samples: Dict[str, Deque[int]] = defaultdict(
+            lambda: deque(maxlen=self.window)
+        )
+        self._undershoots: Dict[str, Deque[bool]] = defaultdict(
+            lambda: deque(maxlen=self.window)
+        )
+
+    @property
+    def enabled(self) -> bool:
+        return self.mode == MODE_ADAPTIVE
+
+    def reconfigure(
+        self,
+        mode: str,
+        quantile: float,
+        min_samples: int,
+        window: int,
+        floor: int,
+        max_undershoot: float,
+    ) -> None:
+        """Apply new settings without forgetting what has been learned.
+
+        A config reload must not send the pool back to leasing ceilings for the next
+        twenty requests of every model. The samples are observed state, not
+        configuration, and they survive exactly as managed keys do.
+        """
+        self.mode = mode
+        self.quantile = min(1.0, max(0.5, float(quantile)))
+        self.min_samples = max(1, int(min_samples))
+        self.floor = max(1, int(floor))
+        self.max_undershoot = max(0.0, min(1.0, float(max_undershoot)))
+        new_window = max(self.min_samples, int(window))
+        if new_window != self.window:
+            self.window = new_window
+            for model, sizes in list(self._samples.items()):
+                self._samples[model] = deque(sizes, maxlen=new_window)
+            for model, missed in list(self._undershoots.items()):
+                self._undershoots[model] = deque(missed, maxlen=new_window)
+
+    def predict(self, model: str, ceiling: int) -> int:
+        """What to lease for this request. Never more than `ceiling`.
+
+        Returning `ceiling` is always a valid answer and is what happens whenever
+        the estimator has no business having an opinion.
+        """
+        ceiling = max(1, int(ceiling))
+        if not self.enabled:
+            return ceiling
+        samples = self._samples.get(model)
+        if samples is None or len(samples) < self.min_samples:
+            return ceiling
+        if self._undershoot_rate(model) > self.max_undershoot:
+            # The distribution moved. Stop predicting until the record of being
+            # beaten ages out of the window.
+            return ceiling
+        predicted = _quantile(samples, self.quantile)
+        return max(self.floor, min(ceiling, predicted))
+
+    def observe(self, model: str, predicted: int, actual: int) -> None:
+        """Record one completion: what it returned, and whether it beat the lease."""
+        if actual <= 0:
+            # No usage came back — an error, or an upstream that reported none.
+            # Neither says anything about how long an answer is.
+            return
+        self._samples[model].append(int(actual))
+        self._undershoots[model].append(int(actual) > int(predicted))
+
+    def _undershoot_rate(self, model: str) -> float:
+        seen = self._undershoots.get(model)
+        if not seen or len(seen) < self.min_samples:
+            return 0.0
+        return sum(1 for missed in seen if missed) / float(len(seen))
+
+    def snapshot(self) -> Dict[str, Any]:
+        """What the estimator currently believes, for the admin API and the console."""
+        models: Dict[str, Any] = {}
+        for model, samples in self._samples.items():
+            if not samples:
+                continue
+            ordered = sorted(samples)
+            models[model] = {
+                "samples": len(ordered),
+                "predicting": len(ordered) >= self.min_samples
+                and self._undershoot_rate(model) <= self.max_undershoot,
+                "median": ordered[len(ordered) // 2],
+                "p95": _quantile(samples, 0.95),
+                "max": ordered[-1],
+                "undershoot_rate": round(self._undershoot_rate(model), 4),
+            }
+        return {
+            "mode": self.mode,
+            "quantile": self.quantile,
+            "min_samples": self.min_samples,
+            "floor": self.floor,
+            "models": models,
+        }
+
+
+def _quantile(values: Any, fraction: float) -> int:
+    """Nearest-rank quantile. No numpy, and no interpolation to argue about."""
+    ordered = sorted(values)
+    if not ordered:
+        return 0
+    index = min(len(ordered) - 1, int(round(fraction * (len(ordered) - 1))))
+    return int(ordered[index])