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
+27 -10
+27 -10
View File
@@ -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(
# 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)
return Response(
content=upstream.content,
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",
},
)