proxy/client.py

import asyncio
import json
import logging
from typing import Any, AsyncIterator

import httpx
from fastapi import HTTPException
from fastapi.responses import StreamingResponse
from starlette.requests import Request

from config import POLL_INTERVAL, POLL_MAX_ATTEMPTS
from proxy.providers import Provider

logger = logging.getLogger(__name__)

_TIMEOUT = httpx.Timeout(10.0, connect=10.0, read=180.0)


def _build_headers(provider: Provider, real_key: str, incoming_headers: dict[str, str]) -> dict[str, str]:
    headers = {k: v for k, v in incoming_headers.items() if k.lower() not in ("host", "authorization", "content-length")}
    headers[provider.auth_header] = provider.auth_value(real_key)
    headers.setdefault("Accept", "application/json")
    headers.setdefault("Content-Type", "application/json")
    return headers


def _extract_request_id(response: httpx.Response) -> str | None:
    try:
        data = response.json()
        if isinstance(data, dict):
            return data.get("requestId") or data.get("id") or data.get("request_id")
    except Exception:
        pass
    return response.headers.get("Location", "").rstrip("/").split("/")[-1] or None


async def _poll_nvidia(client: httpx.AsyncClient, provider: Provider, request_id: str) -> httpx.Response:
    status_url = f"{provider.base_url}/v1/status/{request_id}"
    logger.info("NVIDIA returned 202; polling %s", status_url)

    for attempt in range(1, POLL_MAX_ATTEMPTS + 1):
        await asyncio.sleep(POLL_INTERVAL)
        poll_resp = await client.get(status_url, timeout=_TIMEOUT)
        logger.debug("Poll attempt %d: status %d", attempt, poll_resp.status_code)

        if poll_resp.status_code == 200:
            return poll_resp
        if poll_resp.status_code == 202:
            continue
        poll_resp.raise_for_status()

    raise HTTPException(status_code=504, detail=f"NVIDIA async job {request_id} timed out after polling")


def _upstream_error_message(status_code: int, content: bytes) -> str:
    """Try to extract a readable message from a provider error response."""
    text = content.decode("utf-8", errors="replace").strip()
    if not text:
        return f"Upstream provider returned HTTP {status_code}"
    try:
        data = json.loads(text)
        if isinstance(data, dict):
            if isinstance(data.get("error"), dict):
                return data["error"].get("message") or json.dumps(data["error"])
            return data.get("message") or data.get("error") or text
    except Exception:
        pass
    return text


async def _stream_response(provider: Provider, real_key: str, body: dict[str, Any]) -> AsyncIterator[bytes]:
    """Async generator that manages its own httpx client for streaming."""
    target_url = f"{provider.base_url}/v1/chat/completions"
    headers = _build_headers(provider, real_key, {})
    async with httpx.AsyncClient(timeout=_TIMEOUT, follow_redirects=True) as client:
        async with client.stream("POST", target_url, headers=headers, json=body) as upstream:
            if upstream.status_code >= 400:
                content = await upstream.aread()
                detail = _upstream_error_message(upstream.status_code, content)
                raise HTTPException(status_code=upstream.status_code, detail=detail)
            async for chunk in upstream.aiter_bytes():
                yield chunk


async def forward_request(provider: Provider, real_key: str, request: Request, body: dict[str, Any]) -> Any:
    target_url = f"{provider.base_url}/v1/chat/completions"
    headers = _build_headers(provider, real_key, dict(request.headers))

    if bool(body.get("stream")):
        return StreamingResponse(_stream_response(provider, real_key, body), media_type="text/event-stream")

    async with httpx.AsyncClient(timeout=_TIMEOUT, follow_redirects=True) as client:
        try:
            upstream = await client.post(target_url, headers=headers, json=body)
        except httpx.TimeoutException as exc:
            logger.warning("Upstream timeout for %s: %s", provider.name, exc)
            raise HTTPException(status_code=504, detail="Upstream provider timed out")
        except httpx.RequestError as exc:
            logger.warning("Upstream request error for %s: %s", provider.name, exc)
            raise HTTPException(status_code=502, detail=f"Could not reach provider '{provider.name}': {exc}")

        if provider.handles_202 and upstream.status_code == 202:
            request_id = _extract_request_id(upstream)
            if not request_id:
                raise HTTPException(status_code=502, detail="Provider returned 202 but no requestId was found")
            upstream = await _poll_nvidia(client, provider, request_id)

        if upstream.status_code >= 400:
            content = await upstream.aread()
            logger.warning("Upstream error from %s: %d %s", provider.name, upstream.status_code, content[:500])
            detail = _upstream_error_message(upstream.status_code, content)
            raise HTTPException(status_code=upstream.status_code, detail=detail)

        return {
            "status_code": upstream.status_code,
            "headers": dict(upstream.headers),
            "content": upstream.content,
        }