54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
from typing import Optional, List, Callable
|
|
from sia.web_agent import WebAgentState
|
|
|
|
class MockWebAgent:
|
|
"""Mock WebAgent for testing."""
|
|
def __init__(self):
|
|
self._current_state: Optional[WebAgentState] = None
|
|
self.context: Optional[str] = None
|
|
self.response: Optional[str] = None
|
|
self._validation_error: Optional[str] = None
|
|
self._state_handlers: List[Callable] = []
|
|
self._response_handlers: List[Callable] = []
|
|
|
|
@property
|
|
def current_state(self):
|
|
return self._current_state
|
|
|
|
def add_state_change_handler(self, handler):
|
|
self._state_handlers.append(handler)
|
|
|
|
def add_response_change_handler(self, handler):
|
|
self._response_handlers.append(handler)
|
|
|
|
def set_input(self, input_text):
|
|
pass
|
|
|
|
def _notify_state_handlers(self, new_state: WebAgentState):
|
|
"""Notify all state handlers of a state change."""
|
|
for handler in self._state_handlers:
|
|
handler(new_state)
|
|
|
|
def _notify_response_handlers(self, response: str):
|
|
"""Notify all response handlers of a response change."""
|
|
for handler in self._response_handlers:
|
|
handler(response)
|
|
|
|
def _set_response(self, response: str):
|
|
"""Set response and notify handlers."""
|
|
self.response = response
|
|
self._notify_response_handlers(response)
|
|
|
|
def approve_context(self):
|
|
"""Handle context approval with immediate state updates."""
|
|
# Change to INFERENCE state and notify
|
|
self._current_state = WebAgentState.INFERENCE
|
|
self._notify_state_handlers(WebAgentState.INFERENCE)
|
|
|
|
# Set response and notify
|
|
self._set_response("<reasoning>test</reasoning>")
|
|
|
|
# Change to RESPONSE_APPROVAL state and notify
|
|
self._current_state = WebAgentState.RESPONSE_APPROVAL
|
|
self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL)
|