init

new-planet-backend/app/api/v1/websocket.py

@@ -0,0 +1,109 @@
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
+from typing import Dict, Set, Optional
+import json
+from app.services.chat_service import chat_service
+from app.db.session import AsyncSessionLocal
+from app.api.deps import get_current_user
+from app.core.security import decode_token
+
+router = APIRouter()
+
+# Хранилище активных соединений
+active_connections: Dict[str, Set[WebSocket]] = {}
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: Dict[str, Set[WebSocket]] = {}
+    
+    async def connect(self, websocket: WebSocket, user_id: str):
+        await websocket.accept()
+        if user_id not in self.active_connections:
+            self.active_connections[user_id] = set()
+        self.active_connections[user_id].add(websocket)
+    
+    def disconnect(self, websocket: WebSocket, user_id: str):
+        if user_id in self.active_connections:
+            self.active_connections[user_id].discard(websocket)
+            if not self.active_connections[user_id]:
+                del self.active_connections[user_id]
+    
+    async def send_personal_message(self, message: str, websocket: WebSocket):
+        await websocket.send_text(message)
+    
+    async def broadcast_to_user(self, user_id: str, message: str):
+        if user_id in self.active_connections:
+            for connection in self.active_connections[user_id]:
+                await connection.send_text(message)
+
+
+manager = ConnectionManager()
+
+
+async def get_user_from_token(token: str) -> Optional[str]:
+    """Получить user_id из токена"""
+    payload = decode_token(token)
+    if payload:
+        return payload.get("sub")
+    return None
+
+
+@router.websocket("/ws/chat")
+async def websocket_chat(websocket: WebSocket):
+    """WebSocket endpoint для чата с ИИ"""
+    # Получаем токен из query параметров
+    token = websocket.query_params.get("token")
+    if not token:
+        await websocket.close(code=1008, reason="Token required")
+        return
+    
+    # Проверка токена
+    user_id = await get_user_from_token(token)
+    if not user_id:
+        await websocket.close(code=1008, reason="Unauthorized")
+        return
+    
+    await manager.connect(websocket, user_id)
+    
+    try:
+        async with AsyncSessionLocal() as db:
+            while True:
+                data = await websocket.receive_text()
+                message_data = json.loads(data)
+                
+                # Создаем запрос для chat_service
+                from app.schemas.ai import ChatRequest
+                request = ChatRequest(
+                    message=message_data.get("message", ""),
+                    conversation_id=message_data.get("conversation_id")
+                )
+                
+                # Получаем ответ от ИИ
+                try:
+                    response = await chat_service.chat(
+                        db=db,
+                        user_id=user_id,
+                        request=request
+                    )
+                    
+                    # Отправляем ответ клиенту
+                    await manager.send_personal_message(
+                        json.dumps({
+                            "type": "message",
+                            "response": response.response,
+                            "conversation_id": response.conversation_id,
+                            "tokens_used": response.tokens_used
+                        }),
+                        websocket
+                    )
+                except Exception as e:
+                    await manager.send_personal_message(
+                        json.dumps({
+                            "type": "error",
+                            "message": str(e)
+                        }),
+                        websocket
+                    )
+    except WebSocketDisconnect:
+        manager.disconnect(websocket, user_id)
+