tokenbiryani 0.2.0__py3-none-any.whl

tokenbiryani/proxy/sse.py

@@ -0,0 +1,98 @@
+"""Just enough SSE to account for a streamed response without altering it.
+
+The gateway never rewrites the stream. It watches it go past, so it can tell when
+the first content byte has been committed to the client and what the response cost.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, Iterator, Optional
+
+#: The event that marks the point of no return: after this, the client has content
+#: on screen and a failure can no longer be hidden by retrying elsewhere.
+FIRST_CONTENT_EVENT = b"content_block_delta"
+
+
+def iter_data_objects(text: str) -> Iterator[Dict[str, Any]]:
+    for line in text.splitlines():
+        line = line.strip()
+        if not line.startswith("data:"):
+            continue
+        payload = line[5:].strip()
+        if not payload or payload == "[DONE]":
+            continue
+        try:
+            parsed = json.loads(payload)
+        except ValueError:
+            continue
+        if isinstance(parsed, dict):
+            yield parsed
+
+
+class UsageCollector:
+    """Accumulates token usage from a Messages stream as it passes through."""
+
+    def __init__(self) -> None:
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self.cache_read_tokens = 0
+        self.cache_creation_tokens = 0
+        self.stop_reason: Optional[str] = None
+        self._buffer = ""
+
+    def feed(self, chunk: bytes) -> None:
+        try:
+            self._buffer += chunk.decode("utf-8")
+        except UnicodeDecodeError:
+            self._buffer += chunk.decode("utf-8", "ignore")
+        # Only whole lines are parseable; keep the tail for the next chunk.
+        head, newline, tail = self._buffer.rpartition("\n")
+        if not newline:
+            return
+        self._buffer = tail
+        self._absorb(head)
+
+    def finish(self) -> None:
+        if self._buffer:
+            self._absorb(self._buffer)
+            self._buffer = ""
+
+    def _absorb(self, text: str) -> None:
+        for obj in iter_data_objects(text):
+            kind = obj.get("type")
+            if kind == "message_start":
+                usage = (obj.get("message") or {}).get("usage") or {}
+                self.input_tokens += _int(usage.get("input_tokens"))
+                self.output_tokens += _int(usage.get("output_tokens"))
+                self.cache_read_tokens += _int(usage.get("cache_read_input_tokens"))
+                self.cache_creation_tokens += _int(usage.get("cache_creation_input_tokens"))
+            elif kind == "message_delta":
+                usage = obj.get("usage") or {}
+                self.output_tokens += _int(usage.get("output_tokens"))
+                delta = obj.get("delta") or {}
+                if delta.get("stop_reason"):
+                    self.stop_reason = str(delta["stop_reason"])
+
+    def as_body(self) -> Dict[str, Any]:
+        return {
+            "usage": {
+                "input_tokens": self.input_tokens,
+                "output_tokens": self.output_tokens,
+                "cache_read_input_tokens": self.cache_read_tokens,
+                "cache_creation_input_tokens": self.cache_creation_tokens,
+            }
+        }
+
+
+def error_event(message: str, kind: str = "api_error") -> bytes:
+    """An SSE error frame, for a failure that lands after the stream is committed."""
+    payload = json.dumps({"type": "error", "error": {"type": kind, "message": message}})
+    return ("event: error\ndata: " + payload + "\n\n").encode("utf-8")
+
+
+def _int(value: Any) -> int:
+    try:
+        return int(value or 0)
+    except (TypeError, ValueError):
+        return 0

tokenbiryani/store/base.py

@@ -0,0 +1,150 @@
+"""State store interface.
+
+Everything shared between requests goes through here — affinity, windowed spend,
+per-key request counts. Keeping it behind one interface is what makes the SQLite and
+Redis backends implementations rather than rewrites.
+
+Spend is a *ledger*, not a running total: entries carry a timestamp and callers ask
+for a window. A lifetime total on a persistent store would eventually wedge the
+gateway shut and stay that way.
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import Dict, List, Optional
+
+#: Spend scopes. Keys are billed to whoever presented them; accounts to the credential.
+SCOPE_KEY = "key"
+SCOPE_ACCOUNT = "account"
+
+#: One conversation's cost. A key cap catches a tenant overspending; nothing caught
+#: a single agent loop doing it, and one loop resending a 100k-token prefix a few
+#: hundred times is a whole week of quota with one virtual key's name on it.
+SCOPE_SESSION = "session"
+
+#: Turn *count* per conversation, carried on the same ledger with an amount of 1.0
+#: rather than a table of its own. The ledger is already (scope, name, at, amount)
+#: with a window query over it, which is exactly what counting turns in a window
+#: needs — and it means turn caps hold across instances on Redis for free, where a
+#: counter kept in the process would not.
+SCOPE_SESSION_TURNS = "session_turns"
+
+
+class StateStore(abc.ABC):
+    @abc.abstractmethod
+    async def get_affinity(self, session_key: str) -> Optional[str]:
+        """Which account currently owns this session's prompt cache, if any."""
+
+    @abc.abstractmethod
+    async def set_affinity(self, session_key: str, account_id: str, ttl: float) -> None:
+        ...
+
+    @abc.abstractmethod
+    async def clear_affinity(self, session_key: str) -> None:
+        ...
+
+    @abc.abstractmethod
+    async def clear_affinity_for_account(self, account_id: str) -> int:
+        """Drop every session pinned to this account; return how many. 
+
+        Called when an account is deleted. Without it those sessions keep naming a
+        credential that no longer exists for the rest of the affinity TTL: the
+        router cannot match the owner, so it scores every candidate at zero
+        affinity and reports a cache break on each one — noise attributed to
+        routing for something routing did not do.
+
+        Not called on *disable*, which is usually temporary. A disabled account
+        that comes back should find its conversations still pinned to it.
+        """
+
+    @abc.abstractmethod
+    async def record_key_request(self, key_name: str, now: float) -> int:
+        """Record a request against a virtual key; return its count in the last 60s."""
+
+    @abc.abstractmethod
+    async def add_spend(self, scope: str, name: str, amount: float) -> None:
+        """Append to the spend ledger."""
+
+    @abc.abstractmethod
+    async def get_spend(self, scope: str, name: str, window_seconds: float) -> float:
+        """Total spend for one name over the trailing window."""
+
+    @abc.abstractmethod
+    async def spend_by_scope(self, scope: str, window_seconds: float) -> Dict[str, float]:
+        """Every name in a scope, for hydrating in-memory counters at startup."""
+
+    @abc.abstractmethod
+    async def record_usage(self, sample: Dict[str, object]) -> None:
+        """Append one request's accounting to the usage history.
+
+        Telemetry, not billing — `add_spend` remains the record a cap is enforced
+        against. Losing a row here costs a notch on a chart, nothing more.
+        """
+
+    @abc.abstractmethod
+    async def usage_rows(
+        self, since: float, until: float, account_id: Optional[str] = None
+    ) -> List[Dict[str, object]]:
+        """Raw usage rows in a window. Bucketing happens in `observability.usage`."""
+
+    @abc.abstractmethod
+    async def put_key(self, record: Dict[str, object]) -> None:
+        """Store a managed key record. Records hold a hash, never the key itself."""
+
+    @abc.abstractmethod
+    async def delete_key(self, name: str) -> bool:
+        ...
+
+    @abc.abstractmethod
+    async def list_keys(self) -> List[Dict[str, object]]:
+        ...
+
+    @abc.abstractmethod
+    async def put_account(self, record: Dict[str, object]) -> None:
+        """Store an account added through the API. Credentials arrive encrypted."""
+
+    @abc.abstractmethod
+    async def delete_account(self, account_id: str) -> bool:
+        ...
+
+    @abc.abstractmethod
+    async def list_accounts(self) -> List[Dict[str, object]]:
+        ...
+
+    async def put_setting(self, name: str, value: object) -> None:
+        """Persist one operator setting made from the console.
+
+        Not abstract: a store predating this method keeps working, it just forgets
+        the setting at restart, which is a degradation and not a crash.
+        """
+        return None
+
+    async def get_settings(self) -> Dict[str, object]:
+        return {}
+
+    async def startup(self) -> None:
+        return None
+
+    async def close(self) -> None:
+        return None
+
+
+def build_store(config) -> StateStore:
+    """Construct the store named by `store.backend`."""
+    backend = (config.store.backend or "memory").lower()
+    if backend == "memory":
+        from .memory import MemoryStateStore
+
+        return MemoryStateStore()
+    if backend == "sqlite":
+        from .sqlite import SqliteStateStore
+
+        return SqliteStateStore(config.store.path)
+    if backend == "redis":
+        from .redis_store import RedisStateStore
+
+        return RedisStateStore(config.store.url, config.store.namespace)
+    raise ValueError(
+        f"unknown store backend {backend!r}; known: memory, sqlite, redis"
+    )

tokenbiryani/store/memory.py

@@ -0,0 +1,149 @@
+"""In-process state store. Zero dependencies; loses everything on restart."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import defaultdict, deque
+from typing import Deque, Dict, List, Optional, Tuple
+
+from .base import StateStore
+
+#: How many usage rows this store keeps. Roughly a day of steady traffic, and a
+#: few megabytes. The persistent backends keep 90 days instead.
+USAGE_CAPACITY = 50_000
+
+
+class MemoryStateStore(StateStore):
+    def __init__(self) -> None:
+        self._affinity: Dict[str, Tuple[str, float]] = {}
+        self._requests: Dict[str, Deque[float]] = defaultdict(deque)
+        self._ledger: Dict[str, List[Tuple[float, float]]] = defaultdict(list)
+        self._keys: Dict[str, Dict[str, object]] = {}
+        self._accounts: Dict[str, Dict[str, object]] = {}
+        self._settings: Dict[str, object] = {}
+        self._usage: List[Dict[str, object]] = []
+        self._lock = asyncio.Lock()
+
+    async def get_affinity(self, session_key: str) -> Optional[str]:
+        async with self._lock:
+            entry = self._affinity.get(session_key)
+            if entry is None:
+                return None
+            account_id, expires_at = entry
+            if expires_at and expires_at < time.time():
+                del self._affinity[session_key]
+                return None
+            return account_id
+
+    async def set_affinity(self, session_key: str, account_id: str, ttl: float) -> None:
+        async with self._lock:
+            self._affinity[session_key] = (account_id, time.time() + ttl)
+
+    async def clear_affinity(self, session_key: str) -> None:
+        async with self._lock:
+            self._affinity.pop(session_key, None)
+
+    async def clear_affinity_for_account(self, account_id: str) -> int:
+        async with self._lock:
+            doomed = [
+                session
+                for session, (owner, _) in self._affinity.items()
+                if owner == account_id
+            ]
+            for session in doomed:
+                del self._affinity[session]
+            return len(doomed)
+
+    async def record_key_request(self, key_name: str, now: float) -> int:
+        async with self._lock:
+            bucket = self._requests[key_name]
+            bucket.append(now)
+            cutoff = now - 60.0
+            while bucket and bucket[0] < cutoff:
+                bucket.popleft()
+            return len(bucket)
+
+    async def add_spend(self, scope: str, name: str, amount: float) -> None:
+        async with self._lock:
+            self._ledger[f"{scope}:{name}"].append((time.time(), amount))
+
+    async def get_spend(self, scope: str, name: str, window_seconds: float) -> float:
+        async with self._lock:
+            return self._sum(f"{scope}:{name}", window_seconds)
+
+    async def spend_by_scope(self, scope: str, window_seconds: float) -> Dict[str, float]:
+        prefix = scope + ":"
+        async with self._lock:
+            return {
+                key[len(prefix):]: self._sum(key, window_seconds)
+                for key in list(self._ledger)
+                if key.startswith(prefix)
+            }
+
+    async def record_usage(self, sample: Dict[str, object]) -> None:
+        async with self._lock:
+            self._usage.append(dict(sample))
+            # Bounded, unlike the persistent backends: this store is already
+            # "everything disappears on restart", so the ceiling is memory, not days.
+            if len(self._usage) > USAGE_CAPACITY:
+                del self._usage[: len(self._usage) - USAGE_CAPACITY]
+
+    async def usage_rows(
+        self, since: float, until: float, account_id: Optional[str] = None
+    ) -> List[Dict[str, object]]:
+        async with self._lock:
+            return [
+                dict(row)
+                for row in self._usage
+                if since <= _at(row) < until
+                and (account_id is None or row.get("account_id") == account_id)
+            ]
+
+    async def put_key(self, record: Dict[str, object]) -> None:
+        async with self._lock:
+            self._keys[str(record["name"])] = dict(record)
+
+    async def delete_key(self, name: str) -> bool:
+        async with self._lock:
+            return self._keys.pop(name, None) is not None
+
+    async def list_keys(self) -> List[Dict[str, object]]:
+        async with self._lock:
+            return [dict(record) for record in self._keys.values()]
+
+    async def put_account(self, record: Dict[str, object]) -> None:
+        async with self._lock:
+            self._accounts[str(record["id"])] = dict(record)
+
+    async def delete_account(self, account_id: str) -> bool:
+        async with self._lock:
+            return self._accounts.pop(account_id, None) is not None
+
+    async def list_accounts(self) -> List[Dict[str, object]]:
+        async with self._lock:
+            return [dict(record) for record in self._accounts.values()]
+
+    async def put_setting(self, name: str, value: object) -> None:
+        async with self._lock:
+            self._settings[str(name)] = value
+
+    async def get_settings(self) -> Dict[str, object]:
+        async with self._lock:
+            return dict(self._settings)
+
+    def _sum(self, key: str, window_seconds: float) -> float:
+        cutoff = time.time() - window_seconds
+        entries = self._ledger.get(key)
+        if not entries:
+            return 0.0
+        # Prune while we are here; the ledger is append-only otherwise.
+        kept = [entry for entry in entries if entry[0] >= cutoff]
+        self._ledger[key] = kept
+        return sum(amount for _, amount in kept)
+
+
+def _at(row: Dict[str, object]) -> float:
+    """Rows are `Dict[str, object]` at this boundary; read the timestamp safely."""
+    value = row.get("at")
+    return float(value) if isinstance(value, (int, float)) else 0.0

tokenbiryani/store/redis_store.py

@@ -0,0 +1,222 @@
+"""Redis state store: shared affinity, spend and rate counters across instances.
+
+This is what makes more than one gateway process safe to run. Without it, two
+instances each keep their own affinity map and each enforce half a spend cap.
+
+Named `redis_store` rather than `redis` so nothing in this package can shadow the
+`redis` distribution for a reader skimming imports.
+
+Requires the optional dependency:  pip install "tokenbiryani[redis]"
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+#: Ledger entries are stored as "<amount>|<nonce>" so two identical charges in the
+#: same millisecond remain distinct members of the sorted set. Usage rows use the
+#: same trick, as "<json>|<nonce>".
+_SEPARATOR = "|"
+
+#: Matches the SQLite backend's retention, so a chart looks the same on either.
+USAGE_RETENTION_SECONDS = 90 * 24 * 3600
+
+
+def _connect(url: str) -> Any:
+    try:
+        import redis.asyncio as redis_asyncio
+    except ImportError as exc:  # pragma: no cover - exercised by the error path only
+        raise RuntimeError(
+            "the redis store needs the redis package: pip install 'tokenbiryani[redis]'"
+        ) from exc
+    return redis_asyncio.from_url(url, decode_responses=True)
+
+
+class RedisStateStore(StateStore):
+    def __init__(
+        self,
+        url: str = "redis://127.0.0.1:6379/0",
+        namespace: str = "tokenbiryani",
+        client: Optional[Any] = None,
+    ) -> None:
+        self.url = url
+        self.namespace = namespace
+        self._client = client
+        self._owns_client = client is None
+
+    @property
+    def client(self) -> Any:
+        if self._client is None:
+            self._client = _connect(self.url)
+        return self._client
+
+    def _key(self, *parts: str) -> str:
+        return ":".join((self.namespace,) + parts)
+
+    async def startup(self) -> None:
+        await self.client.ping()
+
+    async def close(self) -> None:
+        if self._client is not None and self._owns_client:
+            closer = getattr(self._client, "aclose", None) or getattr(
+                self._client, "close", None
+            )
+            if closer is not None:
+                await closer()
+            self._client = None
+
+    # ---- affinity ------------------------------------------------------------
+
+    async def get_affinity(self, session_key: str) -> Optional[str]:
+        return await self.client.get(self._key("affinity", session_key))
+
+    async def set_affinity(self, session_key: str, account_id: str, ttl: float) -> None:
+        # px, not ex: a sub-second TTL must not round down to "no expiry".
+        await self.client.set(
+            self._key("affinity", session_key), account_id, px=max(1, int(ttl * 1000))
+        )
+
+    async def clear_affinity(self, session_key: str) -> None:
+        await self.client.delete(self._key("affinity", session_key))
+
+    async def clear_affinity_for_account(self, account_id: str) -> int:
+        """Scan the affinity keyspace and drop the ones this account owns.
+
+        A scan rather than a reverse index: deleting an account is rare and
+        operator-driven, while `set_affinity` runs on every successful request.
+        Paying a little on the rare path keeps the hot one a single SET.
+        """
+        pattern = self._key("affinity", "*")
+        removed = 0
+        async for key in self.client.scan_iter(match=pattern, count=500):
+            if await self.client.get(key) == account_id:
+                removed += int(await self.client.delete(key) or 0)
+        return removed
+
+    # ---- per-key request rate ------------------------------------------------
+
+    async def record_key_request(self, key_name: str, now: float) -> int:
+        key = self._key("requests", key_name)
+        await self.client.zadd(key, {uuid.uuid4().hex: now})
+        # "(" makes the upper bound exclusive. Redis is inclusive by default, which
+        # would evict one entry more than the memory and SQLite backends and quietly
+        # under-count the rate. The window keeps entries with score >= cutoff.
+        await self.client.zremrangebyscore(key, "-inf", f"({now - 60.0}")
+        # Expire the whole set if the key goes quiet, so idle keys cost nothing.
+        await self.client.expire(key, 120)
+        return int(await self.client.zcard(key))
+
+    # ---- spend ledger --------------------------------------------------------
+
+    async def add_spend(self, scope: str, name: str, amount: float) -> None:
+        key = self._key("spend", scope, name)
+        member = f"{amount}{_SEPARATOR}{uuid.uuid4().hex}"
+        await self.client.zadd(key, {member: time.time()})
+        await self.client.sadd(self._key("spend-names", scope), name)
+
+    async def get_spend(self, scope: str, name: str, window_seconds: float) -> float:
+        key = self._key("spend", scope, name)
+        cutoff = time.time() - window_seconds
+        await self.client.zremrangebyscore(key, "-inf", f"({cutoff}")
+        members = await self.client.zrangebyscore(key, cutoff, "+inf")
+        return sum(_amount(member) for member in members)
+
+    # ---- usage history -------------------------------------------------------
+
+    async def record_usage(self, sample: Dict[str, Any]) -> None:
+        # A sorted set scored by timestamp, exactly like the spend ledger. The
+        # nonce keeps two identical requests in the same instant distinct members.
+        member = json.dumps(sample, separators=(",", ":")) + _SEPARATOR + uuid.uuid4().hex
+        await self.client.zadd(self._key("usage"), {member: float(sample.get("at") or 0.0)})
+
+    async def usage_rows(
+        self, since: float, until: float, account_id: Optional[str] = None
+    ) -> List[Dict[str, Any]]:
+        key = self._key("usage")
+        # Trim on read rather than on write: the write is on the request path.
+        await self.client.zremrangebyscore(
+            key, "-inf", f"({time.time() - USAGE_RETENTION_SECONDS}"
+        )
+        # "(" makes the upper bound exclusive so a row exactly on `until` belongs to
+        # the next window, matching the other two backends.
+        members = await self.client.zrangebyscore(key, since, f"({until}")
+        rows: List[Dict[str, Any]] = []
+        for member in members:
+            try:
+                row = json.loads(member.rsplit(_SEPARATOR, 1)[0])
+            except ValueError:
+                continue
+            if account_id is None or row.get("account_id") == account_id:
+                rows.append(row)
+        return rows
+
+    # ---- managed keys --------------------------------------------------------
+
+    async def put_key(self, record: Dict[str, Any]) -> None:
+        await self.client.hset(
+            self._key("keys"), str(record["name"]), json.dumps(record)
+        )
+
+    async def delete_key(self, name: str) -> bool:
+        return bool(await self.client.hdel(self._key("keys"), name))
+
+    async def list_keys(self) -> List[Dict[str, Any]]:
+        raw = await self.client.hgetall(self._key("keys"))
+        records = []
+        for value in (raw or {}).values():
+            try:
+                records.append(json.loads(value))
+            except ValueError:
+                continue
+        return records
+
+    async def put_account(self, record: Dict[str, Any]) -> None:
+        await self.client.hset(self._key("accounts"), str(record["id"]), json.dumps(record))
+
+    async def delete_account(self, account_id: str) -> bool:
+        return bool(await self.client.hdel(self._key("accounts"), account_id))
+
+    async def list_accounts(self) -> List[Dict[str, Any]]:
+        raw = await self.client.hgetall(self._key("accounts"))
+        records = []
+        for value in (raw or {}).values():
+            try:
+                records.append(json.loads(value))
+            except ValueError:
+                continue
+        return sorted(records, key=lambda r: str(r.get("id", "")))
+
+    async def put_setting(self, name: str, value: Any) -> None:
+        await self.client.hset(self._key("settings"), str(name), json.dumps(value))
+
+    async def get_settings(self) -> Dict[str, Any]:
+        raw = await self.client.hgetall(self._key("settings"))
+        out: Dict[str, Any] = {}
+        for name, value in (raw or {}).items():
+            try:
+                out[str(name)] = json.loads(value)
+            except ValueError:
+                continue
+        return out
+
+    async def spend_by_scope(self, scope: str, window_seconds: float) -> Dict[str, float]:
+        names = await self.client.smembers(self._key("spend-names", scope))
+        totals: Dict[str, float] = {}
+        for name in names:
+            total = await self.get_spend(scope, name, window_seconds)
+            if total:
+                totals[name] = total
+        return totals
+
+
+def _amount(member: str) -> float:
+    head = member.split(_SEPARATOR, 1)[0]
+    try:
+        return float(head)
+    except ValueError:
+        return 0.0