diff --git a/readme.md b/readme.md index 1a5f807..1d1cb92 100644 --- a/readme.md +++ b/readme.md @@ -293,13 +293,9 @@ classDiagram -llm: LLMEngine -parser: ResponseParser -validator: XMLValidator - -io_buffer: IOBuffer - -system_prompt: str -action_schema: str #_compile_context() str - #_process_llm_response(response str) Optional~Command~ - #_update() void } class WorkingMemory { @@ -395,13 +391,9 @@ classDiagram -llm: LLMEngine -parser: ResponseParser -validator: XMLValidator - -io_buffer: IOBuffer - -system_prompt: str -action_schema: str #_compile_context() str - #_process_llm_response(response str) Optional~Command~ - #_update() void } class StandardAgent { @@ -410,7 +402,8 @@ classDiagram } class WebAgent { - -current_state: WebAgentState + +context: str + +response: str +WebAgent(model_path str, system_prompt str, action_schema str) +get_current_state() WebAgentState diff --git a/sia/base_agent.py b/sia/base_agent.py index 2c5ba2f..264bfd5 100644 --- a/sia/base_agent.py +++ b/sia/base_agent.py @@ -1,15 +1,12 @@ from abc import ABC, abstractmethod -from typing import Iterator, Callable, Optional, List +from typing import List import xml.etree.ElementTree as ET -from .command import Command from .llm_engine import LlmEngine from .system_metrics import SystemMetrics from .working_memory import WorkingMemory from .xml_validator import XMLValidator from .response_parser import ResponseParser -from .parse_error_entry import ParseErrorEntry -from .io_buffer import IOBuffer class BaseAgent(ABC): """ @@ -17,42 +14,24 @@ class BaseAgent(ABC): Provides core functionality for maintaining working memory, system metrics, and coordinating components for LLM inference. - - Private Attributes: - _working_memory: Collection of current entries - _metrics: System resource monitoring - _llm: LLM inference engine - _parser: XML response parser - _validator: XML response validator - _io_buffer: Input/output operations buffer - _system_prompt: System prompt template - _action_schema: XML schema for action validation """ def __init__(self, - model_path: str, - system_prompt: str, action_schema: str, - io_buffer: IOBuffer): + working_memory: WorkingMemory, + system_metrics: SystemMetrics, + llm: LlmEngine, + validator: XMLValidator, + parser: ResponseParser): """ Initialize agent with required components. - - Args: - model_path: Path to LLM model - system_prompt: System prompt template - action_schema: XML schema for actions - io_buffer: IO buffer implementation to use """ # Initialize components - self._working_memory = WorkingMemory() - self._metrics = SystemMetrics() - self._llm = LlmEngine(model_path) - self._validator = XMLValidator(action_schema) - self._io_buffer = io_buffer - self._parser = ResponseParser(io_buffer) - - # Store prompts - self._system_prompt = system_prompt + self._working_memory = working_memory + self._metrics = system_metrics + self._llm = llm + self._validator = validator + self._parser = parser self._action_schema = action_schema def __del__(self): @@ -68,22 +47,9 @@ class BaseAgent(ABC): Returns: str: Complete context as XML string """ - # Get usage details to include in context - context_size = 0 # TODO: Implement context size tracking - - # Get current system metrics - metrics_context = self._metrics.generate_context(context_size) - - # Get working memory entries memory_context = self._working_memory.generate_context() - - # Create root element - root = ET.Element("state") - - # Add metrics and memory entries - root.append(metrics_context) + context_size = len(memory_context) / 100 + context = self._metrics.generate_context(context_size) for entry in memory_context: - root.append(entry) - - # Convert to string with basic formatting - return ET.tostring(root, encoding="unicode") \ No newline at end of file + context.append(entry) + return ET.tostring(context, encoding="unicode") \ No newline at end of file diff --git a/sia/command.py b/sia/command.py index 883854c..c57b177 100644 --- a/sia/command.py +++ b/sia/command.py @@ -1,9 +1,8 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING -if TYPE_CHECKING: - from .working_memory import WorkingMemory - from .command_result import CommandResult +from .working_memory import WorkingMemory +from .command_result import CommandResult class Command(ABC): """ @@ -12,7 +11,7 @@ class Command(ABC): """ @abstractmethod - def execute(self, memory: 'WorkingMemory') -> 'CommandResult': + def execute(self, memory: WorkingMemory) -> CommandResult: """ Execute the command on the given working memory. diff --git a/sia/web_agent.py b/sia/web_agent.py new file mode 100644 index 0000000..c2ac35d --- /dev/null +++ b/sia/web_agent.py @@ -0,0 +1,143 @@ +from enum import Enum, auto +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. + + States: + UPDATE: Updating system metrics and working memory entries + CONTEXT_APPROVAL: Waiting for human approval of context + INFERENCE: Processing context through LLM + RESPONSE_APPROVAL: Waiting for human approval of LLM response + """ + UPDATE = auto() + CONTEXT_APPROVAL = auto() + INFERENCE = auto() + RESPONSE_APPROVAL = 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.context = "" + self.response = "" + self._system_prompt = system_prompt + self._current_state = WebAgentState.UPDATE + self._validation_error: Optional[str] = None + self._state_change_handlers: List[Callable[[WebAgentState], None]] = [] + self._response_change_handlers: List[Callable[[str], None]] = [] + self._command_result: Optional[CommandResult] = None + + + @property + def current_state(self) -> WebAgentState: + """Get the current state of the agent.""" + return self._current_state + + @property + def command_result(self) -> Optional[CommandResult]: + """Get the result of the last command execution.""" + return self._command_result + + 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_response_change_handler(self, handler: Callable[[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 _notify_state_change(self) -> None: + """Notify all handlers of state change.""" + for handler in self._state_change_handlers: + handler(self._current_state) + + def _notify_response_change(self) -> None: + """Notify all handlers of response change.""" + for handler in self._response_change_handlers: + handler(self._response) + + def _set_response(self, response: str) -> None: + """ + Set the current response and notify handlers. + + Args: + response: New response + """ + self._response = response + self._validation_error = self._validator.validate(self._response) + self._notify_response_change() + + def approve_context(self) -> None: + if self._current_state != WebAgentState.CONTEXT_APPROVAL: + raise Exception("Not in CONTEXT_APPROVAL state") + + self._set_response("") + self._current_state = WebAgentState.INFERENCE + self._notify_state_change() + + response_token_iter = self._llm.infer(self._system_prompt, self.context) + + for token in response_token_iter: + self._set_response(self._response + token) + + self._current_state = WebAgentState.RESPONSE_APPROVAL + self._notify_state_change() + + def approve_response(self) -> None: + if self._current_state != WebAgentState.RESPONSE_APPROVAL: + raise Exception("Not in RESPONSE_APPROVAL state") + + self._current_state = WebAgentState.UPDATE + self._notify_state_change() + + parse_result = self._parser.parse(self._response) + if isinstance(parse_result, Command): + result = parse_result.execute(self._working_memory) + self._command_result = result + else: + self._working_memory.add_entry(parse_result) + + self._working_memory.update() + self._context = self._compile_context() + + self._current_state = WebAgentState.CONTEXT_APPROVAL + self._notify_state_change() \ No newline at end of file diff --git a/test/base_agent_test.py b/test/base_agent_test.py index 2ff4731..12cd5e3 100644 --- a/test/base_agent_test.py +++ b/test/base_agent_test.py @@ -2,6 +2,7 @@ import unittest from datetime import datetime from unittest.mock import Mock, MagicMock, patch import xml.etree.ElementTree as ET +import logging from sia.web_io_buffer import WebIOBuffer from sia.working_memory import WorkingMemory @@ -47,38 +48,48 @@ class BaseAgentTest(unittest.TestCase): self.mock_metrics = Mock(spec=SystemMetrics) self.mock_validator = Mock(spec=XMLValidator) self.io_buffer = WebIOBuffer() + self.working_memory = WorkingMemory() - # Mock metrics context - metrics_elem = ET.Element("context") - metrics_elem.set("time", "2024-10-31T12:00:00Z") - metrics_elem.set("cpu", "10") - metrics_elem.set("gpu", "20") - metrics_elem.set("memory_used", "1000") - metrics_elem.set("memory_total", "2000") - metrics_elem.set("disk_used", "5000") - metrics_elem.set("disk_total", "10000") - metrics_elem.set("context", "50") - metrics_elem.set("stdin", "0") - self.mock_metrics.generate_context.return_value = metrics_elem + # Create parser with IO buffer + self.parser = ResponseParser(self.io_buffer) - # Create test agent with mocked components - with patch('sia.base_agent.LlmEngine') as mock_llm_class, \ - patch('sia.base_agent.SystemMetrics') as mock_metrics_class, \ - patch('sia.base_agent.XMLValidator') as mock_validator_class: + # Mock metrics context generator + def generate_context(context_size): + metrics_elem = ET.Element("context") + metrics_elem.set("time", "2024-10-31T12:00:00Z") + metrics_elem.set("cpu", "10") + metrics_elem.set("gpu", "20") + metrics_elem.set("memory_used", "1000") + metrics_elem.set("memory_total", "2000") + metrics_elem.set("disk_used", "5000") + metrics_elem.set("disk_total", "10000") + metrics_elem.set("context", str(int(context_size * 100))) + metrics_elem.set("stdin", "0") + return metrics_elem - mock_llm_class.return_value = self.mock_llm - mock_metrics_class.return_value = self.mock_metrics - mock_validator_class.return_value = self.mock_validator - - self.agent = TestBaseAgent( - model_path="/test/model", - system_prompt="test prompt", - action_schema="test schema", - io_buffer=self.io_buffer - ) + self.mock_metrics.generate_context.side_effect = generate_context + + # Create test agent with required components + self.agent = TestBaseAgent( + action_schema="test schema", + working_memory=self.working_memory, + system_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) + + # Capture logging output + self.log_records = [] + logging.getLogger().handlers = [] + logging.getLogger().addHandler(self.log_handler) + + def log_handler(self, record): + """Capture log records for testing.""" + self.log_records.append(record) def tearDown(self): """Clean up resources.""" @@ -87,13 +98,11 @@ class BaseAgentTest(unittest.TestCase): def test_initialization(self): """Test agent initialization and component setup.""" - self.assertIsInstance(self.agent._working_memory, WorkingMemory) + self.assertEqual(self.agent._working_memory, self.working_memory) self.assertEqual(self.agent._metrics, self.mock_metrics) self.assertEqual(self.agent._llm, self.mock_llm) self.assertEqual(self.agent._validator, self.mock_validator) - self.assertEqual(self.agent._io_buffer, self.io_buffer) - self.assertIsInstance(self.agent._parser, ResponseParser) - self.assertEqual(self.agent._system_prompt, "test prompt") + self.assertEqual(self.agent._parser, self.parser) self.assertEqual(self.agent._action_schema, "test schema") def test_cleanup(self): @@ -107,16 +116,17 @@ class BaseAgentTest(unittest.TestCase): # Parse context and verify structure root = ET.fromstring(context) - self.assertEqual(root.tag, "state") + self.assertEqual(root.tag, "context") - # Check metrics element - context_elem = root.find("context") - self.assertIsNotNone(context_elem) - self.assertEqual(context_elem.get("cpu"), "10") - self.assertEqual(context_elem.get("gpu"), "20") + # Check metrics + self.assertEqual(root.get("cpu"), "10") + self.assertEqual(root.get("gpu"), "20") + + # Check context size is 0 (empty memory) + self.assertEqual(root.get("context"), "0") # Check no memory entries - self.assertEqual(len(root.findall("*")), 1) # Only metrics element + self.assertEqual(len(list(root)), 0) def test_compile_context_with_entries(self): """Test context compilation with working memory entries.""" @@ -129,8 +139,8 @@ class BaseAgentTest(unittest.TestCase): # Parse and verify context root = ET.fromstring(context) - # Check metrics and entry are present - self.assertEqual(len(root.findall("*")), 2) # Metrics + reasoning + # Check context size reflects one entry + self.assertEqual(root.get("context"), "1") # Verify entry content reasoning_elem = root.find("reasoning") @@ -152,11 +162,12 @@ class BaseAgentTest(unittest.TestCase): context = self.agent._compile_context() root = ET.fromstring(context) - # Verify all entries present - self.assertEqual(len(root.findall("reasoning")), 3) + # Check context size reflects three entries + self.assertEqual(root.get("context"), "3") # Verify entry order maintained reasoning_elems = root.findall("reasoning") + self.assertEqual(len(reasoning_elems), 3) for i, elem in enumerate(reasoning_elems): self.assertEqual(elem.get("id"), f"id-{i}") self.assertEqual(elem.text, f"") \ No newline at end of file diff --git a/test/web_agent_test.py b/test/web_agent_test.py new file mode 100644 index 0000000..a64172b --- /dev/null +++ b/test/web_agent_test.py @@ -0,0 +1,233 @@ +import unittest +from datetime import datetime +from unittest.mock import Mock, MagicMock, patch, call +import xml.etree.ElementTree as ET +from typing import Iterator + +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.""" + # Create mocks + self.mock_llm = Mock(spec=LlmEngine) + self.mock_metrics = Mock(spec=SystemMetrics) + 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 metrics context generator + def generate_context(context_size): + metrics_elem = ET.Element("context") + metrics_elem.set("time", "2024-10-31T12:00:00Z") + metrics_elem.set("cpu", "10") + metrics_elem.set("gpu", "20") + metrics_elem.set("memory_used", "1000") + metrics_elem.set("memory_total", "2000") + metrics_elem.set("disk_used", "5000") + metrics_elem.set("disk_total", "10000") + metrics_elem.set("context", str(int(context_size * 100))) + metrics_elem.set("stdin", "0") + return metrics_elem + + self.mock_metrics.generate_context.side_effect = generate_context + + # Create agent with all components + self.agent = WebAgent( + system_prompt="test prompt", + action_schema="test schema", + working_memory=self.working_memory, + system_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): + """Track response changes for verification.""" + self.response_changes.append(response) + + def test_initialization(self): + """Test initial state of web agent.""" + self.assertEqual(self.agent.current_state, WebAgentState.UPDATE) + self.assertEqual(self.agent.context, "") + self.assertEqual(self.agent.response, "") + 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.""" + # Add handlers + self.agent.add_state_change_handler(self.state_change_handler) + self.agent.add_response_change_handler(self.response_change_handler) + + # Verify handlers added + 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 raises error.""" + 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.""" + # Register handlers + self.agent.add_state_change_handler(self.state_change_handler) + self.agent.add_response_change_handler(self.response_change_handler) + + # Setup LLM response + def mock_infer(prompt: str, context: str) -> Iterator[str]: + yield "test reasoning" + self.mock_llm.infer.side_effect = mock_infer + + # Set initial state + self.agent._current_state = WebAgentState.CONTEXT_APPROVAL + + # Approve context + self.agent.approve_context() + + # Verify state transitions + expected_states = [ + WebAgentState.INFERENCE, + WebAgentState.RESPONSE_APPROVAL + ] + self.assertEqual(self.state_changes, expected_states) + + # Verify response updates + self.assertEqual(len(self.response_changes), 2) # Empty + final response + self.assertEqual(self.response_changes[0], "") + self.assertEqual(self.response_changes[1], "test reasoning") + + def test_response_approval_flow_command(self): + """Test response approval flow with command execution.""" + # Register handlers + self.agent.add_state_change_handler(self.state_change_handler) + + # Setup command result + command_result = CommandResult.success() + mock_command = MockCommand(command_result) + + # Mock parser to return command + self.parser.parse = Mock(return_value=mock_command) + + # Set initial state and response + self.agent._current_state = WebAgentState.RESPONSE_APPROVAL + self.agent._response = "" + + # Approve response + self.agent.approve_response() + + # Verify command executed + self.assertTrue(mock_command.executed) + + # Verify state transitions + expected_states = [ + WebAgentState.UPDATE, + WebAgentState.CONTEXT_APPROVAL + ] + self.assertEqual(self.state_changes, expected_states) + + # Verify command result stored + self.assertEqual(self.agent.command_result, command_result) + + def test_response_approval_flow_entry(self): + """Test response approval flow with entry creation.""" + # Register handlers + self.agent.add_state_change_handler(self.state_change_handler) + + # Set initial state and response + self.agent._current_state = WebAgentState.RESPONSE_APPROVAL + self.agent._response = "test reasoning" + + # Approve response + self.agent.approve_response() + + # Verify entry added to working memory + 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") + + # Verify state transitions + 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.""" + # Register handlers + self.agent.add_response_change_handler(self.response_change_handler) + + # Mock validator to fail + error_message = "Invalid XML" + self.mock_validator.validate.return_value = error_message + + # Set response + self.agent._set_response("") + + # Verify validation error stored + self.assertEqual(self.agent._validation_error, error_message) + + # Verify response updated + self.assertEqual(self.response_changes, [""]) + + def tearDown(self): + """Clean up resources.""" + if hasattr(self, 'agent'): + del self.agent \ No newline at end of file