backend/app/api/websocket.py - master - ArtCoder/jarvis

backend/app/api/websocket.py
186 строк7 KB
Art'Coder
Initial version of project
14 июл 2026, 21:06
14 июл 2026, 21:068b763d5
# websocket.py
from fastapi import WebSocket, WebSocketDisconnect
import json
import asyncio
import uuid
from typing import Dict, Optional
import logging
from datetime import datetime

from ..core.jarvis_brain import JarvisBrain
from ..core.voice_engine import JarvisVoiceEngine
from ..core.action_executor import ActionExecutor
from ..core.exceptions import handle_exception
from ..core.config import settings

logger = logging.getLogger(__name__)

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.connection_data: Dict[str, Dict] = {}
        self.brain = JarvisBrain()
        self.voice = JarvisVoiceEngine()
        self.executor = ActionExecutor()
        self._initialized = False
        self._initialize_components()
    
    def _initialize_components(self):
        try:
            asyncio.create_task(self.brain.initialize())
            self._initialized = True
            logger.info("WebSocket manager initialized")
        except Exception as e:
            logger.error(f"Component initialization failed: {e}")
            self._initialized = False
    
    async def connect(self, websocket: WebSocket, client_id: Optional[str] = None):
        await websocket.accept()
        
        if not client_id:
            client_id = str(uuid.uuid4())
        
        self.active_connections[client_id] = websocket
        self.connection_data[client_id] = {
            "connected_at": datetime.now().isoformat(),
            "messages": 0,
            "last_activity": datetime.now().isoformat()
        }
        
        logger.info(f"Client {client_id} connected")
        
        try:
            await self.send_message(client_id, {
                "type": "system",
                "content": f"J.A.R.V.I.S. {settings.system_version} онлайн. Я на связи, сэр.",
                "voice": True,
                "timestamp": datetime.now().isoformat()
            })
        except Exception as e:
            logger.error(f"Welcome message error: {e}")
    
    def disconnect(self, client_id: str):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
        if client_id in self.connection_data:
            del self.connection_data[client_id]
        logger.info(f"Client {client_id} disconnected")
    
    async def send_message(self, client_id: str, data: Dict):
        if client_id not in self.active_connections:
            return
        
        try:
            if data.get("voice", False) and self.voice:
                content = data.get("content", "")
                if content:
                    asyncio.create_task(self.voice.speak_async(content))
            
            await self.active_connections[client_id].send_json(data)
            
            if client_id in self.connection_data:
                self.connection_data[client_id]["messages"] += 1
                self.connection_data[client_id]["last_activity"] = datetime.now().isoformat()
                
        except WebSocketDisconnect:
            self.disconnect(client_id)
        except Exception as e:
            logger.error(f"Send message error: {e}")
    
    async def handle_message(self, client_id: str, message: str):
        if not self._initialized:
            await self.send_message(client_id, {
                "type": "error",
                "content": "Система не инициализирована, сэр.",
                "voice": True
            })
            return
        
        try:
            start_time = datetime.now()
            
            if self.brain:
                result = await self.brain.process_query(message)
                response_text = result.get("response", "Не удалось обработать запрос, сэр.")
            else:
                response_text = "Мозг системы недоступен, сэр."
            
            action_result = await self._detect_and_execute_action(message)
            if action_result and action_result.get("success"):
                response_text += f" {action_result['result']}"
            
            await self.send_message(client_id, {
                "type": "response",
                "content": response_text,
                "voice": True,
                "timestamp": datetime.now().isoformat(),
                "processing_time": (datetime.now() - start_time).total_seconds()
            })
            
        except Exception as e:
            logger.error(f"Message handling error: {e}")
            error_response = handle_exception(e, "WebSocket.handle_message")
            await self.send_message(client_id, {
                "type": "error",
                "content": error_response.get("message", "Произошла ошибка, сэр."),
                "voice": True
            })
    
    async def _detect_and_execute_action(self, message: str) -> Optional[Dict]:
        if not self.executor:
            return None
        
        message_lower = message.lower()
        
        try:
            if "открой" in message_lower or "open" in message_lower:
                if "ютуб" in message_lower or "youtube" in message_lower:
                    return await self.executor.execute("open_browser", {"url": "https://youtube.com"})
                elif "гугл" in message_lower or "google" in message_lower:
                    return await self.executor.execute("open_browser", {"url": "https://google.com"})
            
            elif "погода" in message_lower:
                return await self.executor.execute("get_weather", {})
            elif "время" in message_lower:
                return await self.executor.execute("get_time", {})
            elif "статус" in message_lower or "system" in message_lower:
                return await self.executor.execute("get_system_info", {})
            elif "скрин" in message_lower or "screenshot" in message_lower:
                return await self.executor.execute("take_screenshot", {})
            elif "громкость" in message_lower or "volume" in message_lower:
                if "выше" in message_lower or "up" in message_lower:
                    return await self.executor.execute("control_volume", {"direction": "up"})
                elif "ниже" in message_lower or "down" in message_lower:
                    return await self.executor.execute("control_volume", {"direction": "down"})
            
            elif "запомни" in message_lower or "reminder" in message_lower:
                text = message_lower.replace("запомни", "").replace("reminder", "").strip()
                if text:
                    return await self.executor.execute("set_reminder", {"text": text})
            
        except Exception as e:
            logger.error(f"Action detection error: {e}")
        
        return None

manager = ConnectionManager()

async def websocket_endpoint(websocket: WebSocket, client_id: str = "tony_stark"):
    await websocket.accept()
    await manager.connect(websocket, client_id)
    
    try:
        while True:
            try:
                data = await websocket.receive_text()
                await manager.handle_message(client_id, data)
            except WebSocketDisconnect:
                manager.disconnect(client_id)
                break
            except Exception as e:
                logger.error(f"WebSocket receive error: {e}")
                break
    except Exception as e:
        logger.error(f"WebSocket endpoint error: {e}")
    finally:
        manager.disconnect(client_id)