diff --git a/app/main.py b/app/main.py index b0c905d..b10dd7b 100644 --- a/app/main.py +++ b/app/main.py @@ -26,7 +26,7 @@ from email.mime.multipart import MIMEMultipart import httpx import jwt from fastapi import FastAPI, Request, Response, HTTPException -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse logging.basicConfig(level=logging.INFO) logger = logging.getLogger("member-portal") @@ -921,29 +921,46 @@ async def inject_key(request: Request, path: str): if not member_key: raise HTTPException(403, "No active membership key") - # Forward the request to LiteLLM with the member's key. - # Use a generous timeout: model generations can take well over the httpx - # default of 5s, and the previous code buffered the full response which - # caused httpx.ReadTimeout on any non-trivial generation. + # Forward the request to LiteLLM with the member's key, STREAMING the + # response back token-by-token (or byte-by-byte) rather than buffering it. + # Buffering collapsed the model's stream and handed OpenWebUI the full + # answer in one blob, so members saw nothing until generation finished — + # which made chat feel slow even though time-to-first-byte was fine. body = await request.body() headers = dict(request.headers) headers["authorization"] = f"Bearer {member_key}" headers.pop("host", None) headers.pop("content-length", None) + headers.pop("accept-encoding", None) # let httpx handle decompression timeout = httpx.Timeout(connect=15.0, read=600.0, write=30.0, pool=15.0) - async with httpx.AsyncClient(timeout=timeout) as client: - upstream = await client.request( - method=request.method, - url=f"{LITELLM_BASE}/v1/{path}", - headers=headers, - content=body, - ) - return Response( - content=upstream.content, + # Send the request upstream, but don't buffer the body — use a stream. + client = httpx.AsyncClient(timeout=timeout) + upstream_req = client.build_request( + method=request.method, + url=f"{LITELLM_BASE}/v1/{path}", + headers=headers, + content=body, + ) + upstream = await client.send(upstream_req, stream=True) + + async def gen(): + try: + async for chunk in upstream.aiter_raw(): + yield chunk + finally: + await upstream.aclose() + await client.aclose() + + # Pass through the content-type (critical for SSE streaming) and status. + return StreamingResponse( + gen(), status_code=upstream.status_code, - headers={"content-type": upstream.headers.get("content-type", "application/json")}, + headers={ + "content-type": upstream.headers.get("content-type", "application/json"), + "cache-control": "no-cache", + }, )