340 lines
13 KiB
Python
340 lines
13 KiB
Python
import unittest
|
|
from datetime import datetime
|
|
from typing import Dict, Iterator, Optional
|
|
from unittest.mock import Mock, MagicMock, patch, call
|
|
import asyncio
|
|
import time
|
|
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.response_parser import ResponseParser
|
|
from sia.iteration_logger import IterationLogger
|
|
from sia.web_agent import WebAgent, AgentState
|
|
from sia.command import Command
|
|
from sia.command_result import CommandResult
|
|
from sia.entry.reasoning_entry import ReasoningEntry
|
|
from sia.entry.parse_error_entry import ParseErrorEntry
|
|
from sia.entry import Entry
|
|
from sia.response_buffer import ResponseBuffer
|
|
from pathlib import Path
|
|
|
|
class MockLlmEngine:
|
|
"""Mock LLM engine for testing."""
|
|
|
|
def __init__(self, output: str = "<reasoning>test reasoning</reasoning>"):
|
|
self.output = output
|
|
self.token_count_retval = 100
|
|
self.token_limit_retval = 1000
|
|
self.infer_called = False
|
|
self.stop_requested = False
|
|
|
|
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: callable = lambda: False) -> Iterator[str]:
|
|
"""Mock implementation of infer that yields specified output."""
|
|
self.infer_called = True
|
|
|
|
for char in self.output:
|
|
if should_stop():
|
|
break
|
|
yield char
|
|
|
|
def token_count(self, system_prompt: str, main_context: str) -> int:
|
|
"""Mock implementation of token_count."""
|
|
return self.token_count_retval
|
|
|
|
def token_limit(self) -> int:
|
|
"""Mock implementation of token_limit."""
|
|
return self.token_limit_retval
|
|
|
|
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."""
|
|
# Mock components
|
|
self.mock_llms = {
|
|
'default': MockLlmEngine(),
|
|
'alternative': MockLlmEngine("<reasoning>alternative reasoning</reasoning>")
|
|
}
|
|
self.io_buffer = WebIOBuffer()
|
|
self.working_memory = WorkingMemory()
|
|
self.work_dir = Path("/tmp") # Use temp directory for tests
|
|
|
|
# Create parser with IO buffer
|
|
self.parser = ResponseParser(self.work_dir, 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",
|
|
"memory_used": 1000,
|
|
"memory_total": 2000,
|
|
"disk_used": 5000,
|
|
"disk_total": 10000
|
|
}
|
|
|
|
# Mock iteration logger
|
|
self.mock_iteration_logger = Mock(spec=IterationLogger)
|
|
|
|
# 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,
|
|
llms=self.mock_llms,
|
|
parser=self.parser,
|
|
iteration_logger=self.mock_iteration_logger
|
|
)
|
|
|
|
# Handler tracking
|
|
self.llm_state_changes = []
|
|
self.context_changes = []
|
|
|
|
def llm_change_handler(self, llm_name: str, state: AgentState):
|
|
"""Track LLM state changes for verification."""
|
|
self.llm_state_changes.append((llm_name, state))
|
|
|
|
def context_change_handler(self, context: str, generated: bool):
|
|
"""Track context changes for verification."""
|
|
self.context_changes.append((context, generated))
|
|
|
|
def test_initialization(self):
|
|
"""Test initial state of web agent."""
|
|
self.assertIsNotNone(self.agent.context)
|
|
self.assertIsNone(self.agent.command_result)
|
|
self.assertIsNone(self.agent.validation_error)
|
|
|
|
# Check LLM states
|
|
llm_states = self.agent.llms
|
|
self.assertEqual(len(llm_states), 2)
|
|
self.assertEqual(llm_states['default'], AgentState.IDLE)
|
|
self.assertEqual(llm_states['alternative'], AgentState.IDLE)
|
|
|
|
# Check response buffer
|
|
self.assertIsInstance(self.agent.response_buffer, ResponseBuffer)
|
|
|
|
def test_run_inference(self):
|
|
"""Test running inference."""
|
|
self.agent.add_llm_change_handler(self.llm_change_handler)
|
|
|
|
# Mock the response_buffer to avoid calling append_text
|
|
with patch.object(self.agent.response_buffer, 'append_text') as mock_append:
|
|
# Run inference
|
|
self.agent.run_inference('default')
|
|
|
|
# Check state changes
|
|
self.assertEqual(len(self.llm_state_changes), 2) # IDLE -> INFERENCE -> IDLE
|
|
self.assertEqual(self.llm_state_changes[0], ('default', AgentState.INFERENCE))
|
|
self.assertEqual(self.llm_state_changes[1], ('default', AgentState.IDLE))
|
|
|
|
# Verify LLM was called
|
|
self.assertTrue(self.mock_llms['default'].infer_called)
|
|
|
|
# Verify append_text was called for each character in the output
|
|
expected_calls = [call(char) for char in "<reasoning>test reasoning</reasoning>"]
|
|
mock_append.assert_has_calls(expected_calls)
|
|
|
|
def test_run_inference_invalid_llm(self):
|
|
"""Test running inference with invalid LLM name."""
|
|
with self.assertRaises(ValueError):
|
|
self.agent.run_inference('nonexistent')
|
|
|
|
def test_run_inference_busy_llm(self):
|
|
"""Test running inference on a busy LLM."""
|
|
# Set LLM state to INFERENCE
|
|
with patch.object(self.agent, '_llm_states', {'default': AgentState.INFERENCE}):
|
|
with self.assertRaises(RuntimeError):
|
|
self.agent.run_inference('default')
|
|
|
|
def test_stop_inference(self):
|
|
"""Test stopping inference."""
|
|
# Create a mock that allows us to control the should_stop function
|
|
mock_engine = MockLlmEngine("<reasoning>this should be interrupted</reasoning>")
|
|
|
|
# Create a new agent with our mock engine to avoid affecting other tests
|
|
test_agent = WebAgent(
|
|
system_prompt="test prompt",
|
|
action_schema="test schema",
|
|
working_memory=WorkingMemory(),
|
|
metrics=self.mock_metrics,
|
|
llms={'test': mock_engine},
|
|
parser=self.parser,
|
|
iteration_logger=self.mock_iteration_logger
|
|
)
|
|
|
|
# Patch the append_text method to avoid errors
|
|
with patch.object(test_agent.response_buffer, 'append_text'):
|
|
# Run inference in a separate thread
|
|
import threading
|
|
def run_inference():
|
|
try:
|
|
test_agent.run_inference('test')
|
|
except Exception as e:
|
|
print(f"Exception in inference thread: {e}")
|
|
|
|
thread = threading.Thread(target=run_inference)
|
|
thread.start()
|
|
|
|
# Wait a bit for inference to start
|
|
time.sleep(0.1)
|
|
|
|
# Stop inference
|
|
test_agent.stop_inference('test')
|
|
|
|
# Wait for thread to complete
|
|
thread.join(timeout=1)
|
|
self.assertFalse(thread.is_alive())
|
|
|
|
def test_modify_context(self):
|
|
"""Test modifying context."""
|
|
# Create a separate agent for this test
|
|
test_agent = WebAgent(
|
|
system_prompt="test prompt",
|
|
action_schema="test schema",
|
|
working_memory=WorkingMemory(),
|
|
metrics=self.mock_metrics,
|
|
llms=self.mock_llms,
|
|
parser=self.parser,
|
|
iteration_logger=self.mock_iteration_logger
|
|
)
|
|
|
|
# Mock the problematic method calls
|
|
with patch.object(test_agent, '_set_llm_state'):
|
|
# Add a handler
|
|
context_handler = Mock()
|
|
test_agent.add_context_change_handler(context_handler)
|
|
|
|
# Update context
|
|
new_context = "<context>modified context</context>"
|
|
test_agent.modify_context(new_context)
|
|
|
|
# Check context was updated
|
|
self.assertEqual(test_agent.context, new_context)
|
|
|
|
# Check context change handler was called
|
|
context_handler.assert_called_once_with(new_context, False)
|
|
|
|
def test_modify_context_generated(self):
|
|
"""Test modifying context with generated flag."""
|
|
# Create a separate agent for this test
|
|
test_agent = WebAgent(
|
|
system_prompt="test prompt",
|
|
action_schema="test schema",
|
|
working_memory=WorkingMemory(),
|
|
metrics=self.mock_metrics,
|
|
llms=self.mock_llms,
|
|
parser=self.parser,
|
|
iteration_logger=self.mock_iteration_logger
|
|
)
|
|
|
|
# Mock the problematic method calls
|
|
with patch.object(test_agent, '_set_llm_state'):
|
|
# Add a handler
|
|
context_handler = Mock()
|
|
test_agent.add_context_change_handler(context_handler)
|
|
|
|
# Update context with generated=True
|
|
new_context = "<context>generated context</context>"
|
|
test_agent.modify_context(new_context, True)
|
|
|
|
# Check context change handler was called with generated=True
|
|
context_handler.assert_called_once_with(new_context, True)
|
|
|
|
def test_approve_response_command(self):
|
|
"""Test approving response that parses to a command."""
|
|
# Create a mock command result
|
|
command_result = CommandResult("Command executed", True, False)
|
|
mock_command = MockCommand(command_result)
|
|
|
|
# Mock parser to return our command
|
|
with patch.object(self.parser, 'parse', return_value=mock_command):
|
|
# Set response buffer content
|
|
self.agent.response_buffer.set_text("<stop/>")
|
|
|
|
# Approve response
|
|
self.agent.approve_response()
|
|
|
|
# Verify command was executed
|
|
self.assertTrue(mock_command.executed)
|
|
|
|
# Check command result was stored
|
|
self.assertEqual(self.agent._command_result, command_result)
|
|
|
|
# Verify iteration was logged
|
|
self.mock_iteration_logger.log_iteration.assert_called_once()
|
|
|
|
def test_approve_response_entry(self):
|
|
"""Test approving response that parses to an entry."""
|
|
# Mock parser to return a reasoning entry
|
|
timestamp = datetime.now()
|
|
with patch('datetime.datetime') as mock_dt:
|
|
mock_dt.now.return_value = timestamp
|
|
|
|
# Mock parser to return an entry instead of a command
|
|
entry = ReasoningEntry("test-id", "test reasoning")
|
|
with patch.object(self.parser, 'parse', return_value=entry):
|
|
# Set response buffer content
|
|
self.agent.response_buffer.set_text("<reasoning>test reasoning</reasoning>")
|
|
|
|
# Approve response
|
|
self.agent.approve_response()
|
|
|
|
# Check entry was added to working memory
|
|
added_entry = self.working_memory.get_entry("test-id")
|
|
self.assertIsNotNone(added_entry)
|
|
self.assertEqual(added_entry.content, "test reasoning")
|
|
|
|
# Verify iteration was logged
|
|
self.mock_iteration_logger.log_iteration.assert_called_once()
|
|
|
|
def test_working_memory_integration(self):
|
|
"""Test integration with working memory updates."""
|
|
# Create a separate agent and working memory for this test
|
|
test_memory = WorkingMemory()
|
|
test_agent = WebAgent(
|
|
system_prompt="test prompt",
|
|
action_schema="test schema",
|
|
working_memory=test_memory,
|
|
metrics=self.mock_metrics,
|
|
llms=self.mock_llms,
|
|
parser=self.parser,
|
|
iteration_logger=self.mock_iteration_logger
|
|
)
|
|
|
|
# Mock modify_context to avoid errors
|
|
with patch.object(test_agent, 'modify_context') as mock_modify:
|
|
# Add entries to memory
|
|
for i in range(3):
|
|
entry = ReasoningEntry(f"id-{i}", f"reasoning {i}")
|
|
test_memory.add_entry(entry)
|
|
|
|
# Verify modify_context was called once for each entry
|
|
self.assertEqual(mock_modify.call_count, 3)
|
|
# Verify all calls had generated=True
|
|
for call_args in mock_modify.call_args_list:
|
|
self.assertTrue(call_args[0][1]) # Second arg is 'generated'
|
|
|
|
def test_get_output(self):
|
|
"""Test getting LLM output."""
|
|
# Directly set the response buffer content instead of running inference
|
|
self.agent.response_buffer.set_text("<reasoning>test reasoning</reasoning>")
|
|
|
|
# Get output
|
|
output = self.agent.response_buffer.get_text()
|
|
|
|
# Check output
|
|
self.assertEqual(output, "<reasoning>test reasoning</reasoning>")
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main() |