Moved context generation to base agent instead of system_metrics

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}")
+            self.assertEqual(elem.text.strip(), f"test {i}")