Spaces:
Sleeping
Sleeping
File size: 1,421 Bytes
43b064c 761bdee 43b064c 761bdee 43b064c f2fee05 761bdee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | import json
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException
from src.services import ConnectionService, ChatService
from src.utils import logger
class ChatController:
def __init__(self):
self.connection_service = ConnectionService
self.chat_service = ChatService
self.router = APIRouter(prefix="/chat", tags=["chat"])
self.router.websocket("/ws")(self.websocket_endpoint)
async def websocket_endpoint(self, websocket: WebSocket):
logger.info("WebSocket connection attempt.")
try:
async with self.connection_service() as connection_service:
await connection_service.connect(websocket)
while True:
data = await websocket.receive_text()
messages = []
parsed_data = json.loads(data)
messages = parsed_data.get("messages", [])
async with self.chat_service() as chat_service:
async for response_message in chat_service.create_response_message(messages):
await connection_service.send_personal_message(response_message, websocket)
except WebSocketDisconnect:
connection_service.disconnect(websocket)
except Exception as e:
logger.error(e)
raise HTTPException(status_code=500, detail=str(e))
|