Spaces:
Running
Running
File size: 5,237 Bytes
a056d23 09e4b39 a056d23 8fc20f8 c6b6b9c 8fc20f8 aae924a 09e4b39 a056d23 c6b6b9c 09e4b39 dbc3826 a056d23 c6b6b9c a056d23 8fc20f8 a056d23 09e4b39 a056d23 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | 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,
}
|