Update main.py

main.py

@@ -90,7 +90,7 @@ async def chat(request: ChatCompletionRequest):
 
     return StreamingResponse(format_response(chat_chunks), media_type="text/event-stream")
 
-async def stream_response(send: Callable) -> None:
+async def stream_response(tokens: Any) -> None:
     async with send:
         try:
             iterator: Generator = llm.generate(tokens)
@@ -123,7 +123,7 @@ async def chatV2(request: Request, body: ChatCompletionRequest):
     combined_messages = ' '.join([message.content for message in body.messages])
     tokens = llm.tokenize(combined_messages)
 
-    return ResponseGenerator(stream_response)
+    return StreamingResponse(stream_response(tokens))
 
 @app.post("/v2/chat/completions")
 async def chatV2_endpoint(request: Request, body: ChatCompletionRequest):
@@ -141,4 +141,4 @@ async def chat(request: ChatCompletionRequestV0, response_mode=None):
     return EventSourceResponse(server_sent_events(tokens, llm))
 
 if __name__ == "__main__":
-  uvicorn.run(app, host="0.0.0.0", port=8000)
+  uvicorn.run(app, host="0.0.0.0", port=8000)