Stream (don't buffer) the chat response so members see tokens immediately

This commit is contained in:
inference-bot committed 2026-09-14 15:15:58 -06:00
1 parent 87374dfeab
commit 841eb70e4d
1 file changed
+32 -15
+32 -15
View File
@@ -26,7 +26,7 @@ from email.mime.multipart import MIMEMultipart
import httpx import httpx
import jwt import jwt
from fastapi import FastAPI, Request, Response, HTTPException from fastapi import FastAPI, Request, Response, HTTPException
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse, StreamingResponse
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("member-portal") logger = logging.getLogger("member-portal")
@@ -921,29 +921,46 @@ async def inject_key(request: Request, path: str):
if not member_key: if not member_key:
raise HTTPException(403, "No active membership key") raise HTTPException(403, "No active membership key")
# Forward the request to LiteLLM with the member's key. # Forward the request to LiteLLM with the member's key, STREAMING the
# Use a generous timeout: model generations can take well over the httpx # response back token-by-token (or byte-by-byte) rather than buffering it.
# default of 5s, and the previous code buffered the full response which # Buffering collapsed the model's stream and handed OpenWebUI the full
# caused httpx.ReadTimeout on any non-trivial generation. # 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() body = await request.body()
headers = dict(request.headers) headers = dict(request.headers)
headers["authorization"] = f"Bearer {member_key}" headers["authorization"] = f"Bearer {member_key}"
headers.pop("host", None) headers.pop("host", None)
headers.pop("content-length", 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) 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( # Send the request upstream, but don't buffer the body — use a stream.
content=upstream.content, 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, 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",
},
) )