streaming_callback.py

"""
Streaming Callback Handler for WebSocket
Streams LangChain agent responses in real-time to WebSocket clients.
"""
from langchain.callbacks.base import AsyncCallbackHandler
from typing import Any, Dict, List, Optional
import asyncio


class WebSocketStreamingCallback(AsyncCallbackHandler):
    """
    Custom callback handler to stream LangChain responses via WebSocket.
    """
    
    def __init__(self, ws_manager, session_id: str):
        """
        Args:
            ws_manager: WebSocket manager instance
            session_id: Target session ID
        """
        self.ws_manager = ws_manager
        self.session_id = session_id
        self.current_tool = None
    
    async def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Called when LLM starts."""
        await self.ws_manager.send_typing(self.session_id)
    
    async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Called when a new token is generated - stream it!"""
        await self.ws_manager.send_message(self.session_id, {
            "type": "response_chunk",
            "content": token
        })
    
    async def on_llm_end(self, response, **kwargs: Any) -> None:
        """Called when LLM finishes."""
        pass
    
    async def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> None:
        """Called when a tool starts executing."""
        tool_name = serialized.get("name", "unknown")
        self.current_tool = tool_name
        
        # Send tool execution notification
        tool_descriptions = {
            "planner": "Analyzing your request and preparing questions...",
            "cad_generator": "Generating your CAD model...",
            "modifier": "Modifying your design...",
            "cad_qa": "Looking up information..."
        }
        
        description = tool_descriptions.get(tool_name, f"Running {tool_name}...")
        await self.ws_manager.send_tool_executing(
            self.session_id,
            tool_name,
            description
        )
    
    async def on_tool_end(self, output: str, **kwargs: Any) -> None:
        """Called when a tool finishes."""
        # Optionally stream tool output
        if output and len(output) < 500:  # Only stream short outputs
            await self.ws_manager.send_message(self.session_id, {
                "type": "tool_output",
                "content": output,
                "tool": self.current_tool
            })
    
    async def on_tool_error(
        self, error: Exception, **kwargs: Any
    ) -> None:
        """Called when a tool errors."""
        await self.ws_manager.send_error(
            self.session_id,
            f"Tool error: {str(error)}"
        )
    
    async def on_agent_action(self, action, **kwargs: Any) -> None:
        """Called when agent decides on an action."""
        # Can be used to show agent's thinking
        pass
    
    async def on_agent_finish(self, finish, **kwargs: Any) -> None:
        """Called when agent finishes."""
        pass
    
    async def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
        """Called when chain errors."""
        await self.ws_manager.send_error(
            self.session_id,
            f"Processing error: {str(error)}"
        )