Add tests for prompt handling

test/agentchat/contrib/test_reasoning_agent.py

@@ -429,5 +429,109 @@ def factorial(n):
         assert response == "The factorial of 5 is 120"
 
 
+def test_prepare_prompt_single_message(reasoning_agent):
+    """
+    Test that when a single message is provided without a conversation history,
+    the _process_prompt method returns the message content as the prompt.
+    """
+    messages = [{"role": "user", "content": TEST_PROMPT}]
+    # Call _process_prompt. Here, we pass the agent itself as sender.
+    prompt, ground_truth = reasoning_agent._process_prompt(messages, sender=reasoning_agent)
+
+    # Since there is only one message, the prompt should equal the message content.
+    assert TEST_PROMPT in prompt
+    # No ground truth is expected.
+    assert ground_truth is None
+
+
+def test_prepare_prompt_with_ground_truth(reasoning_agent):
+    """
+    Test that when a message contains a GROUND_TRUTH marker,
+    _process_prompt splits the content appropriately.
+    """
+    # The message includes the ground truth marker.
+    message_content = f"{TEST_PROMPT} {TEST_GROUND_TRUTH}"
+    messages = [{"role": "user", "content": message_content}]
+    prompt, ground_truth = reasoning_agent._process_prompt(messages, sender=reasoning_agent)
+
+    # The prompt should contain the text before the marker.
+    assert TEST_PROMPT in prompt
+    # The ground truth should start with 'GROUND_TRUTH'
+    assert ground_truth is not None
+    assert ground_truth.startswith("GROUND_TRUTH")
+    # Optionally, you can check that the ground truth includes the expected answer.
+    assert "Paris" in ground_truth
+
+
+def test_prepare_prompt_multi_message(reasoning_agent):
+    """
+    Test that when multiple messages are provided, _process_prompt uses the prompt rewriter.
+    Because the method calls self.send with self._prompt_rewriter as recipient,
+    we override the _prompt_rewriter.last_message method to simulate a rewriter response.
+    """
+    messages = [
+        {"role": "user", "content": TEST_PROMPT},
+        {"role": "assistant", "content": "I believe the answer might be Paris."},
+    ]
+
+    # Monkey-patch the prompt rewriter's last_message to return a predetermined prompt.
+    simulated_rewritten_prompt = (
+        "QUESTION: What is the capital of France?\n\n"
+        "STEPS ALREADY EXECUTED:\n- Asked about the capital\n- Received a hint that it might be Paris"
+    )
+
+    with patch(
+        "autogen.agentchat.conversable_agent.ConversableAgent._generate_oai_reply_from_client"
+    ) as mock_oai_reply:
+
+        def mock_response(*args, **kwargs):
+            return {"content": simulated_rewritten_prompt}
+
+        mock_oai_reply.side_effect = mock_response
+
+        prompt, ground_truth = reasoning_agent._process_prompt(messages, sender=reasoning_agent)
+
+    # The returned prompt should match the simulated rewritten prompt.
+    assert prompt == simulated_rewritten_prompt
+    # Since no ground truth was provided in any message, ground_truth should be None.
+    assert ground_truth is None
+
+
+def test_prepare_prompt_multi_message_with_ground_truth(reasoning_agent):
+    """
+    Test that when multiple messages are provided, _process_prompt uses the prompt rewriter.
+    If a message contains a GROUND_TRUTH marker, the method should split the content appropriately.
+    """
+    messages = [
+        {"role": "user", "content": f"{TEST_PROMPT} {TEST_GROUND_TRUTH}"},
+        {"role": "assistant", "content": "I believe the answer might be Paris."},
+    ]
+
+    # Monkey-patch the prompt rewriter's last_message to return a predetermined prompt.
+    simulated_rewritten_prompt = (
+        "QUESTION: What is the capital of France?\n\n"
+        "STEPS ALREADY EXECUTED:\n- Asked about the capital\n- Received a hint that it might be Paris"
+    )
+
+    with patch(
+        "autogen.agentchat.conversable_agent.ConversableAgent._generate_oai_reply_from_client"
+    ) as mock_oai_reply:
+
+        def mock_response(*args, **kwargs):
+            return {"content": simulated_rewritten_prompt}
+
+        mock_oai_reply.side_effect = mock_response
+
+        prompt, ground_truth = reasoning_agent._process_prompt(messages, sender=reasoning_agent)
+
+    # The returned prompt should match the simulated rewritten prompt.
+    assert prompt == simulated_rewritten_prompt
+    # The ground truth should start with 'GROUND_TRUTH'
+    assert ground_truth is not None
+    assert ground_truth.startswith("GROUND_TRUTH")
+    # Optionally, you can check that the ground truth includes the expected answer.
+    assert "Paris" in ground_truth
+
+
 if __name__ == "__main__":
     pytest.main([__file__])