Update main.py

main.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel
 from typing import List, Dict, Any, Optional
 import time
 import uuid
+from fastapi.responses import StreamingResponse
 
 # Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
@@ -194,27 +195,45 @@ async def chat_completions(request: ChatRequest):
         messages=messages
     )
 
-    response_content = ""
-    async for chunk in async_generator:
-        response_content += chunk if isinstance(chunk, str) else chunk.content  # Concatenate response
-
-    # Generate a unique ID and use the current timestamp
     response_id = f"chatcmpl-{uuid.uuid4()}"
     created_timestamp = int(time.time())
 
-    return {
-        "id": response_id,
-        "object": "chat.completion",
-        "created": created_timestamp,
-        "model": request.model,
-        "choices": [
-            {
-                "message": {
-                    "role": "assistant",
-                    "content": response_content
-                },
-                "finish_reason": "stop",
-                "index": 0
-            }
-        ]
-    }
+    async def generate_response():
+        # Stream the response content
+        async for chunk in async_generator:
+            if isinstance(chunk, str):
+                yield {
+                    "id": response_id,
+                    "object": "chat.completion.chunk",
+                    "created": created_timestamp,
+                    "model": request.model,
+                    "choices": [
+                        {
+                            "message": {
+                                "role": "assistant",
+                                "content": chunk
+                            },
+                            "finish_reason": None,
+                            "index": 0
+                        }
+                    ]
+                }
+            else:
+                yield {
+                    "id": response_id,
+                    "object": "chat.completion.chunk",
+                    "created": created_timestamp,
+                    "model": request.model,
+                    "choices": [
+                        {
+                            "message": {
+                                "role": "assistant",
+                                "content": chunk.content
+                            },
+                            "finish_reason": None,
+                            "index": 0
+                        }
+                    ]
+                }
+
+    return StreamingResponse(generate_response(), media_type="application/json")