tokenbiryani 0.2.0__py3-none-any.whl

tokenbiryani/providers/base.py

@@ -0,0 +1,173 @@
+"""Upstream protocol.
+
+An upstream is anything that can serve a Messages request and report its limits.
+Adding Bedrock or Vertex means implementing this and nothing else.
+"""
+
+from __future__ import annotations
+
+import abc
+from dataclasses import dataclass, field
+from typing import Any, Dict, Mapping, Optional
+
+import httpx
+
+#: Headers we never forward from the client — ours to set, or hop-by-hop.
+BLOCKED_REQUEST_HEADERS = frozenset(
+    {
+        "host", "content-length", "connection", "keep-alive", "transfer-encoding",
+        "upgrade", "te", "trailer", "proxy-authorization", "proxy-authenticate",
+        "authorization", "x-api-key", "accept-encoding",
+    }
+)
+
+#: Headers we never pass back to the client verbatim.
+BLOCKED_RESPONSE_HEADERS = frozenset(
+    {"content-length", "content-encoding", "connection", "transfer-encoding", "keep-alive"}
+)
+
+
+@dataclass
+class UpstreamResult:
+    status: int
+    headers: Dict[str, str] = field(default_factory=dict)
+    body: Optional[Dict[str, Any]] = None
+    raw: bytes = b""
+
+
+def forwardable(headers: Mapping[str, str]) -> Dict[str, str]:
+    """Pass the client's request headers through, minus the ones we own.
+
+    The gateway is a router, not a rewriter: anything the client sent that isn't
+    ours to control travels untouched, including beta flags we've never heard of.
+    """
+    return {k: v for k, v in headers.items() if k.lower() not in BLOCKED_REQUEST_HEADERS}
+
+
+def _as_result(response: httpx.Response, parse_json: bool = True) -> UpstreamResult:
+    body = None
+    if parse_json:
+        try:
+            parsed = response.json()
+            body = parsed if isinstance(parsed, dict) else None
+        except ValueError:
+            body = None
+    return UpstreamResult(
+        status=response.status_code,
+        headers=dict(response.headers),
+        body=body,
+        raw=response.content,
+    )
+
+
+class Upstream(abc.ABC):
+    """One credential's worth of capacity."""
+
+    def __init__(self, account_id: str) -> None:
+        self.account_id = account_id
+
+    @abc.abstractmethod
+    def url(self, path: str) -> str:
+        ...
+
+    @abc.abstractmethod
+    def auth_headers(self) -> Dict[str, str]:
+        ...
+
+    def request_headers(self, client_headers: Mapping[str, str]) -> Dict[str, str]:
+        headers = forwardable(client_headers)
+        headers.update(self.auth_headers())
+        return headers
+
+    async def send(
+        self,
+        client: httpx.AsyncClient,
+        path: str,
+        payload: Mapping[str, Any],
+        client_headers: Mapping[str, str],
+        timeout: Optional[float] = None,
+    ) -> UpstreamResult:
+        response = await client.post(
+            self.url(path),
+            json=payload,
+            headers=self.request_headers(client_headers),
+            timeout=timeout,
+        )
+        body: Optional[Dict[str, Any]]
+        try:
+            parsed = response.json()
+            body = parsed if isinstance(parsed, dict) else None
+        except ValueError:
+            body = None
+        return UpstreamResult(
+            status=response.status_code,
+            headers=dict(response.headers),
+            body=body,
+            raw=response.content,
+        )
+
+    # ---- Message Batches ----------------------------------------------------
+    #
+    # The spill lane. Same credential, separate rate-limit pool upstream, so batch
+    # traffic does not eat the interactive budget the router is protecting.
+
+    async def submit_batch(
+        self,
+        client: httpx.AsyncClient,
+        requests: Any,
+        client_headers: Mapping[str, str],
+    ) -> UpstreamResult:
+        return await self.send(
+            client, "/v1/messages/batches", {"requests": requests}, client_headers
+        )
+
+    async def poll_batch(
+        self, client: httpx.AsyncClient, batch_id: str, client_headers: Mapping[str, str]
+    ) -> UpstreamResult:
+        response = await client.get(
+            self.url("/v1/messages/batches/" + batch_id),
+            headers=self.request_headers(client_headers),
+        )
+        return _as_result(response)
+
+    async def fetch_batch_results(
+        self, client: httpx.AsyncClient, batch_id: str, client_headers: Mapping[str, str]
+    ) -> UpstreamResult:
+        response = await client.get(
+            self.url("/v1/messages/batches/" + batch_id + "/results"),
+            headers=self.request_headers(client_headers),
+        )
+        return _as_result(response, parse_json=False)
+
+    async def cancel_batch(
+        self, client: httpx.AsyncClient, batch_id: str, client_headers: Mapping[str, str]
+    ) -> UpstreamResult:
+        return await self.send(
+            client, "/v1/messages/batches/" + batch_id + "/cancel", {}, client_headers
+        )
+
+    def iter_sse(self, response: httpx.Response):
+        """Yield the response as Anthropic SSE bytes.
+
+        The default is a straight passthrough, because that is what the request
+        already is. Bedrock is the exception: it frames its stream in AWS
+        event-stream binary, so its adapter re-emits SSE here.
+        """
+        return response.aiter_bytes()
+
+    def open_stream(
+        self,
+        client: httpx.AsyncClient,
+        path: str,
+        payload: Mapping[str, Any],
+        client_headers: Mapping[str, str],
+        timeout: Optional[float] = None,
+    ):
+        """Return an httpx streaming context manager; the caller drives it."""
+        return client.stream(
+            "POST",
+            self.url(path),
+            json=payload,
+            headers=self.request_headers(client_headers),
+            timeout=timeout,
+        )

tokenbiryani/providers/bedrock.py

@@ -0,0 +1,182 @@
+"""AWS Bedrock upstream.
+
+Two things make Bedrock the awkward one. It addresses the model in the URL and signs
+every request with SigV4, so the request cannot simply be forwarded. And it frames its
+stream in AWS event-stream binary rather than SSE, so `iter_sse` decodes each frame
+and re-emits the Anthropic event the client is expecting.
+
+    pip install "tokenbiryani[bedrock]"
+
+    accounts:
+      - id: acct-bedrock
+        type: bedrock
+        options:
+          region: us-east-1
+          profile: default          # optional; omit for the ambient credential chain
+          model_map:
+            claude-test-1: anthropic.claude-3-5-sonnet-20241022-v2:0
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from typing import Any, Dict, Mapping, Optional
+
+import httpx
+
+from ..config import AccountConfig
+from .base import Upstream, UpstreamResult, _as_result
+from .translate import BEDROCK_VERSION, map_model, to_platform_body
+
+SERVICE = "bedrock"
+
+
+class BedrockSigner:
+    """SigV4 over botocore. Signing needs the body, so it happens per request."""
+
+    def __init__(
+        self, region: str, profile: Optional[str] = None, credentials: Optional[Any] = None
+    ) -> None:
+        self.region = region
+        self.profile = profile
+        self._credentials = credentials
+
+    def credentials(self) -> Any:
+        if self._credentials is None:
+            try:
+                from botocore.session import Session
+            except ImportError as exc:  # pragma: no cover - error path only
+                raise RuntimeError(
+                    "the bedrock adapter needs botocore: "
+                    "pip install 'tokenbiryani[bedrock]'"
+                ) from exc
+            session = Session(profile=self.profile) if self.profile else Session()
+            resolved = session.get_credentials()
+            if resolved is None:
+                raise RuntimeError(
+                    "no AWS credentials found for the bedrock adapter"
+                )
+            self._credentials = resolved.get_frozen_credentials()
+        return self._credentials
+
+    def sign(self, url: str, body: bytes, headers: Mapping[str, str]) -> Dict[str, str]:
+        from botocore.auth import SigV4Auth
+        from botocore.awsrequest import AWSRequest
+
+        request = AWSRequest(method="POST", url=url, data=body, headers=dict(headers))
+        SigV4Auth(self.credentials(), SERVICE, self.region).add_auth(request)
+        return dict(request.headers)
+
+
+class BedrockUpstream(Upstream):
+    def __init__(self, config: AccountConfig, signer: Optional[BedrockSigner] = None) -> None:
+        super().__init__(config.id)
+        self.config = config
+        options = config.options or {}
+        self.region = str(options.get("region") or "us-east-1")
+        self.model_map: Dict[str, str] = dict(options.get("model_map") or {})
+        self.base_url = (config.base_url or "").rstrip("/") or (
+            f"https://bedrock-runtime.{self.region}.amazonaws.com"
+        )
+        self._signer = signer or BedrockSigner(self.region, options.get("profile"))
+
+    # ---- addressing ----------------------------------------------------------
+
+    def model_url(self, model: str, stream: bool) -> str:
+        verb = "invoke-with-response-stream" if stream else "invoke"
+        return f"{self.base_url}/model/{map_model(model, self.model_map)}/{verb}"
+
+    def url(self, path: str) -> str:
+        return self.base_url + "/" + path.lstrip("/")
+
+    def auth_headers(self) -> Dict[str, str]:
+        # Signing produces the auth headers; nothing static to add.
+        return dict(self.config.headers)
+
+    def _headers(self, client_headers: Mapping[str, str], url: str, body: bytes):
+        headers = super().request_headers(client_headers)
+        headers.pop("anthropic-version", None)
+        headers["content-type"] = "application/json"
+        headers["host"] = httpx.URL(url).host
+        return self._signer.sign(url, body, headers)
+
+    # ---- transport -----------------------------------------------------------
+
+    async def send(
+        self,
+        client: httpx.AsyncClient,
+        path: str,
+        payload: Mapping[str, Any],
+        client_headers: Mapping[str, str],
+        timeout: Optional[float] = None,
+    ) -> UpstreamResult:
+        url = self.model_url(str(payload.get("model") or ""), stream=False)
+        body = json.dumps(to_platform_body(payload, BEDROCK_VERSION)).encode("utf-8")
+        response = await client.post(
+            url, content=body, headers=self._headers(client_headers, url, body), timeout=timeout
+        )
+        return _as_result(response)
+
+    def open_stream(
+        self,
+        client: httpx.AsyncClient,
+        path: str,
+        payload: Mapping[str, Any],
+        client_headers: Mapping[str, str],
+        timeout: Optional[float] = None,
+    ):
+        url = self.model_url(str(payload.get("model") or ""), stream=True)
+        body = json.dumps(to_platform_body(payload, BEDROCK_VERSION)).encode("utf-8")
+        return client.stream(
+            "POST",
+            url,
+            content=body,
+            headers=self._headers(client_headers, url, body),
+            timeout=timeout,
+        )
+
+    async def iter_sse(self, response: httpx.Response):
+        """Decode AWS event-stream frames and re-emit them as Anthropic SSE.
+
+        Each frame's payload is `{"bytes": "<base64 of one Anthropic event>"}`. The
+        gateway's stream machinery — the commit boundary, usage accounting — reads
+        ordinary SSE, so the translation stops here.
+        """
+        buffer = _event_stream_buffer()
+        async for chunk in response.aiter_bytes():
+            buffer.add_data(chunk)
+            for event in buffer:
+                frame = _to_sse(event)
+                if frame:
+                    yield frame
+
+
+def _event_stream_buffer():
+    try:
+        from botocore.eventstream import EventStreamBuffer
+    except ImportError as exc:  # pragma: no cover - error path only
+        raise RuntimeError(
+            "the bedrock adapter needs botocore: pip install 'tokenbiryani[bedrock]'"
+        ) from exc
+    return EventStreamBuffer()
+
+
+def _to_sse(event: Any) -> bytes:
+    payload = getattr(event, "payload", None)
+    if not payload:
+        return b""
+    try:
+        wrapper = json.loads(payload.decode("utf-8"))
+    except (ValueError, AttributeError):
+        return b""
+    encoded = wrapper.get("bytes")
+    if not encoded:
+        return b""
+    try:
+        inner = base64.b64decode(encoded)
+        parsed = json.loads(inner.decode("utf-8"))
+    except (ValueError, TypeError):
+        return b""
+    kind = parsed.get("type") or "message"
+    return f"event: {kind}\ndata: {json.dumps(parsed)}\n\n".encode()

tokenbiryani/providers/oauth.py

@@ -0,0 +1,165 @@
+"""`type: oauth` — an account backed by a Claude subscription session.
+
+The token normally comes from a login the gateway ran itself and holds per account,
+which is what pooling more than one subscription requires. It can also come from a
+credentials file or an environment variable — see `oauth_credentials` — for the case
+where something else already obtained it.
+
+**What an oauth account costs you:** a subscription session does not send the
+`anthropic-ratelimit-{requests,input-tokens,output-tokens}-*` triples an API key
+does, so `observable_limits` is forced false rather than left to the operator — an
+account whose windows never populate would read as permanently full and win every
+routing comparison.
+
+It does send `anthropic-ratelimit-unified-*`: how much of the rolling 5-hour and
+7-day windows is spent. That is a real headroom number and the mirror reads it, so
+these accounts do get headroom-aware routing. What stays dark is anything needing an
+absolute token count: leases cannot reserve against a percentage, and the capacity
+horizon cannot plot one. Failover and prompt-cache affinity are unaffected.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, Mapping, Optional
+
+from ..config import AccountConfig
+from .base import Upstream
+from .oauth_credentials import CredentialError, StaticToken, TokenSource, build_source
+
+#: Account options that name a token the gateway reads rather than owns. Any of
+#: these means "do not use a gateway-managed session for this account".
+EXTERNAL_SOURCE_OPTIONS = ("access_token", "token_env", "credentials_path")
+
+logger = logging.getLogger("tokenbiryani")
+
+DEFAULT_BASE_URL = "https://api.anthropic.com"
+
+#: Unverified, and deliberately configurable. This project has never been run against
+#: a real subscription session, so this is a starting value rather than a fact. If the
+#: gateway logs 4xx mentioning a beta flag, set `options.beta_header`.
+DEFAULT_BETA = "oauth-2025-04-20"
+
+ANTHROPIC_VERSION = "2023-06-01"
+
+
+class OAuthUpstream(Upstream):
+    """Bearer-token upstream. The token is supplied by a callback, not held here.
+
+    The gateway owns token lifecycle — storage, refresh, expiry — so this class asks
+    for the current one on each request rather than caching a copy that could go
+    stale behind a successful refresh.
+    """
+
+    def __init__(
+        self,
+        config: AccountConfig,
+        token_provider: Optional[Callable[[], str]] = None,
+    ) -> None:
+        super().__init__(config.id)
+        self.config = config
+        options: Dict[str, Any] = dict(config.options or {})
+        self.base_url = (config.base_url or DEFAULT_BASE_URL).rstrip("/")
+        self.beta = str(options.get("beta_header") or DEFAULT_BETA)
+
+        # Two ways an oauth account gets a token, and explicit config wins.
+        #
+        # `type: oauth` originally meant only "read the token the Claude CLI already
+        # wrote". Those configs keep working verbatim, so any account naming a source
+        # keeps using it; only accounts with no source configured use the session the
+        # gateway logged in for and refreshes.
+        self._source: Optional[TokenSource] = None
+        if any(options.get(k) for k in EXTERNAL_SOURCE_OPTIONS):
+            self._source = build_source(options)
+        elif config.api_key:
+            # A token pasted into the console — `claude setup-token` mints one that
+            # outlives a CLI session. It arrives in `api_key` rather than in options
+            # because that is the field the store encrypts and the admin API strips
+            # by name; a secret in `options` would be neither.
+            self._source = StaticToken(config.api_key)
+        self._token_provider = token_provider
+
+        if config.observable_limits:
+            # Not merely cosmetic: an account whose classic windows never populate
+            # reads as full, wins every comparison against accounts reporting an
+            # honest partial budget, and absorbs the pool's traffic until it 429s.
+            # Its real budget arrives as unified utilisation, which the mirror reads
+            # whatever this flag says.
+            logger.warning(
+                "account %r is an oauth account but observable_limits is true. "
+                "Subscription sessions send no per-window limit/remaining headers, so "
+                "its headroom would read as permanently full. Forcing it false — its "
+                "unified utilisation is read either way.",
+                config.id,
+            )
+            config.observable_limits = False
+
+    def token(self) -> str:
+        if self._source is not None:
+            # CredentialError propagates: the gateway classifies it as a transport
+            # failure, and the reason lands in the request inspector.
+            return self._source.token()
+        if self._token_provider is None:
+            raise RuntimeError(
+                f"account {self.account_id!r} has no token: log in from the console, "
+                "point it at this machine's Claude Code session, paste a token from "
+                "`claude setup-token`, or set options.credentials_path / "
+                "options.token_env / options.access_token"
+            )
+        return self._token_provider()
+
+    def describe_source(self) -> str:
+        if self._source is not None:
+            return self._source.describe()
+        return "session held by the gateway"
+
+    def source_expiry(self) -> Optional[float]:
+        """When the token this account reads goes stale, if the source knows.
+
+        A credentials file states its own expiry, and showing it is the difference
+        between "this will stop working at 21:56 unless the CLI renews it" and a
+        401 nobody expected.
+        """
+        getter = getattr(self._source, "expires_at", None)
+        if getter is None:
+            return None
+        try:
+            return getter()
+        except CredentialError:
+            return None
+
+    @property
+    def source_kind(self) -> str:
+        """Which of the four token sources this account is using, for the console."""
+        if isinstance(self._source, StaticToken):
+            return "token"
+        if self._source is not None:
+            return "external"
+        return "login"
+
+    def url(self, path: str) -> str:
+        return self.base_url + "/" + path.lstrip("/")
+
+    def auth_headers(self) -> Dict[str, str]:
+        headers = {
+            "authorization": "Bearer " + self.token(),
+            "content-type": "application/json",
+            "anthropic-beta": self.beta,
+            "anthropic-version": ANTHROPIC_VERSION,
+        }
+        headers.update(self.config.headers)
+        return headers
+
+    def request_headers(self, client_headers: Mapping[str, str]) -> Dict[str, str]:
+        headers = super().request_headers(client_headers)
+        # A subscription session authenticates with a bearer token; an x-api-key
+        # forwarded from the caller would be ignored at best and confusing at worst.
+        for name in [k for k in headers if k.lower() == "x-api-key"]:
+            headers.pop(name)
+        client_beta = next(
+            (v for k, v in client_headers.items() if k.lower() == "anthropic-beta"), None
+        )
+        if client_beta and self.beta not in client_beta:
+            # Keep the caller's beta flags and add ours.
+            headers["anthropic-beta"] = client_beta + "," + self.beta
+        return headers