import asyncio import json import logging from typing import Optional from fastapi import Query, Request, HTTPException, Depends, APIRouter, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse from fastapi.security import HTTPAuthorizationCredentials from pydantic import BaseModel, ConfigDict from broadcaster import Broadcast from models.chat import ChatCompletionRequest from models.user import User from services.openai_service.openai_service import generate_completion from auth.security import get_current_user from config.messages import SuccessResponse from redis import Redis logger = logging.getLogger(__name__) chat_router = APIRouter() redis_client = Redis(host='localhost', port=6379, db=0, decode_responses=True) @chat_router.websocket("/ws") async def chat_irc_endpoint(websocket: WebSocket): """WebSocket endpoint for real-time chat interactions""" global connected_users # Backend de broadcast con Redis broadcast = Broadcast("redis://localhost:6379") logger.info("New WebSocket connection attempt") logger.debug(f"WebSocket query params: {websocket.query_params}") current_user = await get_current_user( HTTPAuthorizationCredentials( scheme="Bearer", credentials=websocket.query_params.get("token", "") ) ) if not current_user: await websocket.close(code=1008) # Policy Violation return await websocket.accept() logger.info(f"User {current_user.email} connected to WebSocket chat.") username = None async def reader(): """Lee mensajes desde Redis y los manda al WebSocket""" logger.info("Iniciando lector de mensajes") async with broadcast.subscribe(channel="chat") as subscriber: logger.info("Subscribed to Redis channel 'chat'") logger.debug("Starting message read loop") async for event in subscriber: try: logger.debug(f"Broadcasting message to WebSocket: {event.message}") await websocket.send_text(event.message) except Exception: logger.error("Error sending message to WebSocket, closing connection.") break await broadcast.connect() reader_task = asyncio.create_task(reader()) try: while True: try: data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0) payload = json.loads(data) event_type = payload.get("type") logger.debug(f"Received event: {event_type} with payload: {payload}") if event_type == "join": logger.info(f"User {payload['username']} joined the chat.") data = { "mail": current_user.email, "username": payload["username"] } redis_client.sadd("connected_users", json.dumps(data)) response = {"type": "join", "username": payload["username"]} elif event_type == "message": logger.debug(f"Message from {payload['username']}: {payload['message']}") message_username = payload["username"] message = payload["message"] response = {"type": "message", "username": message_username, "message": message} elif event_type == "leave": logger.info(f"User {payload['username']} left the chat.") message_username = payload["username"] users = redis_client.smembers("connected_users") for user in users: user_data = json.loads(user) if user_data["username"] == message_username: redis_client.srem("connected_users", user) break response = {"type": "leave", "username": message_username} await websocket.close() break elif event_type == "ai_message": logger.debug(f"IA Message from {payload['username']}: {payload['messages']}") message_username = payload["username"] messages = payload["messages"] response_content = await generate_completion(messages, current_user) response = {"type": "message", "username": "IAKlein", "message": response_content} elif event_type == "mention": logger.debug(f"Mention to {payload['username']}") mention_username = payload["username"] logger.debug(f"Mention username: {mention_username}, Current username: {username}") await broadcast.publish(channel="chat", message=json.dumps({"type": "mentioned", "username": mention_username})) continue elif event_type == "pong": logger.debug(f"Received pong from user {current_user.email}") continue # Publicar en Redis await broadcast.publish(channel="chat", message=json.dumps(response)) except asyncio.TimeoutError: await websocket.send_text(json.dumps({"type": "ping"})) except WebSocketDisconnect: logger.info(f"User {current_user.email} disconnected from WebSocket chat.") users = redis_client.smembers("connected_users") for user in users: user_data = json.loads(user) if user_data["mail"] == current_user.email: redis_client.srem("connected_users", user) break response = {"type": "leave", "username": user_data["username"]} await broadcast.publish(channel="chat", message=json.dumps(response)) finally: reader_task.cancel() await broadcast.disconnect() @chat_router.get("/users") async def get_connected_users(q: str = Optional[Query(...)], _: User = Depends(get_current_user)): """Get a list of connected users (solo local al worker)""" # return {"users": [user.username for user in connected_users if q.lower() in user.username.lower()]} all_users = redis_client.smembers("connected_users") all_users = [json.loads(user)["username"] for user in all_users] if q is None or q.strip() == "": return {"users": all_users} filtered_users = [user for user in all_users if q.lower() in user.lower()] return {"users": filtered_users} @chat_router.get("/onlines") async def get_online_user_count(_: User = Depends(get_current_user)): """Get the count of online users (solo local al worker)""" all_users = redis_client.smembers("connected_users") return {"count": len(all_users)}