55 lines
2.1 KiB
Python
55 lines
2.1 KiB
Python
import asyncio
|
|
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)
|
|
|
|
async def _notify_state_handlers(self, new_state: WebAgentState):
|
|
"""Notify all state handlers of a state change."""
|
|
for handler in self._state_handlers:
|
|
await handler(new_state)
|
|
|
|
async def _notify_response_handlers(self, response: str):
|
|
"""Notify all response handlers of a response change."""
|
|
for handler in self._response_handlers:
|
|
await handler(response)
|
|
|
|
async def _set_response(self, response: str):
|
|
"""Set response and notify handlers."""
|
|
self.response = response
|
|
await self._notify_response_handlers(response)
|
|
|
|
async def approve_context(self):
|
|
"""Handle context approval with immediate state updates."""
|
|
# Change to INFERENCE state and notify
|
|
self._current_state = WebAgentState.INFERENCE
|
|
for handler in self._state_handlers:
|
|
await handler(WebAgentState.INFERENCE)
|
|
|
|
# Set response and notify
|
|
self.response = "<reasoning>test</reasoning>"
|
|
for handler in self._response_handlers:
|
|
await handler(self.response)
|
|
|
|
# Change to RESPONSE_APPROVAL state and notify
|
|
self._current_state = WebAgentState.RESPONSE_APPROVAL
|
|
for handler in self._state_handlers:
|
|
await handler(WebAgentState.RESPONSE_APPROVAL) |