Update api/utils.py

api/utils.py

@@ -17,42 +17,46 @@ async def process_streaming_response(request: ChatRequest):
     messages = [msg.dict() for msg in request.messages]
 
     try:
-        async for content in AmigoChat.generate_response(
-            model=request.model,
-            messages=messages,
-            stream=True
-        ):
-            timestamp = int(datetime.now().timestamp())
-            yield json.dumps({
+        async def event_generator():
+            async for content in AmigoChat.generate_response(
+                model=request.model,
+                messages=messages,
+                stream=True
+            ):
+                timestamp = int(datetime.now().timestamp())
+                chunk = {
+                    "id": f"chatcmpl-{uuid.uuid4()}",
+                    "object": "chat.completion.chunk",
+                    "created": timestamp,
+                    "model": request.model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": content},
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+                yield f"data: {json.dumps(chunk)}\n\n"
+
+            # Indicate the end of the stream
+            end_chunk = {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion.chunk",
-                "created": timestamp,
+                "created": int(datetime.now().timestamp()),
                 "model": request.model,
                 "choices": [
                     {
                         "index": 0,
-                        "delta": {"content": content},
-                        "finish_reason": None,
+                        "delta": {},
+                        "finish_reason": "stop",
                     }
                 ],
-                "usage": None,
-            }) + "\n"
+            }
+            yield f"data: {json.dumps(end_chunk)}\n\n"
+            yield "data: [DONE]\n\n"
 
-        # Indicate the end of the stream
-        yield json.dumps({
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion.chunk",
-            "created": int(datetime.now().timestamp()),
-            "model": request.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "delta": {},
-                    "finish_reason": "stop",
-                }
-            ],
-            "usage": None,
-        }) + "\n"
+        return event_generator()
 
     except Exception as e:
         logger.error(f"Error in streaming response: {e}")