Update main.py

main.py

@@ -91,33 +91,25 @@ async def chat(request: ChatCompletionRequest):
     return StreamingResponse(format_response(chat_chunks), media_type="text/event-stream")
 
 async def stream_response(tokens: Any) -> None:
-    async with send:
-        try:
-            iterator: Generator = llm.generate(tokens)
-            async for chat_chunk in iterator:
-                response = {
-                    'choices': [
-                        {
-                            'message': {
-                                'role': 'system',
-                                'content': llm.detokenize(chat_chunk)
-                            },
-                            'finish_reason': 'stop' if llm.detokenize(chat_chunk) == "[DONE]" else 'unknown'
-                        }
-                    ]
-                }
-                await send({
-                    "type": "http.response.body",
-                    "body": f"data: {json.dumps(response)}\n\n".encode("utf-8"),
-                    "more_body": True,
-                })
-            await send({
-                "type": "http.response.body",
-                "body": b"event: done\ndata: {}\n\n",
-                "more_body": False,
-            })
-        except Exception as e:
-            print(f"Exception in event publisher: {str(e)}")
+    try:
+        iterator: Generator = llm.generate(tokens)
+        for chat_chunk in iterator:
+            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 b"event: done\ndata: {}\n\n"
+    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])