Moved context generation to base agent instead of system_metrics

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