196 lines
7.2 KiB
Python
196 lines
7.2 KiB
Python
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 .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,
|
|
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._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()
|
|
|
|
@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(f"{self._system_prompt}\n{self._action_schema}", self.context)
|
|
response = ""
|
|
for token in response_token_iter:
|
|
response += token
|
|
self._set_response(response)
|
|
print(f"{token}")
|
|
self._set_state(WebAgentState.RESPONSE_APPROVAL)
|
|
|
|
def _approve_response_thread(self) -> None:
|
|
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
|
|
else:
|
|
self._working_memory.add_entry(parse_result)
|
|
|
|
self._working_memory.update()
|
|
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() |