hermes-agent/api_endpoint/websocket_logger.py

#!/usr/bin/env python3
"""
WebSocket Logger Client

Simple client for sending agent events to the logging server via WebSocket.
Used by the agent to log events in real-time during execution.
"""

import json
import asyncio
from typing import Dict, Any, Optional
from datetime import datetime
import websockets


class WebSocketLogger:
    """
    Client for logging agent events via WebSocket.
    
    Usage:
        logger = WebSocketLogger("unique-session-id")
        await logger.connect()
        await logger.log_query("What is Python?", model="gpt-4")
        await logger.log_api_call(call_number=1)
        await logger.log_response(call_number=1, content="Python is...")
        await logger.disconnect()
    """
    
    def __init__(
        self, 
        session_id: str,
        server_url: str = "ws://localhost:8000/ws",
        enabled: bool = True
    ):
        """
        Initialize WebSocket logger.
        
        Args:
            session_id: Unique identifier for this agent session
            server_url: WebSocket server URL (default: ws://localhost:8000/ws)
            enabled: Whether logging is enabled (default: True)
        """
        self.session_id = session_id
        self.server_url = server_url
        self.enabled = enabled
        self.websocket: Optional[websockets.WebSocketClientProtocol] = None
        self.connected = False
        self.reconnect_count = 0  # Track reconnections for debugging
    
    async def connect(self):
        """
        Connect to the WebSocket logging server.
        
        Establishes WebSocket connection and sends initial session_start event.
        If connection fails, gracefully disables logging (agent continues normally).
        """
        if not self.enabled:
            return
        
        try:
            # Establish WebSocket connection to the server
            # Use VERY LONG ping intervals to avoid timeout during long tool execution
            # The event loop is blocked during tool execution, so we can't process pings
            # Setting to very large values (1 hour) effectively disables it
            self.websocket = await websockets.connect(
                self.server_url,
                ping_interval=3600,      # 1 hour - effectively disabled (event loop blocked anyway)
                ping_timeout=3600,       # 1 hour timeout for pong response
                close_timeout=10,        # Timeout for close handshake
                max_size=10 * 1024 * 1024,  # 10MB max message size (for large raw_results)
                open_timeout=10          # Timeout for initial connection
            )
            self.connected = True
            print(f"✅ Connected to logging server (ping/pong: 3600s intervals): {self.server_url}")
            
            # Send initial session_start event
            # This tells the server to create a new SessionLogger for this session
            await self._send_event("session_start", {
                "session_id": self.session_id,
                "start_time": datetime.now().isoformat()
            })
            
        except Exception as e:
            # Connection failed - disable logging but don't crash the agent
            print(f"⚠️ Failed to connect to logging server: {e}")
            print(f"   Logging will be disabled for this session.")
            self.enabled = False
            self.connected = False
    
    async def disconnect(self):
        """Disconnect from the WebSocket server."""
        if self.websocket and self.connected:
            try:
                await self.websocket.close()
                self.connected = False
                print(f"✅ Disconnected from logging server")
            except Exception as e:
                print(f"⚠️ Error disconnecting: {e}")
    
    async def _send_event(self, event_type: str, data: Dict[str, Any]):
        """
        Send an event to the logging server.
        
        This is the core method that sends all events via WebSocket.
        Creates a standardized message format and handles acknowledgments.
        
        Args:
            event_type: Type of event (query, api_call, response, tool_call, tool_result, error, complete)
            data: Event data dictionary containing event-specific information
        """
        # Safety check: Don't send if logging is disabled
        if not self.enabled:
            return
        
        # Auto-reconnect if connection was lost
        if not self.connected or not self.websocket:
            try:
                self.reconnect_count += 1
                print(f"🔄 Reconnecting to logging server (attempt #{self.reconnect_count})...")
                await self.connect()
                print(f"✅ Reconnected successfully!")
            except Exception as e:
                print(f"⚠️ Failed to reconnect: {e}")
                self.enabled = False  # Disable logging after failed reconnect
                return
        
        try:
            # Create standardized message structure
            # All events follow this format for consistent server-side handling
            message = {
                "session_id": self.session_id,      # Links event to specific agent session
                "event_type": event_type,            # Identifies what kind of event this is
                "data": data                         # Event-specific payload
            }
            
            # Send message as JSON string over WebSocket
            await self.websocket.send(json.dumps(message))
            
            # Wait for server acknowledgment (with 1 second timeout)
            # This ensures the server received and processed the event
            try:
                response = await asyncio.wait_for(
                    self.websocket.recv(),
                    timeout=1.0
                )
                # Server sends back: {"status": "logged", "session_id": "...", "event_type": "..."}
                # We don't need to process it, just confirms receipt
            except asyncio.TimeoutError:
                # No response within 1 second - that's okay, continue anyway
                # Server might be busy or network slow, but event was likely sent
                pass
                
        except Exception as e:
            # Log error but don't crash - graceful degradation
            # Agent should continue working even if logging fails
            error_str = str(e)
            
            # Check if connection was closed (error 1011 = keepalive ping timeout)
            if "1011" in error_str or "closed" in error_str.lower():
                print(f"⚠️ WebSocket connection closed: {error_str}")
                self.connected = False  # Mark as disconnected
                # Don't try to send more events - connection is dead
            else:
                print(f"⚠️ Error sending event to logging server: {e}")
            # Don't disable entirely or try to reconnect - just continue with logging disabled
    
    # Convenience methods for specific event types
    
    async def log_query(
        self, 
        query: str, 
        model: str = None,
        toolsets: list = None
    ):
        """
        Log a user query (the question/task given to the agent).
        
        This is typically the first event in a session after connection.
        Captures what the user asked and which model/tools will be used.
        """
        await self._send_event("query", {
            "query": query,          # The user's question/instruction
            "model": model,          # Which AI model is being used
            "toolsets": toolsets     # Which tool categories are enabled
        })
    
    async def log_api_call(
        self,
        call_number: int,
        message_count: int = None,
        has_tools: bool = None
    ):
        """
        Log an API call to the AI model.
        
        Called right before sending a request to the model (OpenAI/Anthropic/etc).
        Helps track how many API calls are being made and conversation length.
        """
        await self._send_event("api_call", {
            "call_number": call_number,      # Sequential number (1, 2, 3...)
            "message_count": message_count,  # How many messages in conversation so far
            "has_tools": has_tools          # Whether tools are available to the model
        })
    
    async def log_response(
        self,
        call_number: int,
        content: str = None,
        has_tool_calls: bool = False,
        tool_call_count: int = 0,
        duration: float = None
    ):
        """
        Log an assistant response from the AI model.
        
        Called after receiving a response from the API.
        Captures what the model said and whether it wants to use tools.
        """
        await self._send_event("response", {
            "call_number": call_number,          # Which API call this response is from
            "content": content,                   # What the model said (text response)
            "has_tool_calls": has_tool_calls,    # Did model request tool execution?
            "tool_call_count": tool_call_count,  # How many tools does it want to call?
            "duration": duration                  # How long the API call took (seconds)
        })
    
    async def log_tool_call(
        self,
        call_number: int,
        tool_index: int,
        tool_name: str,
        parameters: Dict[str, Any],
        tool_call_id: str = None
    ):
        """
        Log a tool call (before executing the tool).
        
        Captures which tool is being called and with what parameters.
        This happens BEFORE the tool runs, so no results yet.
        """
        await self._send_event("tool_call", {
            "call_number": call_number,      # Which API call requested this tool
            "tool_index": tool_index,        # Which tool in the sequence (if multiple)
            "tool_name": tool_name,          # Name of tool (e.g., "web_search", "web_extract")
            "parameters": parameters,        # Arguments passed to the tool (e.g., {"query": "Python", "limit": 5})
            "tool_call_id": tool_call_id    # Unique ID to link call with result
        })
    
    async def log_tool_result(
        self,
        call_number: int,
        tool_index: int,
        tool_name: str,
        result: str = None,
        error: str = None,
        duration: float = None,
        tool_call_id: str = None,
        raw_result: str = None  # NEW: Full untruncated result for verification
    ):
        """
        Log a tool result (output from tool execution).
        
        Captures both a truncated preview (for UI display) and the full raw result
        (for verification and debugging). This is especially important for web tools
        where you want to see what was scraped vs what the LLM processed.
        
        Args:
            call_number: Which API call this tool was part of
            tool_index: Which tool in the sequence (1st, 2nd, etc.)
            tool_name: Name of the tool that was executed
            result: Tool output (will be truncated to 1000 chars for preview)
            error: Error message if tool failed
            duration: How long the tool took to execute (seconds)
            tool_call_id: Unique ID linking this result to the tool call
            raw_result: NEW - Full untruncated result for verification/debugging
        """
        await self._send_event("tool_result", {
            "call_number": call_number,
            "tool_index": tool_index,
            "tool_name": tool_name,
            "result": result[:1000] if result else None,  # Truncated preview (1000 chars max)
            "raw_result": raw_result,  # NEW: Full result - can be 100KB+ for web scraping
            "error": error,
            "duration": duration,
            "tool_call_id": tool_call_id
        })
    
    async def log_error(
        self,
        error_message: str,
        call_number: int = None
    ):
        """
        Log an error that occurred during agent execution.
        
        Captures exceptions, API failures, or other issues.
        """
        await self._send_event("error", {
            "error_message": error_message,  # Description of what went wrong
            "call_number": call_number       # Which API call caused the error (if applicable)
        })
    
    async def log_complete(
        self,
        final_response: str = None,
        total_calls: int = None,
        completed: bool = True
    ):
        """
        Log session completion (final event before disconnecting).
        
        Marks the end of the agent's execution and provides summary info.
        """
        await self._send_event("complete", {
            "final_response": final_response[:500] if final_response else None,  # Truncated summary of final answer
            "total_calls": total_calls,      # How many API calls were made total
            "completed": completed           # Did it complete successfully? (true/false)
        })


# Synchronous wrapper for convenience
class SyncWebSocketLogger:
    """
    Synchronous wrapper around WebSocketLogger.
    
    For use in synchronous code - creates an event loop internally.
    """
    
    def __init__(self, session_id: str, server_url: str = "ws://localhost:8000/ws", enabled: bool = True):
        self.logger = WebSocketLogger(session_id, server_url, enabled)
        self.loop = None
    
    def connect(self):
        """Connect to server (synchronous)."""
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.loop.run_until_complete(self.logger.connect())
    
    def disconnect(self):
        """Disconnect from server (synchronous)."""
        if self.loop:
            self.loop.run_until_complete(self.logger.disconnect())
            self.loop.close()
    
    def _run_async(self, coro):
        """
        Run an async coroutine synchronously.
        
        Bridge between sync code (agent) and async code (WebSocket).
        Uses event loop to execute async operations in sync context.
        """
        if self.loop and self.loop.is_running():
            # Already in event loop, just await
            asyncio.create_task(coro)
        else:
            # Run in current loop
            if self.loop:
                self.loop.run_until_complete(coro)
    
    def log_query(self, query: str, model: str = None, toolsets: list = None):
        self._run_async(self.logger.log_query(query, model, toolsets))
    
    def log_api_call(self, call_number: int, message_count: int = None, has_tools: bool = None):
        self._run_async(self.logger.log_api_call(call_number, message_count, has_tools))
    
    def log_response(self, call_number: int, content: str = None, has_tool_calls: bool = False, 
                    tool_call_count: int = 0, duration: float = None):
        self._run_async(self.logger.log_response(call_number, content, has_tool_calls, 
                                                 tool_call_count, duration))
    
    def log_tool_call(self, call_number: int, tool_index: int, tool_name: str, 
                     parameters: Dict[str, Any], tool_call_id: str = None):
        self._run_async(self.logger.log_tool_call(call_number, tool_index, tool_name, 
                                                  parameters, tool_call_id))
    
    def log_tool_result(self, call_number: int, tool_index: int, tool_name: str,
                       result: str = None, error: str = None, duration: float = None,
                       tool_call_id: str = None, raw_result: str = None):
        self._run_async(self.logger.log_tool_result(call_number, tool_index, tool_name,
                                                    result, error, duration, tool_call_id, raw_result))
    
    def log_error(self, error_message: str, call_number: int = None):
        self._run_async(self.logger.log_error(error_message, call_number))
    
    def log_complete(self, final_response: str = None, total_calls: int = None, completed: bool = True):
        self._run_async(self.logger.log_complete(final_response, total_calls, completed))