Added web agent

This commit is contained in:
2024-11-01 18:25:37 +01:00
parent f48b87f037
commit aadb841eb8
6 changed files with 448 additions and 103 deletions

View File

@@ -293,13 +293,9 @@ classDiagram
-llm: LLMEngine -llm: LLMEngine
-parser: ResponseParser -parser: ResponseParser
-validator: XMLValidator -validator: XMLValidator
-io_buffer: IOBuffer
-system_prompt: str
-action_schema: str -action_schema: str
#_compile_context() str #_compile_context() str
#_process_llm_response(response str) Optional~Command~
#_update() void
} }
class WorkingMemory { class WorkingMemory {
@@ -395,13 +391,9 @@ classDiagram
-llm: LLMEngine -llm: LLMEngine
-parser: ResponseParser -parser: ResponseParser
-validator: XMLValidator -validator: XMLValidator
-io_buffer: IOBuffer
-system_prompt: str
-action_schema: str -action_schema: str
#_compile_context() str #_compile_context() str
#_process_llm_response(response str) Optional~Command~
#_update() void
} }
class StandardAgent { class StandardAgent {
@@ -410,7 +402,8 @@ classDiagram
} }
class WebAgent { class WebAgent {
-current_state: WebAgentState +context: str
+response: str
+WebAgent(model_path str, system_prompt str, action_schema str) +WebAgent(model_path str, system_prompt str, action_schema str)
+get_current_state() WebAgentState +get_current_state() WebAgentState

View File

@@ -1,15 +1,12 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Iterator, Callable, Optional, List from typing import List
import xml.etree.ElementTree as ET import xml.etree.ElementTree as ET
from .command import Command
from .llm_engine import LlmEngine from .llm_engine import LlmEngine
from .system_metrics import SystemMetrics from .system_metrics import SystemMetrics
from .working_memory import WorkingMemory from .working_memory import WorkingMemory
from .xml_validator import XMLValidator from .xml_validator import XMLValidator
from .response_parser import ResponseParser from .response_parser import ResponseParser
from .parse_error_entry import ParseErrorEntry
from .io_buffer import IOBuffer
class BaseAgent(ABC): class BaseAgent(ABC):
""" """
@@ -17,42 +14,24 @@ class BaseAgent(ABC):
Provides core functionality for maintaining working memory, system metrics, Provides core functionality for maintaining working memory, system metrics,
and coordinating components for LLM inference. 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, def __init__(self,
model_path: str,
system_prompt: str,
action_schema: str, action_schema: str,
io_buffer: IOBuffer): working_memory: WorkingMemory,
system_metrics: SystemMetrics,
llm: LlmEngine,
validator: XMLValidator,
parser: ResponseParser):
""" """
Initialize agent with required components. 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 # Initialize components
self._working_memory = WorkingMemory() self._working_memory = working_memory
self._metrics = SystemMetrics() self._metrics = system_metrics
self._llm = LlmEngine(model_path) self._llm = llm
self._validator = XMLValidator(action_schema) self._validator = validator
self._io_buffer = io_buffer self._parser = parser
self._parser = ResponseParser(io_buffer)
# Store prompts
self._system_prompt = system_prompt
self._action_schema = action_schema self._action_schema = action_schema
def __del__(self): def __del__(self):
@@ -68,22 +47,9 @@ class BaseAgent(ABC):
Returns: Returns:
str: Complete context as XML string 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() memory_context = self._working_memory.generate_context()
context_size = len(memory_context) / 100
# Create root element context = self._metrics.generate_context(context_size)
root = ET.Element("state")
# Add metrics and memory entries
root.append(metrics_context)
for entry in memory_context: for entry in memory_context:
root.append(entry) context.append(entry)
return ET.tostring(context, encoding="unicode")
# Convert to string with basic formatting
return ET.tostring(root, encoding="unicode")

View File

@@ -1,9 +1,8 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
if TYPE_CHECKING: from .working_memory import WorkingMemory
from .working_memory import WorkingMemory from .command_result import CommandResult
from .command_result import CommandResult
class Command(ABC): class Command(ABC):
""" """
@@ -12,7 +11,7 @@ class Command(ABC):
""" """
@abstractmethod @abstractmethod
def execute(self, memory: 'WorkingMemory') -> 'CommandResult': def execute(self, memory: WorkingMemory) -> CommandResult:
""" """
Execute the command on the given working memory. Execute the command on the given working memory.

143
sia/web_agent.py Normal file
View File

@@ -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()

View File

@@ -2,6 +2,7 @@ import unittest
from datetime import datetime from datetime import datetime
from unittest.mock import Mock, MagicMock, patch from unittest.mock import Mock, MagicMock, patch
import xml.etree.ElementTree as ET import xml.etree.ElementTree as ET
import logging
from sia.web_io_buffer import WebIOBuffer from sia.web_io_buffer import WebIOBuffer
from sia.working_memory import WorkingMemory from sia.working_memory import WorkingMemory
@@ -47,8 +48,13 @@ class BaseAgentTest(unittest.TestCase):
self.mock_metrics = Mock(spec=SystemMetrics) self.mock_metrics = Mock(spec=SystemMetrics)
self.mock_validator = Mock(spec=XMLValidator) self.mock_validator = Mock(spec=XMLValidator)
self.io_buffer = WebIOBuffer() self.io_buffer = WebIOBuffer()
self.working_memory = WorkingMemory()
# Mock metrics context # 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 = ET.Element("context")
metrics_elem.set("time", "2024-10-31T12:00:00Z") metrics_elem.set("time", "2024-10-31T12:00:00Z")
metrics_elem.set("cpu", "10") metrics_elem.set("cpu", "10")
@@ -57,29 +63,34 @@ class BaseAgentTest(unittest.TestCase):
metrics_elem.set("memory_total", "2000") metrics_elem.set("memory_total", "2000")
metrics_elem.set("disk_used", "5000") metrics_elem.set("disk_used", "5000")
metrics_elem.set("disk_total", "10000") metrics_elem.set("disk_total", "10000")
metrics_elem.set("context", "50") metrics_elem.set("context", str(int(context_size * 100)))
metrics_elem.set("stdin", "0") metrics_elem.set("stdin", "0")
self.mock_metrics.generate_context.return_value = metrics_elem return metrics_elem
# Create test agent with mocked components self.mock_metrics.generate_context.side_effect = generate_context
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_llm_class.return_value = self.mock_llm
mock_metrics_class.return_value = self.mock_metrics
mock_validator_class.return_value = self.mock_validator
# Create test agent with required components
self.agent = TestBaseAgent( self.agent = TestBaseAgent(
model_path="/test/model",
system_prompt="test prompt",
action_schema="test schema", action_schema="test schema",
io_buffer=self.io_buffer 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 # Set timestamp for entries
self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0) 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): def tearDown(self):
"""Clean up resources.""" """Clean up resources."""
if hasattr(self, 'agent'): if hasattr(self, 'agent'):
@@ -87,13 +98,11 @@ class BaseAgentTest(unittest.TestCase):
def test_initialization(self): def test_initialization(self):
"""Test agent initialization and component setup.""" """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._metrics, self.mock_metrics)
self.assertEqual(self.agent._llm, self.mock_llm) self.assertEqual(self.agent._llm, self.mock_llm)
self.assertEqual(self.agent._validator, self.mock_validator) self.assertEqual(self.agent._validator, self.mock_validator)
self.assertEqual(self.agent._io_buffer, self.io_buffer) self.assertEqual(self.agent._parser, self.parser)
self.assertIsInstance(self.agent._parser, ResponseParser)
self.assertEqual(self.agent._system_prompt, "test prompt")
self.assertEqual(self.agent._action_schema, "test schema") self.assertEqual(self.agent._action_schema, "test schema")
def test_cleanup(self): def test_cleanup(self):
@@ -107,16 +116,17 @@ class BaseAgentTest(unittest.TestCase):
# Parse context and verify structure # Parse context and verify structure
root = ET.fromstring(context) root = ET.fromstring(context)
self.assertEqual(root.tag, "state") self.assertEqual(root.tag, "context")
# Check metrics element # Check metrics
context_elem = root.find("context") self.assertEqual(root.get("cpu"), "10")
self.assertIsNotNone(context_elem) self.assertEqual(root.get("gpu"), "20")
self.assertEqual(context_elem.get("cpu"), "10")
self.assertEqual(context_elem.get("gpu"), "20") # Check context size is 0 (empty memory)
self.assertEqual(root.get("context"), "0")
# Check no memory entries # 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): def test_compile_context_with_entries(self):
"""Test context compilation with working memory entries.""" """Test context compilation with working memory entries."""
@@ -129,8 +139,8 @@ class BaseAgentTest(unittest.TestCase):
# Parse and verify context # Parse and verify context
root = ET.fromstring(context) root = ET.fromstring(context)
# Check metrics and entry are present # Check context size reflects one entry
self.assertEqual(len(root.findall("*")), 2) # Metrics + reasoning self.assertEqual(root.get("context"), "1")
# Verify entry content # Verify entry content
reasoning_elem = root.find("reasoning") reasoning_elem = root.find("reasoning")
@@ -152,11 +162,12 @@ class BaseAgentTest(unittest.TestCase):
context = self.agent._compile_context() context = self.agent._compile_context()
root = ET.fromstring(context) root = ET.fromstring(context)
# Verify all entries present # Check context size reflects three entries
self.assertEqual(len(root.findall("reasoning")), 3) self.assertEqual(root.get("context"), "3")
# Verify entry order maintained # Verify entry order maintained
reasoning_elems = root.findall("reasoning") reasoning_elems = root.findall("reasoning")
self.assertEqual(len(reasoning_elems), 3)
for i, elem in enumerate(reasoning_elems): for i, elem in enumerate(reasoning_elems):
self.assertEqual(elem.get("id"), f"id-{i}") self.assertEqual(elem.get("id"), f"id-{i}")
self.assertEqual(elem.text, f"<![CDATA[test {i}]]>") self.assertEqual(elem.text, f"<![CDATA[test {i}]]>")

233
test/web_agent_test.py Normal file
View File

@@ -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 "<reasoning>test reasoning</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], "<reasoning>test reasoning</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 = "<stop/>"
# 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 = "<reasoning>test reasoning</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("<invalid>")
# Verify validation error stored
self.assertEqual(self.agent._validation_error, error_message)
# Verify response updated
self.assertEqual(self.response_changes, ["<invalid>"])
def tearDown(self):
"""Clean up resources."""
if hasattr(self, 'agent'):
del self.agent