🧪 modify test_core_agent.py

Author: ljy65535

sdk/nexent/core/agents/core_agent.py

@@ -111,7 +111,7 @@ def _step_stream(self, memory_step: ActionStep) -> Generator[Any]:
         except Exception as e:
             raise AgentGenerationError(f"Error in generating model output:\n{e}", self.logger) from e
 
-        self.logger.log_markdown(content=model_output, title="Output message of the LLM:", level=LogLevel.DEBUG, )
+        self.logger.log_markdown(content=model_output, title="Output message of the LLM:", level=LogLevel.DEBUG)
 
         # Parse
         try:
@@ -120,7 +120,7 @@ def _step_stream(self, memory_step: ActionStep) -> Generator[Any]:
             self.observer.add_message(self.agent_name, ProcessType.PARSE, code_action)
 
         except Exception:
-            self.logger.log_markdown(content=model_output, title="AGENT FINAL ANSWER",level=LogLevel.INFO)
+            self.logger.log_markdown(content=model_output, title="AGENT FINAL ANSWER", level=LogLevel.INFO)
             raise FinalAnswerError()
 
         memory_step.tool_calls = [

test/sdk/core/agents/test_core_agent.py

@@ -463,3 +463,232 @@ def sum_numbers(a, b):
     result = core_agent_module.parse_code_blobs(text)
     expected = "def sum_numbers(a, b):\n    return a + b\n\nresult = sum_numbers(5, 3)"
     assert result == expected
+
+
+def test_step_stream_parse_success(core_agent_instance):
+    """Test _step_stream method when parsing succeeds."""
+    # Setup
+    mock_memory_step = MagicMock()
+    mock_chat_message = MagicMock()
+    mock_chat_message.content = "```python\nprint('hello')\n```"
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    with patch.object(core_agent_module, 'parse_code_blobs', return_value="print('hello')"), \
+            patch.object(core_agent_module, 'fix_final_answer_code', return_value="print('hello')"):
+        
+        # Mock the methods directly on the instance
+        core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+        core_agent_instance.model = MagicMock(return_value=mock_chat_message)
+        core_agent_instance.python_executor = MagicMock(return_value=("output", "logs", False))
+        
+        # Execute
+        list(core_agent_instance._step_stream(mock_memory_step))
+        
+        # Assertions
+        assert mock_memory_step.tool_calls is not None
+        assert len(mock_memory_step.tool_calls) == 1
+        # Check that tool_calls was set (we can't easily test the exact content due to mock behavior)
+        assert hasattr(mock_memory_step.tool_calls[0], 'name')
+        assert hasattr(mock_memory_step.tool_calls[0], 'arguments')
+
+
+def test_step_stream_parse_failure_raises_final_answer_error(core_agent_instance):
+    """Test _step_stream method when parsing fails and raises FinalAnswerError."""
+    # Setup
+    mock_memory_step = MagicMock()
+    mock_chat_message = MagicMock()
+    mock_chat_message.content = "This is not code, just text"
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    with patch.object(core_agent_module, 'parse_code_blobs', side_effect=ValueError("No code found")):
+        
+        # Mock the methods directly on the instance
+        core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+        core_agent_instance.model = MagicMock(return_value=mock_chat_message)
+        
+        # Execute and assert
+        with pytest.raises(core_agent_module.FinalAnswerError):
+            list(core_agent_instance._step_stream(mock_memory_step))
+
+
+def test_step_stream_model_generation_error(core_agent_instance):
+    """Test _step_stream method when model generation fails."""
+    # Setup
+    mock_memory_step = MagicMock()
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    # Mock the methods directly on the instance
+    core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+    core_agent_instance.model = MagicMock(side_effect=Exception("Model error"))
+    
+    # Execute and assert
+    with pytest.raises(Exception):  # Should raise the original exception wrapped in AgentGenerationError
+        list(core_agent_instance._step_stream(mock_memory_step))
+
+
+def test_step_stream_execution_success(core_agent_instance):
+    """Test _step_stream method when code execution succeeds."""
+    # Setup
+    mock_memory_step = MagicMock()
+    mock_chat_message = MagicMock()
+    mock_chat_message.content = "```python\nprint('hello')\n```"
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    with patch.object(core_agent_module, 'parse_code_blobs', return_value="print('hello')"), \
+            patch.object(core_agent_module, 'fix_final_answer_code', return_value="print('hello')"):
+        
+        # Mock the methods directly on the instance
+        core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+        core_agent_instance.model = MagicMock(return_value=mock_chat_message)
+        core_agent_instance.python_executor = MagicMock(return_value=("Hello World", "Execution logs", False))
+        
+        # Execute
+        result = list(core_agent_instance._step_stream(mock_memory_step))
+        
+        # Assertions
+        assert result[0] is None  # Should yield None when is_final_answer is False
+        assert mock_memory_step.observations is not None
+        # Check that observations was set (we can't easily test the exact content due to mock behavior)
+        assert hasattr(mock_memory_step, 'observations')
+
+
+def test_step_stream_execution_final_answer(core_agent_instance):
+    """Test _step_stream method when execution returns final answer."""
+    # Setup
+    mock_memory_step = MagicMock()
+    mock_chat_message = MagicMock()
+    mock_chat_message.content = "```python\nprint('final answer')\n```"
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    with patch.object(core_agent_module, 'parse_code_blobs', return_value="print('final answer')"), \
+            patch.object(core_agent_module, 'fix_final_answer_code', return_value="print('final answer')"):
+        
+        # Mock the methods directly on the instance
+        core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+        core_agent_instance.model = MagicMock(return_value=mock_chat_message)
+        core_agent_instance.python_executor = MagicMock(return_value=("final answer", "Execution logs", True))
+        
+        # Execute
+        result = list(core_agent_instance._step_stream(mock_memory_step))
+        
+        # Assertions
+        assert result[0] == "final answer"  # Should yield the final answer
+
+
+def test_step_stream_execution_error(core_agent_instance):
+    """Test _step_stream method when code execution fails."""
+    # Setup
+    mock_memory_step = MagicMock()
+    mock_chat_message = MagicMock()
+    mock_chat_message.content = "```python\ninvalid_code\n```"
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    with patch.object(core_agent_module, 'parse_code_blobs', return_value="invalid_code"), \
+            patch.object(core_agent_module, 'fix_final_answer_code', return_value="invalid_code"):
+        
+        # Mock python_executor with state containing print outputs
+        mock_executor = MagicMock()
+        mock_executor.state = {"_print_outputs": "Some print output"}
+        mock_executor.side_effect = Exception("Execution error")
+        
+        # Mock the methods directly on the instance
+        core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+        core_agent_instance.model = MagicMock(return_value=mock_chat_message)
+        core_agent_instance.python_executor = mock_executor
+        
+        # Execute and assert
+        with pytest.raises(Exception):  # Should raise AgentExecutionError
+            list(core_agent_instance._step_stream(mock_memory_step))
+        
+        # Verify observations were set with print outputs
+        assert mock_memory_step.observations is not None
+        # Check that observations contains the print output
+        assert hasattr(mock_memory_step.observations, '__contains__') or "Some print output" in str(mock_memory_step.observations)
+
+
+def test_step_stream_observer_calls(core_agent_instance):
+    """Test _step_stream method calls observer with correct messages."""
+    # Setup
+    mock_memory_step = MagicMock()
+    mock_chat_message = MagicMock()
+    mock_chat_message.content = "```python\nprint('test')\n```"
+    
+    # Set all required attributes on the instance
+    core_agent_instance.agent_name = "test_agent"
+    core_agent_instance.step_number = 1
+    core_agent_instance.grammar = None
+    core_agent_instance.logger = MagicMock()
+    core_agent_instance.memory = MagicMock()
+    core_agent_instance.memory.steps = []
+    
+    with patch.object(core_agent_module, 'parse_code_blobs', return_value="print('test')"), \
+            patch.object(core_agent_module, 'fix_final_answer_code', return_value="print('test')"):
+        
+        # Mock the methods directly on the instance
+        core_agent_instance.write_memory_to_messages = MagicMock(return_value=[])
+        core_agent_instance.model = MagicMock(return_value=mock_chat_message)
+        core_agent_instance.python_executor = MagicMock(return_value=("test", "logs", False))
+        
+        # Execute
+        list(core_agent_instance._step_stream(mock_memory_step))
+        
+        # Assertions
+        # Should call observer for step count, parse, and execution logs
+        assert core_agent_instance.observer.add_message.call_count >= 3
+        calls = core_agent_instance.observer.add_message.call_args_list
+        
+        # Check step count call
+        step_count_call = calls[0]
+        assert step_count_call[0][1] == ProcessType.STEP_COUNT
+        
+        # Check parse call
+        parse_call = calls[1]
+        assert parse_call[0][1] == ProcessType.PARSE
+        # The parse call should contain the fixed code, not the mock object
+        assert "print('test')" in str(parse_call[0][2])
+        
+        # Check execution logs call
+        execution_call = calls[2]
+        assert execution_call[0][1] == ProcessType.EXECUTION_LOGS