Update main.py

main.py

@@ -6,7 +6,7 @@ from functools import partial
 import fastapi
 import uvicorn
 from fastapi import HTTPException, Depends, Request
-from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
 from anyio import create_memory_object_stream
@@ -87,7 +87,7 @@ async def chat(request: ChatCompletionRequest):
 
     return EventSourceResponse(generate_response(chat_chunks, llm))
 
-async def stream_response(tokens: Any) -> None:
+async def stream_response(tokens, llm):
     try:
         iterator: Generator = llm.generate(tokens)
         for chat_chunk in iterator:
@@ -102,22 +102,17 @@ async def stream_response(tokens: Any) -> None:
                     }
                 ]
             }
-            yield f"data: {json.dumps(response)}\n\n"
-        yield b"event: done\ndata: {}\n\n"
+            yield dict(data=json.dumps(response))
+        yield dict(data="[DONE]")
     except Exception as e:
         print(f"Exception in event publisher: {str(e)}")
 
-
-async def chatV2(request: Request, body: ChatCompletionRequest):
-    combined_messages = ' '.join([message.content for message in body.messages])
-    tokens = llm.tokenize(combined_messages)
-
-    return StreamingResponse(stream_response(tokens))
-
 @app.post("/v2/chat/completions")
 async def chatV2_endpoint(request: Request, body: ChatCompletionRequest):
-    return await chatV2(request, body)
+    combined_messages = ' '.join([message.content for message in body.messages])
+    tokens = llm.tokenize(combined_messages)
 
+    return EventSourceResponse(stream_response(tokens, llm))
 
 @app.post("/v0/chat/completions")
 async def chat(request: ChatCompletionRequestV0, response_mode=None):