195 lines
7.6 KiB
Python
195 lines
7.6 KiB
Python
from datetime import datetime
|
|
from typing import Iterator
|
|
from unittest.mock import Mock, MagicMock, patch, call
|
|
import asyncio
|
|
import time
|
|
import unittest
|
|
import xml.etree.ElementTree as ET
|
|
|
|
from sia.web_io_buffer import WebIOBuffer
|
|
from sia.working_memory import WorkingMemory
|
|
from sia.system_metrics import SystemMetrics
|
|
from sia.llm_engine import LlmEngine
|
|
from sia.xml_validator import XMLValidator
|
|
from sia.response_parser import ResponseParser
|
|
from sia.web_agent import WebAgent, WebAgentState
|
|
from sia.command import Command
|
|
from sia.command_result import CommandResult
|
|
from sia.reasoning_entry import ReasoningEntry
|
|
from sia.parse_error_entry import ParseErrorEntry
|
|
from sia.entry import Entry
|
|
|
|
class MockCommand(Command):
|
|
"""Mock command for testing execution flow."""
|
|
def __init__(self, result: CommandResult):
|
|
self.result = result
|
|
self.executed = False
|
|
|
|
def execute(self, memory: WorkingMemory) -> CommandResult:
|
|
self.executed = True
|
|
return self.result
|
|
|
|
class WebAgentTest(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test cases with mocked components."""
|
|
self.mock_llm = Mock(spec=LlmEngine)
|
|
self.mock_validator = Mock(spec=XMLValidator)
|
|
self.io_buffer = WebIOBuffer()
|
|
self.working_memory = WorkingMemory()
|
|
|
|
# Mock validator to pass by default
|
|
self.mock_validator.validate.return_value = None
|
|
|
|
# Create parser with IO buffer
|
|
self.parser = ResponseParser(self.io_buffer)
|
|
|
|
# Mock system metrics
|
|
self.mock_metrics = Mock(spec=SystemMetrics)
|
|
self.mock_metrics.get_metrics.return_value = {
|
|
"timestamp": "2024-10-31T12:00:00Z",
|
|
"cpu": 10,
|
|
"gpu": 20,
|
|
"memory_used": 1000,
|
|
"memory_total": 2000,
|
|
"disk_used": 5000,
|
|
"disk_total": 10000
|
|
}
|
|
|
|
# Mock LLM response generator
|
|
def mock_infer(prompt: str, context: str):
|
|
yield "<reasoning>test reasoning</reasoning>"
|
|
|
|
self.mock_llm.infer.side_effect = mock_infer
|
|
self.mock_llm.token_count.return_value = 100
|
|
self.mock_llm.token_limit.return_value = 1000
|
|
|
|
# Create agent with all components
|
|
self.agent = WebAgent(
|
|
system_prompt="test prompt",
|
|
action_schema="test schema",
|
|
working_memory=self.working_memory,
|
|
metrics=self.mock_metrics,
|
|
llm=self.mock_llm,
|
|
validator=self.mock_validator,
|
|
parser=self.parser
|
|
)
|
|
|
|
# Set timestamp for entries
|
|
self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0)
|
|
|
|
# Handler tracking
|
|
self.state_changes = []
|
|
self.response_changes = []
|
|
|
|
def state_change_handler(self, state: WebAgentState):
|
|
"""Track state changes for verification."""
|
|
self.state_changes.append(state)
|
|
|
|
def response_change_handler(self, response: str, validation_error: str):
|
|
"""Track response changes for verification."""
|
|
self.response_changes.append(response)
|
|
|
|
def test_initialization(self):
|
|
"""Test initial state of web agent."""
|
|
self.assertEqual(self.agent.state, WebAgentState.CONTEXT_APPROVAL)
|
|
self.assertEqual(self.agent.response, "")
|
|
self.assertIsNotNone(self.agent.context)
|
|
self.assertIsNone(self.agent.validation_error)
|
|
self.assertEqual(len(self.agent._state_change_handlers), 0)
|
|
self.assertEqual(len(self.agent._response_change_handlers), 0)
|
|
|
|
def test_handler_registration(self):
|
|
"""Test adding state and response change handlers."""
|
|
self.agent.add_state_change_handler(self.state_change_handler)
|
|
self.agent.add_response_change_handler(self.response_change_handler)
|
|
|
|
self.assertEqual(len(self.agent._state_change_handlers), 1)
|
|
self.assertEqual(len(self.agent._response_change_handlers), 1)
|
|
|
|
# Adding same handler twice should not duplicate
|
|
self.agent.add_state_change_handler(self.state_change_handler)
|
|
self.agent.add_response_change_handler(self.response_change_handler)
|
|
|
|
self.assertEqual(len(self.agent._state_change_handlers), 1)
|
|
self.assertEqual(len(self.agent._response_change_handlers), 1)
|
|
|
|
def test_approve_context_state_error(self):
|
|
"""Test approving context in wrong state."""
|
|
self.agent._state = WebAgentState.UPDATE
|
|
|
|
with self.assertRaises(Exception) as context:
|
|
self.agent.approve_context()
|
|
self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
|
|
|
|
def test_approve_response_state_error(self):
|
|
"""Test approving response in wrong state raises error."""
|
|
with self.assertRaises(Exception) as context:
|
|
self.agent.approve_response()
|
|
self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
|
|
|
|
def test_context_approval_flow(self):
|
|
"""Test complete context approval flow with state transitions."""
|
|
self.agent.add_state_change_handler(self.state_change_handler)
|
|
self.agent.add_response_change_handler(self.response_change_handler)
|
|
|
|
self.agent.approve_context()
|
|
|
|
expected_states = [
|
|
WebAgentState.INFERENCE,
|
|
WebAgentState.RESPONSE_APPROVAL
|
|
]
|
|
self.assertEqual(self.state_changes, expected_states)
|
|
self.assertEqual(self.response_changes[-1], "<reasoning>test reasoning</reasoning>")
|
|
|
|
def test_response_approval_flow_command(self):
|
|
"""Test response approval flow with command execution."""
|
|
self.agent.add_state_change_handler(self.state_change_handler)
|
|
|
|
command_result = CommandResult.success()
|
|
mock_command = MockCommand(command_result)
|
|
self.parser.parse = Mock(return_value=mock_command)
|
|
|
|
self.agent._state = WebAgentState.RESPONSE_APPROVAL
|
|
self.agent._response = "<stop/>"
|
|
|
|
self.agent.approve_response()
|
|
|
|
self.assertTrue(mock_command.executed)
|
|
|
|
expected_states = [
|
|
WebAgentState.UPDATE,
|
|
WebAgentState.CONTEXT_APPROVAL
|
|
]
|
|
self.assertEqual(self.state_changes, expected_states)
|
|
self.assertEqual(self.agent.command_result, command_result)
|
|
|
|
def test_response_approval_flow_entry(self):
|
|
"""Test response approval flow with entry creation."""
|
|
self.agent.add_state_change_handler(self.state_change_handler)
|
|
|
|
self.agent._state = WebAgentState.RESPONSE_APPROVAL
|
|
self.agent._set_response("<reasoning>test reasoning</reasoning>")
|
|
|
|
self.agent.approve_response()
|
|
time.sleep(1)
|
|
|
|
self.assertEqual(self.working_memory.get_entries_count(), 1)
|
|
entry = self.working_memory.get_entries()[0]
|
|
self.assertIsInstance(entry, ReasoningEntry)
|
|
self.assertEqual(entry.content, "test reasoning")
|
|
|
|
expected_states = [
|
|
WebAgentState.UPDATE,
|
|
WebAgentState.CONTEXT_APPROVAL
|
|
]
|
|
self.assertEqual(self.state_changes, expected_states)
|
|
|
|
def test_response_validation(self):
|
|
"""Test response validation during response setting."""
|
|
self.agent.add_response_change_handler(self.response_change_handler)
|
|
error_message = "Invalid XML"
|
|
self.mock_validator.validate.return_value = error_message
|
|
self.agent._set_response("<invalid>")
|
|
self.assertEqual(self.agent.validation_error, error_message)
|
|
self.assertEqual(self.response_changes, ["<invalid>"])
|