110 lines
3.9 KiB
Python
110 lines
3.9 KiB
Python
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)
|
||
|