diff --git a/sia/base_agent.py b/sia/base_agent.py
index f50e371..b10f86f 100644
--- a/sia/base_agent.py
+++ b/sia/base_agent.py
@@ -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)
\ No newline at end of file
+
+ return pretty_print_element(context)
diff --git a/sia/system_metrics.py b/sia/system_metrics.py
index c4210cb..e0e778e 100644
--- a/sia/system_metrics.py
+++ b/sia/system_metrics.py
@@ -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.
+ Get current system metrics.
+ Clears usage samples after calculating averages.
- Args:
- context_usage: Current context size usage (0-1)
-
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."""
diff --git a/test/base_agent_test.py b/test/base_agent_test.py
index 17ac4c8..69532a8 100644
--- a/test/base_agent_test.py
+++ b/test/base_agent_test.py
@@ -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}")
\ No newline at end of file
+ self.assertEqual(elem.text.strip(), f"test {i}")
diff --git a/test/system_metrics_test.py b/test/system_metrics_test.py
index 2fe6328..561b756 100644
--- a/test/system_metrics_test.py
+++ b/test/system_metrics_test.py
@@ -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)
diff --git a/test/web_agent_test.py b/test/web_agent_test.py
index d92dfb1..426c813 100644
--- a/test/web_agent_test.py
+++ b/test/web_agent_test.py
@@ -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 "test 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 = ""
- # 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("test 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("")
self.assertEqual(self.agent.validation_error, error_message)
self.assertEqual(self.response_changes, [""])
-
- def tearDown(self):
- """Clean up resources."""
- if hasattr(self, 'agent'):
- del self.agent
diff --git a/test/web_socket_manager_test.py b/test/web_socket_manager_test.py
index 5caa729..e92680f 100644
--- a/test/web_socket_manager_test.py
+++ b/test/web_socket_manager_test.py
@@ -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)