src/services/router.py

"""Unified router: one retry loop for every model and provider.   # module docstring start

Each target is a plain dict {url, key, upstream_model}. The router picks one   # how it works
via round-robin, swaps the model field, forwards, and cools down bad targets.   # load distribution
No provider-specific branches — subnet nodes and API keys look identical.   # module docstring end
"""
from __future__ import annotations   # defer annotation evaluation (PEP 563)

import json   # stdlib json for parsing the request body to extract the model field
import logging   # structured logging for request/retry/error events
from typing import Optional   # Optional for the return types

from fastapi.responses import JSONResponse, Response, StreamingResponse   # response types

from ..models.config import ModelConfig   # the config that knows all models + targets
from .forwarder import SSEFilter, clean_response_headers, forward, make_client, swap_model   # forwarding helpers
from .health import HealthTracker   # per-target cooldown tracker
from .round_robin import RoundRobin   # thread-safe cursor for even distribution
from .billing.gate import extract_usage as gate_extract_usage   # parse usage from a response body
from .billing.usage_watcher import UsageWatcher   # SSE usage scanner for streaming metering

logger = logging.getLogger("router")   # named logger for this module

MAX_ATTEMPTS = 5   # how many targets to try before giving up on a request
RETRY_STATUSES = {401, 402, 403, 429, 500, 502, 503, 504}   # statuses that mean "try another target"
COOLDOWN = 60.0   # seconds to skip a target after it returns a bad status


def extract_model(body_bytes: bytes) -> Optional[str]:   # pull the "model" field from a JSON request body
    if not body_bytes:   # GET/DELETE have no body
        return None   # no model to extract
    try:   # attempt to parse the body as JSON
        parsed = json.loads(body_bytes)   # parse to a Python object
    except Exception:   # body isn't valid JSON
        return None   # can't extract model
    if isinstance(parsed, dict):   # only JSON objects carry a "model" field
        return parsed.get("model")   # return the model string (or None)
    return None   # JSON arrays/scalars have no model field


def target_id(target: dict) -> str:   # build a unique key for the health tracker
    """Unique key for the health tracker: url + key combo."""   # docstring
    return f"{target['url']}::{target['key']}"   # url+key uniquely identifies one target slot


class Router:   # the unified proxy: one retry loop for all models and providers
    def __init__(self, config: ModelConfig, client=None):   # construct from a loaded config + optional shared client
        self.config = config   # keep a reference to resolve model names
        self.client = client or make_client()   # shared httpx pool (injected for tests, created if absent)
        self.health = HealthTracker(cooldown_seconds=COOLDOWN)   # per-target cooldown tracker
        self.pools = {}   # model name -> RoundRobin cursor over its targets
        self._target_keys = {}   # model name -> list of target_ids (for healthy count)
        for name, cfg in config.models.items():   # iterate all loaded models
            targets = cfg.get("targets", [])   # get this model's target list
            if targets:   # skip models with no targets
                self.pools[name] = RoundRobin(targets)   # wrap targets in a rotating cursor
                self._target_keys[name] = [target_id(t) for t in targets]   # precompute ids for health checks

    def pick_target(self, exact_name: str) -> Optional[dict]:   # choose a healthy target
        """Pick a healthy target. Returns None if the pool is empty OR fully cooled down."""   # docstring
        rr = self.pools.get(exact_name)   # get the round-robin for this model
        if rr is None or len(rr) == 0:   # pool missing or empty
            return None   # caller returns 503
        total = len(rr)   # pool size
        for _ in range(total):   # at most one full pass over the pool
            t = rr.next()   # get the next target by rotation
            if t is None:   # defensive: empty pool
                return None   # give up
            if not self.health.is_down(target_id(t)):   # target is not in cooldown
                return t   # use this healthy target
        return None   # ALL targets cooling down — signal exhaustion (don't waste attempts)

    async def handle(self, method, path, body_bytes, client_headers, query_params,   # main entry point
                     user_id=None, is_admin=False):   # billing identity from the gate
        model = extract_model(body_bytes)   # pull the model name from the body
        if model is None:   # client didn't specify a model
            return JSONResponse(   # 400 bad request
                {"error": {"message": "missing 'model' field in request body", "type": "invalid_request"}},   # OpenAI-style error
                status_code=400,
            )

        exact, targets = self.config.resolve(model)   # resolve the client name to exact name + targets
        if exact is None:   # unknown model requested
            available = list(self.config.models.keys())   # all valid model names (for the error message)
            return JSONResponse(   # 404 model not found
                {"error": {"message": f"model '{model}' not found. Available: {available}",   # helpful message
                           "type": "model_not_found"}},
                status_code=404,
            )

        attempts = 0   # how many targets we've tried
        last_error = None   # remember the latest error for the final 502
        while attempts < MAX_ATTEMPTS:   # retry loop across targets
            target = self.pick_target(exact)   # pick a healthy target
            if target is None:   # pool empty or fully cooled down
                healthy = self.health.healthy_count(self._target_keys.get(exact, []))   # count usable targets
                return JSONResponse(   # 503 service unavailable (fail fast, don't burn attempts)
                    {"error": {
                        "message": f"no targets available for '{exact}' "   # include pool diagnostics
                                   f"(pool={len(self.pools.get(exact, []))}, healthy={healthy})",
                        "type": "server_error"}},
                    status_code=503,
                )

            upstream = target.get("upstream_model")   # what name to rewrite the model field to
            swapped = swap_model(body_bytes, upstream) if upstream else body_bytes   # rewrite (or pass through)
            tid = target_id(target)   # unique key for the health tracker

            logger.info("%s %s model=%s target=%s attempt=%d",   # log each attempt
                        method, path, exact, target["url"], attempts + 1)

            try:   # open the upstream connection
                response = await forward(self.client, target, method, path, swapped,   # send to the picked target
                                         client_headers, query_params)
            except Exception as exc:   # connection refused / timeout / DNS error
                logger.warning("forward failed: %s -> %s", tid, exc)   # log the failure
                self.health.mark_down(tid)   # cool down this target
                last_error = str(exc)   # save the message
                attempts += 1   # try the next target
                continue   # back to top of retry loop

            if response.status_code in RETRY_STATUSES:   # this target is bad (auth/quota/rate-limit/5xx)
                body_preview = b""   # capture a short error body for logging
                try:   # drain the response body so the connection can close
                    body_preview = await response.aread()
                except Exception:   # body read failed (e.g. already closed)
                    pass   # skip the preview
                await response.aclose()   # close the upstream response
                logger.warning("upstream %d: %s -> %s",   # log the bad status + truncated body
                               response.status_code, tid, body_preview[:200])
                self.health.mark_down(tid)   # cool down the bad target
                last_error = f"upstream returned {response.status_code}"   # remember for the 502
                attempts += 1   # try the next target
                continue   # back to retry loop

            self.health.mark_up(tid)   # success path: mark target healthy (clears cooldown)
            ctype = response.headers.get("content-type", "")   # inspect the upstream content type
            is_stream = "text/event-stream" in ctype   # SSE streams need filtering + chunked passthrough
            resp_headers = clean_response_headers(response)   # strip framing headers Starlette sets itself
            billing_on = self._billing_on()   # cache the billing-enabled check (one call per response)
            meter = user_id if (billing_on and not is_admin and user_id) else None   # who to bill (None = unmetered)

            if is_stream:   # streaming (SSE) response
                sse = SSEFilter()   # one filter per stream (holds partial-line buffer)
                watcher = UsageWatcher() if meter else None   # meter streaming usage if billing

                async def gen():   # generator that yields filtered upstream chunks
                    try:   # ensure cleanup runs even if client disconnects
                        async for chunk in response.aiter_raw():   # iterate raw (undecoded) bytes
                            if chunk:   # skip empty keepalive chunks
                                filtered = sse.feed(chunk)   # filter to data: lines only
                                if watcher:   # we're metering this stream
                                    watcher.feed(filtered)   # scan for the usage chunk
                                if filtered:   # may be empty if chunk ended mid-line
                                    yield filtered   # forward the filtered bytes
                        tail = sse.flush()   # emit any buffered partial line at stream end
                        if tail:   # there's a qualifying leftover
                            if watcher:   # check the tail for usage too
                                watcher.feed(tail)   # scan final partial line
                            yield tail   # forward it
                    except Exception as exc:   # stream broke mid-flight
                        logger.error("stream error: %s", exc)   # log it (can't recover)
                    finally:   # always release the upstream connection
                        if watcher and watcher.has_usage:   # we captured usage from the stream
                            await self._burn(meter, exact, watcher.tok_in, watcher.tok_out)   # charge
                        await response.aclose()   # close upstream response

                return StreamingResponse(gen(), status_code=response.status_code,   # streaming response
                                         headers=resp_headers, media_type="text/event-stream")   # SSE media type

            raw = await response.aread()   # non-stream: buffer the whole body
            await response.aclose()   # close the upstream response
            if meter:   # metered non-streaming: parse usage + charge
                await self._burn_from_body(meter, exact, raw)   # extract + burn
            return Response(content=raw, status_code=response.status_code,   # buffered response
                            headers=resp_headers, media_type=ctype or None)   # forward content type

        return JSONResponse(   # all attempts exhausted
            {"error": {"message": f"all {MAX_ATTEMPTS} attempts failed: {last_error}",   # report last failure
                       "type": "upstream_error"}},   # OpenAI-style error
            status_code=502,   # bad gateway
        )

    @staticmethod
    def _billing_on() -> bool:   # is the billing system configured?
        """Check if billing is enabled (gist store + wallet present)."""   # docstring
        from .billing import gate   # local import (avoid a hard dep at module load)
        return gate.is_enabled()   # True when both secrets are set

    async def _burn_from_body(self, user_id: str, model_name: str, body: bytes) -> None:   # meter a buffered response
        """Parse usage from a buffered JSON response body and charge the user."""   # docstring
        try:   # the body may not be JSON (e.g. an image response)
            parsed = json.loads(body)   # parse the response body
        except Exception:   # not JSON — can't meter
            return   # skip billing (non-JSON responses aren't metered)
        tok_in, tok_out = gate_extract_usage(parsed)   # pull prompt/completion tokens
        if tok_in or tok_out:   # only charge if there were tokens
            await self._burn(user_id, model_name, tok_in, tok_out)   # charge

    async def _burn(self, user_id: str, model_name: str, tok_in: int, tok_out: int) -> None:   # charge one request
        """Charge a user for a request. Best-effort — never raises."""   # docstring
        from .billing import gate   # local import
        try:   # billing failures must NEVER break the response
            await gate.burn(user_id, self.config.models, model_name, tok_in, tok_out)   # spend + log
        except Exception as exc:   # any billing error
            logger.warning("burn failed for %s: %s", user_id, exc)   # log and move on