websocket_manager.py

"""
WebSocket Manager for real-time communication with frontend.
Handles streaming responses token-by-token like GitHub Copilot.
"""
from fastapi import WebSocket, WebSocketDisconnect
from typing import Dict, Any, Optional
import json
import asyncio


class WebSocketManager:
    """
    Manages active WebSocket connections and streaming.
    """
    
    def __init__(self):
        # Store active connections: session_id -> WebSocket
        self.active_connections: Dict[str, WebSocket] = {}
    
    async def connect(self, session_id: str, websocket: WebSocket):
        """Accept and store a WebSocket connection."""
        await websocket.accept()
        self.active_connections[session_id] = websocket
        print(f"WebSocket connected: {session_id}")
    
    def disconnect(self, session_id: str):
        """Remove a WebSocket connection."""
        if session_id in self.active_connections:
            del self.active_connections[session_id]
            print(f"WebSocket disconnected: {session_id}")
    
    async def send_message(self, session_id: str, message: Dict[str, Any]):
        """Send a JSON message to a specific session."""
        if session_id in self.active_connections:
            try:
                await self.active_connections[session_id].send_json(message)
            except Exception as e:
                print(f"Error sending message to {session_id}: {e}")
                self.disconnect(session_id)
    
    async def stream_text(self, session_id: str, text: str, chunk_size: int = 3):
        """
        Stream text token-by-token to create typing effect.
        
        Args:
            session_id: Target session
            text: Full text to stream
            chunk_size: Characters per chunk (default 3 for smooth streaming)
        """
        if session_id not in self.active_connections:
            return
        
        # Split text into chunks
        for i in range(0, len(text), chunk_size):
            chunk = text[i:i + chunk_size]
            await self.send_message(session_id, {
                "type": "response_chunk",
                "content": chunk
            })
            await asyncio.sleep(0.01)  # Small delay for smooth streaming
    
    async def send_event(self, session_id: str, event_type: str, data: Dict[str, Any] = None):
        """
        Send a specific event to the client.
        
        Event types:
        - typing: Bot is thinking
        - tool_executing: A tool is running (CAD generation, etc.)
        - complete: Response is complete
        - error: An error occurred
        """
        message = {
            "type": event_type,
            "data": data or {}
        }
        await self.send_message(session_id, message)
    
    async def send_typing(self, session_id: str):
        """Indicate bot is thinking."""
        await self.send_event(session_id, "typing")
    
    async def send_tool_executing(self, session_id: str, tool_name: str, description: str = ""):
        """Indicate a tool is executing."""
        await self.send_event(session_id, "tool_executing", {
            "tool": tool_name,
            "description": description
        })
    
    async def send_complete(self, session_id: str, metadata: Dict[str, Any] = None):
        """Indicate response is complete."""
        await self.send_event(session_id, "complete", metadata or {})
    
    async def send_error(self, session_id: str, error_message: str):
        """Send error message."""
        await self.send_event(session_id, "error", {"message": error_message})
    
    async def stream_with_metadata(
        self, 
        session_id: str, 
        text: str, 
        intent: str = None,
        tool: str = None,
        metadata: Dict[str, Any] = None
    ):
        """
        Stream a complete response with metadata.
        
        This is the main method for streaming responses from the agent.
        """
        # Send typing indicator
        await self.send_typing(session_id)
        
        # Build metadata
        meta = metadata or {}
        if intent:
            meta["intent"] = intent
        if tool:
            meta["tool"] = tool
        
        # Stream the text
        if session_id not in self.active_connections:
            return
        
        # Send full chunks with metadata for each part
        words = text.split()
        current_chunk = ""
        
        for i, word in enumerate(words):
            current_chunk += word + " "
            
            # Send every 3-5 words for smooth but not too chatty streaming
            if (i + 1) % 4 == 0 or i == len(words) - 1:
                await self.send_message(session_id, {
                    "type": "response_chunk",
                    "content": current_chunk.strip() + " ",
                    "metadata": meta
                })
                current_chunk = ""
                await asyncio.sleep(0.02)


# Singleton instance
_ws_manager = None

def get_ws_manager() -> WebSocketManager:
    """Get or create the singleton WebSocket manager."""
    global _ws_manager
    if _ws_manager is None:
        _ws_manager = WebSocketManager()
    return _ws_manager