Stream (don't buffer) the chat response so members see tokens immediately
This commit is contained in:
1 parent
87374dfeab
commit
841eb70e4d
1 file changed
+32
-15
+32
-15
@@ -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",
|
||||||
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user