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

22
sia/base_agent.py
View File

@@ -1,6 +1,7 @@
from abc import ABC, abstractmethod
from typing import List
import xml.etree.ElementTree as ET
import time
from .llm_engine import LlmEngine
from .response_parser import ResponseParser
@@ -27,7 +28,6 @@ class BaseAgent(ABC):
"""
Initialize agent with required components.
"""
# Initialize components
self._working_memory = working_memory
self._metrics = system_metrics
self._llm = llm
@@ -48,9 +48,27 @@ class BaseAgent(ABC):
Returns:
str: Complete context as XML string
"""
# Get memory context and calculate size
memory_context = self._working_memory.generate_context()
context_size = len(memory_context) / 100
context = self._metrics.generate_context(context_size, self._parser.io_buffer.buffer_length())
# Get system metrics
metrics_data = self._metrics.get_metrics()
# Create context element with metrics
context = ET.Element("context")
context.set("time", metrics_data["timestamp"])
context.set("cpu", str(metrics_data["cpu"]))
context.set("gpu", str(metrics_data["gpu"]))
context.set("memory_used", str(metrics_data["memory_used"]))
context.set("memory_total", str(metrics_data["memory_total"]))
context.set("disk_used", str(metrics_data["disk_used"]))
context.set("disk_total", str(metrics_data["disk_total"]))
context.set("context", str(round(context_size * 100)))
context.set("stdin", str(self._parser.io_buffer.buffer_length()))
# Add memory entries
for entry in memory_context:
context.append(entry)
return pretty_print_element(context)

54
sia/system_metrics.py
View File

@@ -1,19 +1,18 @@
import threading
import time
from typing import Optional
import xml.etree.ElementTree as ET
from typing import Optional, Dict
import psutil
class SystemMetrics:
"""
Tracks system metrics including CPU, GPU, memory and disk usage.
Maintains rolling averages for CPU and GPU between context generations.
Maintains rolling averages for CPU and GPU.
Attributes:
_stop_event: Threading event to signal monitor thread to stop
_monitor_thread: Background thread that collects usage data
_cpu_samples: List of CPU usage samples since last context
_gpu_samples: List of GPU usage samples since last context
_cpu_samples: List of CPU usage samples since last reading
_gpu_samples: List of GPU usage samples since last reading
_metrics_lock: Lock for thread-safe access to samples
"""
@@ -69,52 +68,47 @@ class SystemMetrics:
time.sleep(self._sample_interval)
def generate_context(self, context_usage: float, stdin_buffer_length: int) -> ET.Element:
def get_metrics(self) -> Dict:
"""
Generate XML context element with current system metrics.
Clears usage samples after generating averages.
Args:
context_usage: Current context size usage (0-1)
Get current system metrics.
Clears usage samples after calculating averages.
Returns:
ET.Element: Context element with system metrics
Dict containing:
- cpu: Average CPU usage percentage
- gpu: Average GPU usage percentage
- memory_used: Used memory in bytes
- memory_total: Total memory in bytes
- disk_used: Used disk space in bytes
- disk_total: Total disk space in bytes
- timestamp: Current UTC timestamp in ISO format
"""
# Create context element
context = ET.Element("context")
metrics = {}
# Add timestamp
context.set("time", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
metrics["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
# Add usage metrics
with self._metrics_lock:
# CPU average (or 0 if no samples)
cpu_avg = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
context.set("cpu", str(cpu_avg))
metrics["cpu"] = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
self._cpu_samples.clear()
# GPU average (or 0 if no samples)
gpu_avg = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
context.set("gpu", str(gpu_avg))
metrics["gpu"] = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
self._gpu_samples.clear()
# Memory usage in bytes
memory = psutil.virtual_memory()
context.set("memory_used", str(memory.used))
context.set("memory_total", str(memory.total))
metrics["memory_used"] = memory.used
metrics["memory_total"] = memory.total
# Disk usage in bytes (root partition)
disk = psutil.disk_usage('/')
context.set("disk_used", str(disk.used))
context.set("disk_total", str(disk.total))
metrics["disk_used"] = disk.used
metrics["disk_total"] = disk.total
# Context usage (0-100)
context.set("context_usage", str(round(context_usage * 100)))
# Standard input buffer size
context.set("stdin", stdin_buffer_length)
return context
return metrics
def stop(self):
"""Stop the monitoring thread and clean up."""

46
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
# 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
}
self.mock_metrics.generate_context.side_effect = generate_context
# 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)

76
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,7 +23,6 @@ class SystemMetricsTest(unittest.TestCase):
def validate_timestamp(self, timestamp: str):
"""Validate timestamp format and reasonableness"""
try:
# Check format
self.assertRegex(
timestamp,
@@ -43,9 +41,6 @@ class SystemMetricsTest(unittest.TestCase):
self.assertLess(abs(now_secs - timestamp_secs), 60)
except Exception as e:
self.fail(f"Invalid timestamp format: {timestamp}")
def validate_memory_metrics(self, used: int, total: int):
"""Validate memory metrics are reasonable"""
# Check types
@@ -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,55 +88,47 @@ 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):
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)
initial_metrics = [self.metrics.get_metrics() for _ in range(3)]
initial_cpu = min(m["cpu"] for m in initial_metrics)
# Generate CPU load in separate process
process = multiprocessing.Process(target=cpu_load_process)
@@ -151,28 +137,16 @@ def test_cpu_usage_detection(self):
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)
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,
f"CPU usage should increase under load. Initial: {initial_cpu}, Load: {load_cpu}")
self.assertGreater(load_cpu, initial_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]
for value in test_values:
context = self.metrics.generate_context(value)
self.assertEqual(
int(context.get("context_usage")),
round(value * 100)
)
def test_cleanup(self):
"""Test monitoring stops properly"""
self.metrics.stop()
@@ -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)