llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Moved context generation to base agent instead of system_metrics

Browse Source
This commit is contained in:
Niels Geens
2024-11-14 14:51:54 +01:00
parent 325870a81a
commit ede0a642d3
6 changed files with 139 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
test/base_agent_test.py
View File

@@ -55,33 +55,25 @@ class TestBaseAgent(BaseAgent):
class BaseAgentTest(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()
# Create parser with IO buffer
self.parser = ResponseParser(self.io_buffer)
# Mock metrics context generator
def generate_context(context_size, stdin_buffer_length):
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", stdin_buffer_length)
return metrics_elem
self.mock_metrics.generate_context.side_effect = generate_context
# 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
}
# Create test agent with required components
# Create test agent
self.agent = TestBaseAgent(
action_schema="test schema",
working_memory=self.working_memory,
@@ -99,13 +91,8 @@ class BaseAgentTest(unittest.TestCase):
logger = logging.getLogger()
logger.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."""
# Remove log handler
logger = logging.getLogger()
logger.removeHandler(self.log_handler)
@@ -137,6 +124,10 @@ class BaseAgentTest(unittest.TestCase):
# Check metrics
self.assertEqual(root.get("cpu"), "10")
self.assertEqual(root.get("gpu"), "20")
self.assertEqual(root.get("memory_used"), "1000")
self.assertEqual(root.get("memory_total"), "2000")
self.assertEqual(root.get("disk_used"), "5000")
self.assertEqual(root.get("disk_total"), "10000")
# Check context size is 0 (empty memory)
self.assertEqual(root.get("context"), "0")
@@ -146,13 +137,10 @@ class BaseAgentTest(unittest.TestCase):
def test_compile_context_with_entries(self):
"""Test context compilation with a single entry."""
# Add test entry to working memory - note id and timestamp are now first
entry = ReasoningEntry("test-id", self.test_timestamp, "test reasoning")
self.agent._working_memory.add_entry(entry)
context = self.agent._compile_context()
# Parse and verify context
root = ET.fromstring(context)
# Check context size reflects one entry
@@ -162,12 +150,10 @@ class BaseAgentTest(unittest.TestCase):
reasoning_elem = root.find("reasoning")
self.assertIsNotNone(reasoning_elem)
self.assertEqual(reasoning_elem.get("id"), "test-id")
# Strip whitespace before comparing
self.assertEqual(reasoning_elem.text.strip(), "test reasoning")
def test_multiple_entries(self):
"""Test handling multiple entries in working memory."""
# Add multiple entries - note id and timestamp are now first
entries = [
ReasoningEntry(f"id-{i}", self.test_timestamp, f"test {i}")
for i in range(3)
@@ -187,4 +173,4 @@ class BaseAgentTest(unittest.TestCase):
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.strip(), f"test {i}")
self.assertEqual(elem.text.strip(), f"test {i}")

138
test/system_metrics_test.py
View File

@@ -1,6 +1,5 @@
import unittest
import time
import xml.etree.ElementTree as ET
import psutil
import multiprocessing
import math
@@ -24,27 +23,23 @@ class SystemMetricsTest(unittest.TestCase):
def validate_timestamp(self, timestamp: str):
"""Validate timestamp format and reasonableness"""
try:
# Check format
self.assertRegex(
timestamp,
r'^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$'
)
# Check timestamp is recent (within last minute)
time_parts = timestamp[:-1].split('T') # Remove Z
time_str = f"{time_parts[0]} {time_parts[1]}"
timestamp_time = time.strptime(time_str, "%Y-%m-%d %H:%M:%S")
now = time.gmtime()
# Convert to seconds since epoch for comparison
timestamp_secs = time.mktime(timestamp_time)
now_secs = time.mktime(now)
self.assertLess(abs(now_secs - timestamp_secs), 60)
except Exception as e:
self.fail(f"Invalid timestamp format: {timestamp}")
# Check format
self.assertRegex(
timestamp,
r'^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$'
)
# Check timestamp is recent (within last minute)
time_parts = timestamp[:-1].split('T') # Remove Z
time_str = f"{time_parts[0]} {time_parts[1]}"
timestamp_time = time.strptime(time_str, "%Y-%m-%d %H:%M:%S")
now = time.gmtime()
# Convert to seconds since epoch for comparison
timestamp_secs = time.mktime(timestamp_time)
now_secs = time.mktime(now)
self.assertLess(abs(now_secs - timestamp_secs), 60)
def validate_memory_metrics(self, used: int, total: int):
"""Validate memory metrics are reasonable"""
@@ -61,7 +56,6 @@ class SystemMetricsTest(unittest.TestCase):
memory = psutil.virtual_memory()
self.assertEqual(total, memory.total)
# Used memory should be within 20% of current usage
# (allowing for normal fluctuation)
self.assertLess(abs(used - memory.used) / memory.total, 0.2)
def validate_disk_metrics(self, used: int, total: int):
@@ -94,84 +88,64 @@ class SystemMetricsTest(unittest.TestCase):
f"{name} above valid range: {value}"
)
def test_initial_context(self):
"""Test initial context generation"""
context = self.metrics.generate_context(0.5, 10)
# Verify it's an XML element
self.assertIsInstance(context, ET.Element)
self.assertEqual(context.tag, "context")
def test_initial_metrics(self):
"""Test initial metrics generation"""
metrics = self.metrics.get_metrics()
# Validate timestamp
self.validate_timestamp(context.get("time"))
self.validate_timestamp(metrics["timestamp"])
# Validate memory metrics
self.validate_memory_metrics(
int(context.get("memory_used")),
int(context.get("memory_total"))
metrics["memory_used"],
metrics["memory_total"]
)
# Validate disk metrics
self.validate_disk_metrics(
int(context.get("disk_used")),
int(context.get("disk_total"))
metrics["disk_used"],
metrics["disk_total"]
)
# Validate usage percentages
self.validate_usage_values([
("CPU", int(context.get("cpu"))),
("GPU", int(context.get("gpu"))),
("Context", int(context.get("context_usage")))
("CPU", metrics["cpu"]),
("GPU", metrics["gpu"])
])
# Validate stdin buffer size
self.assertEqual(context.get("stdin"), 10)
def test_multiple_context_generations(self):
"""Test multiple context generations have different timestamps"""
context1 = self.metrics.generate_context(0.5, 10)
def test_multiple_metric_readings(self):
"""Test multiple metric readings have different timestamps"""
metrics1 = self.metrics.get_metrics()
time.sleep(1)
context2 = self.metrics.generate_context(0.5, 10)
metrics2 = self.metrics.get_metrics()
self.assertNotEqual(
context1.get("time"),
context2.get("time")
metrics1["timestamp"],
metrics2["timestamp"]
)
def test_cpu_usage_detection(self):
"""Test CPU usage increases under load"""
# Get initial CPU usage - take average of multiple samples
initial_contexts = [self.metrics.generate_context(0.5) for _ in range(3)]
initial_cpu = min(int(ctx.get("cpu")) for ctx in initial_contexts)
# Generate CPU load in separate process
process = multiprocessing.Process(target=cpu_load_process)
process.start()
try:
# Wait for load to register and take multiple samples
time.sleep(1.0)
load_contexts = [self.metrics.generate_context(0.5) for _ in range(3)]
load_cpu = max(int(ctx.get("cpu")) for ctx in load_contexts)
def test_cpu_usage_detection(self):
"""Test CPU usage increases under load"""
# Get initial CPU usage - take average of multiple samples
initial_metrics = [self.metrics.get_metrics() for _ in range(3)]
initial_cpu = min(m["cpu"] for m in initial_metrics)
# Verify load increases CPU usage
self.assertGreater(load_cpu, initial_cpu,
f"CPU usage should increase under load. Initial: {initial_cpu}, Load: {load_cpu}")
finally:
process.terminate()
process.join()
def test_context_usage_reflection(self):
"""Test context usage parameter is reflected accurately"""
test_values = [0.0, 0.42, 0.8, 1.0]
# Generate CPU load in separate process
process = multiprocessing.Process(target=cpu_load_process)
process.start()
for value in test_values:
context = self.metrics.generate_context(value)
self.assertEqual(
int(context.get("context_usage")),
round(value * 100)
)
try:
# Wait for load to register and take multiple samples
time.sleep(1.0)
load_metrics = [self.metrics.get_metrics() for _ in range(3)]
load_cpu = max(m["cpu"] for m in load_metrics)
# Verify load increases CPU usage
self.assertGreater(load_cpu, initial_cpu)
finally:
process.terminate()
process.join()
def test_cleanup(self):
"""Test monitoring stops properly"""
@@ -180,6 +154,6 @@ def test_cpu_usage_detection(self):
# Monitor thread should finish
self.assertFalse(self.metrics._monitor_thread.is_alive())
# Should still generate valid context after stopping
context = self.metrics.generate_context(0.5)
self.assertIsInstance(context, ET.Element)
# Should still generate valid metrics after stopping
metrics = self.metrics.get_metrics()
self.assertIsInstance(metrics, dict)

62
test/web_agent_test.py
View File

@@ -32,9 +32,7 @@ class MockCommand(Command):
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()
@@ -45,21 +43,17 @@ class WebAgentTest(unittest.TestCase):
# Create parser with IO buffer
self.parser = ResponseParser(self.io_buffer)
# Mock metrics context generator
def generate_context(context_size, stdin_buffer_length):
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", stdin_buffer_length)
return metrics_elem
self.mock_metrics.generate_context.side_effect = generate_context
# 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):
@@ -104,11 +98,9 @@ class WebAgentTest(unittest.TestCase):
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)
@@ -121,11 +113,10 @@ class WebAgentTest(unittest.TestCase):
def test_approve_context_state_error(self):
"""Test approving context in wrong state."""
# Force agent into UPDATE state
self.agent._state = WebAgentState.UPDATE
with self.assertRaises(Exception) as context:
asyncio.run(self.agent.approve_context())
self.agent.approve_context()
self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
def test_approve_response_state_error(self):
@@ -136,20 +127,11 @@ class WebAgentTest(unittest.TestCase):
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]:
# Return complete response at once for testing
yield "<reasoning>test reasoning</reasoning>"
self.mock_llm.infer.side_effect = mock_infer
# Approve context
self.agent.approve_context()
# Verify state transitions
expected_states = [
WebAgentState.INFERENCE,
WebAgentState.RESPONSE_APPROVAL
@@ -159,56 +141,41 @@ class WebAgentTest(unittest.TestCase):
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._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._state = WebAgentState.RESPONSE_APPROVAL
self.agent._set_response("<reasoning>test reasoning</reasoning>")
# Approve response
self.agent.approve_response()
time.sleep(1)
# 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
@@ -223,8 +190,3 @@ class WebAgentTest(unittest.TestCase):
self.agent._set_response("<invalid>")
self.assertEqual(self.agent.validation_error, error_message)
self.assertEqual(self.response_changes, ["<invalid>"])
def tearDown(self):
"""Clean up resources."""
if hasattr(self, 'agent'):
del self.agent

22
test/web_socket_manager_test.py
View File

@@ -22,21 +22,17 @@ class WebSocketManagerTest(AioHTTPTestCase):
self.io_buffer = WebIOBuffer()
self.working_memory = WorkingMemory()
# Create mock metrics with proper generate_context implementation
# Create mock metrics with proper get_metrics implementation
self.mock_metrics = Mock(spec=SystemMetrics)
metrics_elem = ET.Element("context")
metrics_elem.attrib = {
"time": "2024-10-31T12:00:00Z",
"cpu": "10",
"gpu": "20",
"memory_used": "1000",
"memory_total": "2000",
"disk_used": "5000",
"disk_total": "10000",
"context": "0",
"stdin": "0"
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
}
self.mock_metrics.generate_context.return_value = metrics_elem
# Create minimal mocks for other components
self.mock_llm = Mock(spec=LlmEngine)