src/controllers/_chat_controller.py

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))