:bulb: [Test] Websocket API of create conversation

apis/chat_api.py

@@ -1,6 +1,7 @@
 import uvicorn
 
-from fastapi import FastAPI
+from fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnect
+from fastapi.routing import APIRoute
 from pydantic import BaseModel, Field
 from conversations import ConversationSession
 
@@ -43,11 +44,26 @@ class ChatAPIApp:
         ]
         return self.available_models
 
+    async def create_conversation_session(
+        self, websocket: WebSocket, conversation_style="precise"
+    ):
+        await websocket.accept()
+        conversation_session = ConversationSession(conversation_style)
+        conversation_session.open()
+        while True:
+            try:
+                data = await websocket.receive_text()
+                response = await conversation_session.chat(data)
+                await websocket.send_text(response)
+            except Exception as e:
+                print(e)
+                break
+
     def setup_routes(self):
-        self.app.get(
-            "/models",
-            summary="Get available models",
-        )(self.get_available_models)
+        self.router = APIRouter()
+        self.router.add_api_route("/models", self.get_available_models)
+        self.router.add_websocket_route("/create", self.create_conversation_session)
+        self.app.include_router(self.router)
 
 
 app = ChatAPIApp().app

conversations/conversation_session.py

@@ -34,11 +34,11 @@ class ConversationSession:
     def close(self):
         self.event_loop.close()
 
-    def chat(self, prompt):
+    async def chat(self, prompt):
         logger.success(f"\n[User]: ", end="")
         logger.mesg(f"{prompt}")
         logger.success(f"[Bing]:")
-        self.event_loop.run_until_complete(self.connector.stream_chat(prompt=prompt))
+        return await self.connector.stream_chat(prompt=prompt)
 
 
 if __name__ == "__main__":

examples/chat_with_websocket.py

@@ -0,0 +1,19 @@
+import asyncio
+import websockets
+
+
+async def chat_with_websocket(url, message):
+    async with websockets.connect(url) as websocket:
+        await websocket.send(message)
+        response = await websocket.recv()
+        print(response)
+
+
+server = "localhost"
+port = 22222
+ws_url = f"ws://{server}:{port}/create"
+
+prompt = "Hello, who are you?"
+event_loop = asyncio.get_event_loop()
+event_loop.run_until_complete(chat_with_websocket(ws_url, prompt))
+event_loop.close()