src/services/forwarder.py

"""HTTP forwarding helpers — target dicts + shared connection pool + SSE filter."""   # module docstring
from __future__ import annotations   # defer annotation evaluation (PEP 563)

import json   # stdlib json for the model-field swap
import logging   # structured logging for connection events
from typing import Dict   # type alias for header dicts
from urllib.parse import urlparse   # extract the host from a target URL

import httpx   # async HTTP client + connection pool

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

# Headers that must not be forwarded end-to-end (RFC 7230).   # spec-mandated strip list
HOP_BY_HOP = {   # request headers to strip before forwarding
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",   # connection-level
    "te", "trailers", "transfer-encoding", "upgrade", "host", "content-length",   # framing + host
    "content-encoding",   # upstream may re-encode; don't double-apply
}

# Response headers Starlette sets itself — strip these from upstream.   # avoid header conflicts
RESPONSE_STRIP = {   # response headers to strip before returning to client
    "content-length", "transfer-encoding", "connection",   # framing (Starlette recomputes)
    "content-encoding", "keep-alive",   # encoding + connection
}


def clean_headers(headers) -> Dict[str, str]:   # strip hop-by-hop headers before forwarding to upstream
    return {k: v for k, v in headers.items() if k.lower() not in HOP_BY_HOP}   # keep only end-to-end


def clean_response_headers(response) -> Dict[str, str]:   # strip framing headers Starlette will set itself
    return {k: v for k, v in response.headers.items() if k.lower() not in RESPONSE_STRIP}   # dedupe


def swap_model(body_bytes: bytes, new_model: str) -> bytes:   # rewrite the JSON 'model' field
    """Rewrite the JSON 'model' field. Passes through non-JSON bodies unchanged."""   # docstring
    if not body_bytes:   # GET/DELETE have no body
        return body_bytes   # nothing to rewrite
    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 (e.g. form data)
        return body_bytes   # pass it through untouched
    if isinstance(parsed, dict):   # only JSON objects carry a "model" field
        parsed["model"] = new_model   # overwrite the client's model with the target's upstream name
        return json.dumps(parsed).encode("utf-8")   # re-serialize to bytes
    return body_bytes   # JSON arrays/scalars have no model field; leave as-is


def make_client() -> httpx.AsyncClient:   # create the shared connection pool
    """Shared connection pool with split timeouts (long read for slow LLMs)."""   # docstring
    return httpx.AsyncClient(   # one client reused across all requests (connection pooling)
        timeout=httpx.Timeout(connect=10.0, read=300.0, write=30.0, pool=5.0),   # split: fast connect, long read
        limits=httpx.Limits(max_connections=200, max_keepalive_connections=50),   # cap concurrent sockets
    )


# Module-level singleton so auth routes (which run before the proxy router is
# wired) can reuse the same connection pool without a circular import.
_shared_client: httpx.AsyncClient | None = None   # set by app.py at startup


def set_shared_client(client: httpx.AsyncClient) -> None:   # called once at startup
    """Store the shared httpx client so other modules (auth) can access it."""   # docstring
    global _shared_client   # mutate the module global
    _shared_client = client   # store the reference


def get_shared_client() -> httpx.AsyncClient:   # fetch the shared client
    """Return the shared client. Falls back to a fresh one if not yet wired."""   # docstring
    if _shared_client is None:   # not yet initialized (shouldn't happen in prod)
        return make_client()   # create a throwaway so the call never crashes
    return _shared_client   # the real shared pool


async def forward(   # open a streaming upstream connection
    client: httpx.AsyncClient,   # the shared connection pool
    target: dict,   # {url, key, upstream_model} describing where to send
    method: str,   # HTTP verb (GET/POST/...)
    path: str,   # path after the host, e.g. "v1/chat/completions"
    body_bytes: bytes,   # already-model-swapped request body
    client_headers: dict,   # cleaned client headers to forward
    query_params: dict,   # query string to forward unchanged
) -> httpx.Response:   # streaming response; caller owns .aclose()
    """Open a streaming upstream connection. Caller owns response.aclose()."""   # docstring
    # The client sends /v1/... but each target URL already has its own version path   # different providers use /v1, /v4, /v1beta/openai, etc.
    suffix = path[3:] if path.startswith("v1/") else path.lstrip("/")   # strip "v1/" prefix (keep rest)
    url = f"{target['url'].rstrip('/')}/{suffix}"   # build the full upstream URL
    headers = clean_headers(client_headers)   # drop hop-by-hop headers
    headers["authorization"] = f"Bearer {target['key']}"   # inject this target's key (replaces client's)
    headers["host"] = urlparse(target["url"]).netloc   # set Host to the upstream authority
    req = client.build_request(method, url, content=body_bytes, headers=headers, params=query_params)   # build without sending
    return await client.send(req, stream=True)   # send and keep open for streaming


class SSEFilter:   # buffers raw bytes and emits only 'data:' lines + blank separators
    """Buffers raw bytes and emits only 'data:' lines + blank separators.   # docstring start

    Filters out SSE comments (':'), 'event:', 'id:', 'retry:' lines that some   # why we filter
    clients (notably pi) can't parse.   # docstring end
    """

    def __init__(self):   # construct with an empty byte buffer
        self._buf = b""   # holds a partial line spanning chunk boundaries

    def feed(self, chunk: bytes) -> bytes:   # process an incoming chunk
        """Process a chunk. Returns the filtered bytes ready to forward."""   # docstring
        self._buf += chunk   # append to whatever was left from the last chunk
        out = b""   # collect the filtered output
        while b"\n" in self._buf:   # as long as we have at least one complete line
            line, self._buf = self._buf.split(b"\n", 1)   # peel off one line, keep the rest
            if self._keep(line):   # should this line be forwarded?
                out += line + b"\n"   # re-append the newline we split on
        return out   # may be empty if the chunk ended mid-line

    def flush(self) -> bytes:   # emit any buffered partial line at stream end
        """Emit any buffered partial line if it qualifies."""   # docstring
        if self._buf and self._keep(self._buf):   # leftover bytes that are a valid line
            result = self._buf   # take them
            self._buf = b""   # clear the buffer
            return result   # emit
        self._buf = b""   # discard non-qualifying leftover
        return b""   # nothing to emit

    @staticmethod   # doesn't need self
    def _keep(line: bytes) -> bool:   # should this line be forwarded?
        """Should this line be forwarded? data: lines and blank separators only."""   # docstring
        stripped = line.strip()   # drop trailing \r and surrounding whitespace
        return stripped.startswith(b"data:") or stripped == b""   # keep data: lines + blank SSE separators