from enum import Enum, auto from threading import Thread, Lock from typing import Callable, List, Optional from .base_agent import BaseAgent from .command import Command from .command_result import CommandResult from .iteration_logger import IterationLogger 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. """ UPDATE = auto() """Updating system metrics and working memory entries""" CONTEXT_APPROVAL = auto() """Waiting for human approval of context""" INFERENCE = auto() """Processing context through LLM""" RESPONSE_APPROVAL = auto() """Waiting for human approval of LLM response""" STOPPED = 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, metrics: SystemMetrics, llm: LlmEngine, validator: XMLValidator, parser: ResponseParser, iteration_logger: IterationLogger, ): """ Initialize web agent with required components. """ super().__init__( system_prompt, action_schema, working_memory, metrics, llm, validator, parser ) self._response = "" self._state = WebAgentState.CONTEXT_APPROVAL self._validation_error: Optional[str] = None self._state_change_handlers: List[Callable[[WebAgentState], None]] = [] self._context_change_handlers: List[Callable[[str], None]] = [] self._response_change_handlers: List[Callable[[str, str], None]] = [] self._command_result: Optional[CommandResult] = None self._context = self._compile_context() self._state_lock = Lock() self._iteration_logger = iteration_logger @property def state(self) -> WebAgentState: """Get the current state of the agent.""" return self._state @property def command_result(self) -> Optional[CommandResult]: """Get the result of the last command execution.""" return self._command_result @property def validation_error(self) -> Optional[str]: """Get the validation error, if any.""" return self._validation_error @property def response(self) -> str: """Get the current response.""" return self._response @property def context(self) -> str: """Get the current context.""" return self._context 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_context_change_handler(self, handler: Callable[[str], None]) -> None: """ Add a callback for context changes. Args: handler: Function to call with new context """ if handler not in self._context_change_handlers: self._context_change_handlers.append(handler) def add_response_change_handler(self, handler: Callable[[str, 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 modify_context(self, context: str) -> None: with self._state_lock: if self._state != WebAgentState.CONTEXT_APPROVAL: error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})" raise Exception(error_msg) Thread(target=self._set_context, args=(context,)).start() def modify_response(self, response: str) -> None: started = False try: self._state_lock.acquire() if self._state != WebAgentState.RESPONSE_APPROVAL: error_msg = f"Not in RESPONSE_APPROVAL state (current: {self._state})" raise Exception(error_msg) Thread(target=self._set_response, args=(response, self._state_lock)).start() started = True finally: if not started: self._state_lock.release() def approve_context(self) -> None: with self._state_lock: if self._state != WebAgentState.CONTEXT_APPROVAL: error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})" raise Exception(error_msg) self._set_state(WebAgentState.INFERENCE) Thread(target=self._approve_context_thread).start() def approve_response(self) -> None: with self._state_lock: if self._state != WebAgentState.RESPONSE_APPROVAL: raise Exception("Not in RESPONSE_APPROVAL state") self._set_state(WebAgentState.UPDATE) Thread(target=self._approve_response_thread).start() def _approve_context_thread(self) -> None: self._set_response("") response_token_iter = self._llm.infer(self.system_prompt, self.context) response = "" for token in response_token_iter: response += token self._set_response(response) print(f"{token}", end='') self._set_state(WebAgentState.RESPONSE_APPROVAL) print() def _approve_response_thread(self) -> None: self._iteration_logger.log_iteration(self._context, self._response) parse_result = self._parser.parse(self._response) if isinstance(parse_result, Command): result = parse_result.execute(self._working_memory) self._command_result = result if result.should_stop: self._set_state(WebAgentState.STOPPED) return self._working_memory.update() else: parse_result.update() self._working_memory.update() self._working_memory.add_entry(parse_result) self._set_context(self._compile_context()) self._set_state(WebAgentState.CONTEXT_APPROVAL) def _set_state(self, state) -> None: """Notify all handlers of state change.""" self._state = state for handler in self._state_change_handlers: handler(self._state) def _set_context(self, context) -> None: """Notify all handlers of context change.""" self._context = context for handler in self._context_change_handlers: handler(context) def _set_response(self, response: str, lock: Optional[Lock] = None) -> None: """Set response and notify handlers.""" try: self._response = response validation_error = self._validator.validate(response) self._validation_error = validation_error for handler in self._response_change_handlers: handler(response, validation_error) finally: if lock is not None: lock.release()