Moved context generation to base agent instead of system_metrics

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)

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.
-        
-        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."""

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
+        # 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)

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 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()
+        # 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)
+        # 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}")
+        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)
+    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)
         
-    # Generate CPU load in separate process
-    process = multiprocessing.Process(target=cpu_load_process)
-    process.start()
+        # 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)
+        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, 
-            f"CPU usage should increase under load. Initial: {initial_cpu}, Load: {load_cpu}")
+            # Verify load increases CPU usage
+            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)
-            )
+        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)

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 "<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

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)