jze9/code-memory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix streaming TransferEncodingError

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jze9
2026-05-28 20:35:03 +05:00
parent 7a0c819eca
commit 3f392bffca
1 changed files with 18 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
main.py
View File

@@ -129,22 +129,28 @@ async def chat_completions(req: ChatRequest, bg: BackgroundTasks):
if m.role == "user"), "") if m.role == "user"), "")
if req.stream: if req.stream:
collected: list[str] = []
async def gen() -> AsyncIterator[bytes]: async def gen() -> AsyncIterator[bytes]:
full_response = [] try:
async for chunk in oll.chat_stream(messages, model=req.model): async for chunk in oll.chat_stream(messages, model=req.model):
full_response.append(chunk) collected.append(chunk)
data = json.dumps({"choices": [{"delta": {"content": chunk}, data = json.dumps({"choices": [{"delta": {"content": chunk},
"finish_reason": None}]}) "finish_reason": None}]})
yield f"data: {data}\n\n".encode() yield f"data: {data}\n\n".encode()
yield b"data: [DONE]\n\n" yield b"data: [DONE]\n\n"
# save after stream except Exception as e:
response_text = "".join(full_response) 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, "user", user_content)
await mem.save_turn(req.session_id, "assistant", response_text) 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(user_content, user_id=req.user_id)
await mem.extract_and_store(response_text, 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) await mem.summarize_old_turns(req.session_id)
bg.add_task(_save_stream)
return StreamingResponse(gen(), media_type="text/event-stream") return StreamingResponse(gen(), media_type="text/event-stream")
response_text = await oll.chat(messages, model=req.model) response_text = await oll.chat(messages, model=req.model)