From 3f392bffca41d2464ab2103b44e0d19dbd32553d Mon Sep 17 00:00:00 2001 From: jze9 Date: Thu, 28 May 2026 20:35:03 +0500 Subject: [PATCH] Fix streaming TransferEncodingError DB saves were happening inside the generator after [DONE], causing connection abort if any await failed. Moved to BackgroundTask. Co-Authored-By: Claude Sonnet 4.6 --- main.py | 30 ++++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/main.py b/main.py index b03e739..2dba296 100644 --- a/main.py +++ b/main.py @@ -129,22 +129,28 @@ async def chat_completions(req: ChatRequest, bg: BackgroundTasks): if m.role == "user"), "") if req.stream: + collected: list[str] = [] + async def gen() -> AsyncIterator[bytes]: - full_response = [] - async for chunk in oll.chat_stream(messages, model=req.model): - full_response.append(chunk) - data = json.dumps({"choices": [{"delta": {"content": chunk}, - "finish_reason": None}]}) - yield f"data: {data}\n\n".encode() - yield b"data: [DONE]\n\n" - # save after stream - response_text = "".join(full_response) - await mem.save_turn(req.session_id, "user", user_content) + try: + async for chunk in oll.chat_stream(messages, model=req.model): + collected.append(chunk) + data = json.dumps({"choices": [{"delta": {"content": chunk}, + "finish_reason": None}]}) + yield f"data: {data}\n\n".encode() + yield b"data: [DONE]\n\n" + except Exception as e: + yield f"data: {json.dumps({'error': str(e)})}\n\n".encode() + + async def _save_stream(): + response_text = "".join(collected) + await mem.save_turn(req.session_id, "user", user_content) await mem.save_turn(req.session_id, "assistant", response_text) - await mem.extract_and_store(user_content, user_id=req.user_id) - await mem.extract_and_store(response_text, user_id=req.user_id) + await mem.extract_and_store(user_content, user_id=req.user_id) + await mem.extract_and_store(response_text, user_id=req.user_id) await mem.summarize_old_turns(req.session_id) + bg.add_task(_save_stream) return StreamingResponse(gen(), media_type="text/event-stream") response_text = await oll.chat(messages, model=req.model)