"""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