Update main.py

main.py

@@ -6,9 +6,10 @@ from fastapi import HTTPException
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
+from starlette.responses import StreamingResponse
 from ctransformers import AutoModelForCausalLM
 from pydantic import BaseModel
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Generator
 
 
 llm = AutoModelForCausalLM.from_pretrained("TheBloke/WizardCoder-15B-1.0-GGML",
@@ -60,7 +61,7 @@ async def completion(request: ChatCompletionRequest, response_mode=None):
     return response
 
 @app.post("/v1/chat/completions")
-async def chat(request: ChatCompletionRequestV2):
+async def chat(request: ChatCompletionRequest):
     tokens = llm.tokenize([message.content for message in request.messages])
     
     try:
@@ -68,21 +69,23 @@ async def chat(request: ChatCompletionRequestV2):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-    def format_response(chat_chunks) -> Dict[str, Any]:
-        response = {
-            'choices': []
-        }
+    async def format_response(chat_chunks: Generator) -> Any:
         for chat_chunk in chat_chunks:
-            response['choices'].append({
-                'message': {
-                    'role': 'system',
-                    'content': llm.detokenize(chat_chunk)
-                },
-                'finish_reason': 'stop' if llm.detokenize(chat_chunk) == "[DONE]" else 'unknown'
-            })
-        return response
+            response = {
+                'choices': [
+                    {
+                        'message': {
+                            'role': 'system',
+                            'content': llm.detokenize(chat_chunk)
+                        },
+                        'finish_reason': 'stop' if llm.detokenize(chat_chunk) == "[DONE]" else 'unknown'
+                    }
+                ]
+            }
+            yield f"data: {json.dumps(response)}\n\n"
+        yield "event: done\ndata: {}\n\n"
 
-    return format_response(chat_chunks)
+    return StreamingResponse(format_response(chat_chunks), media_type="text/event-stream")
 
 @app.post("/v0/chat/completions")
 async def chat(request: ChatCompletionRequest, response_mode=None):