Fixed streaming in chat_completion

main/routes.py

@@ -70,12 +70,10 @@ async def create_chat_completion(request: ChatCompletionRequest):
         last_message = request.messages[-1].content
 
         if request.stream:
-            # For streaming, we need to create a generator that yields OpenAI-compatible chunks
             async def generate_stream():
                 async for chunk in api.generate_stream(
                         prompt=last_message,
                 ):
-                    # Create a streaming response chunk in OpenAI format
                     response_chunk = {
                         "id": "chatcmpl-123",
                         "object": "chat.completion.chunk",
@@ -89,16 +87,18 @@ async def create_chat_completion(request: ChatCompletionRequest):
                             "finish_reason": None
                         }]
                     }
+                    # Need to format this exactly as SSE requires
                     yield f"data: {json.dumps(response_chunk)}\n\n"
-
-                # Send the final chunk
-                yield f"data: [DONE]\n\n"
+                yield "data: [DONE]\n\n"
 
             return StreamingResponse(
                 generate_stream(),
-                media_type="text/event-stream"
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                }
             )
-
         else:
             # For non-streaming, generate the full response
             response_text = await api.generate_response(