SIA/sia/web_agent.py

from enum import Enum, auto
from typing import Callable, List, Optional

from .base_agent import BaseAgent
from .command import Command
from .command_result import CommandResult
from .llm_engine import LlmEngine
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator

class WebAgentState(Enum):
    """
    States for the web agent state machine.
    
    States:
        UPDATE: Updating system metrics and working memory entries
        CONTEXT_APPROVAL: Waiting for human approval of context
        INFERENCE: Processing context through LLM
        RESPONSE_APPROVAL: Waiting for human approval of LLM response
    """
    UPDATE = auto()
    CONTEXT_APPROVAL = auto()
    INFERENCE = auto()
    RESPONSE_APPROVAL = auto()

class WebAgent(BaseAgent):
    """
    Agent implementation for interactive web interface.
    
    Uses a state machine to allow human intervention between steps.
    Broadcasts state changes to registered handlers.
    """
    
    def __init__(self,
                system_prompt: str,
                action_schema: str, 
                working_memory: WorkingMemory,
                system_metrics: SystemMetrics,
                llm: LlmEngine,
                validator: XMLValidator,
                parser: ResponseParser):
        """
        Initialize web agent with required components.
        """
        super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
        self.response = ""
        self._system_prompt = system_prompt
        self._current_state = WebAgentState.CONTEXT_APPROVAL
        self._validation_error: Optional[str] = None
        self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
        self._response_change_handlers: List[Callable[[str], None]] = []
        self._command_result: Optional[CommandResult] = None
        self.context = self._compile_context()

    @property
    def current_state(self) -> WebAgentState:
        """Get the current state of the agent."""
        return self._current_state
    
    @property
    def command_result(self) -> Optional[CommandResult]:
        """Get the result of the last command execution."""
        return self._command_result
        
    def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
        """
        Add a callback for state changes.
        
        Args:
            handler: Function to call with new state
        """
        if handler not in self._state_change_handlers:
            self._state_change_handlers.append(handler)
    
    def add_response_change_handler(self, handler: Callable[[str], None]) -> None:
        """
        Add a callback for response changes.

        Args:
            handler: Function to call with new response
        """
        if handler not in self._response_change_handlers:
            self._response_change_handlers.append(handler)
            
    def _notify_state_change(self) -> None:
        """Notify all handlers of state change."""
        for handler in self._state_change_handlers:
            handler(self._current_state)
    
    def _notify_response_change(self) -> None:
        """Notify all handlers of response change."""
        for handler in self._response_change_handlers:
            handler(self.response)
    
    def _set_response(self, response: str) -> None:
        """Set response and notify handlers."""
        self.response = response
        self._validation_error = self._validator.validate(response)
        for handler in self._response_change_handlers:
            try:
                handler(response)
            except Exception as e:
                print(f"Error in response handler: {e}")

    def _set_state(self, new_state: WebAgentState) -> None:
        """
        Set the current state and notify handlers.

        Args:
            new_state: New state to set
        """
        self._current_state = new_state
        self._notify_state_change()
            
    async def approve_context(self) -> None:
        if self._current_state != WebAgentState.CONTEXT_APPROVAL:
            error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})"
            raise Exception(error_msg)
                
        self._set_state(WebAgentState.INFERENCE)
        self._set_response("")
        
        response_token_iter = self._llm.infer(self._system_prompt, self.context)
        response = ""
        try:
            for token in response_token_iter:
                response += token
                await self._set_response(response)
                print(f"{token}")
        except Exception as e:
            print(f"Error during inference: {e}")
            self._set_state(WebAgentState.CONTEXT_APPROVAL)
            return
        
        await self._set_response(response)
        self._set_state(WebAgentState.RESPONSE_APPROVAL)

    async def _set_response(self, response: str) -> None:
        """Set response and notify handlers asynchronously."""
        self.response = response
        self._validation_error = self._validator.validate(response)
        for handler in self._response_change_handlers:
            try:
                await handler(response)
            except Exception as e:
                print(f"Error in response handler: {e}")

    def approve_response(self) -> None:
        if self._current_state != WebAgentState.RESPONSE_APPROVAL:
            raise Exception("Not in RESPONSE_APPROVAL state")

        self._set_state(WebAgentState.UPDATE)

        parse_result = self._parser.parse(self.response)
        if isinstance(parse_result, Command):
            result = parse_result.execute(self._working_memory)
            self._command_result = result
        else:
            self._working_memory.add_entry(parse_result)
        
        self._working_memory.update()
        self.context = self._compile_context()

        self._set_state(WebAgentState.CONTEXT_APPROVAL)