Spaces:
Running
Running
| """ | |
| Test integration of calculator prompt enhancer with Fixed GAIA Agent. | |
| Verifies that exponentiation operations are properly enhanced. | |
| """ | |
| import pytest | |
| import logging | |
| from unittest.mock import Mock, patch, MagicMock | |
| # Configure test environment | |
| from utils.environment_setup import setup_test_environment | |
| setup_test_environment() | |
| from agents.fixed_enhanced_unified_agno_agent import FixedGAIAAgent | |
| from utils.calculator_prompt_enhancer import CalculatorPromptEnhancer | |
| logger = logging.getLogger(__name__) | |
| class TestAgentPromptEnhancerIntegration: | |
| """Test integration of prompt enhancer with GAIA agent.""" | |
| def setup_method(self): | |
| """Set up test fixtures.""" | |
| self.enhancer = CalculatorPromptEnhancer() | |
| def test_agent_has_prompt_enhancer(self): | |
| """Test that agent initializes with prompt enhancer.""" | |
| with patch.dict('os.environ', {'MISTRAL_API_KEY': 'test-key'}): | |
| with patch('agents.fixed_enhanced_unified_agno_agent.MistralChat'): | |
| with patch('agents.fixed_enhanced_unified_agno_agent.Agent'): | |
| agent = FixedGAIAAgent() | |
| # Verify prompt enhancer is initialized | |
| assert hasattr(agent, 'prompt_enhancer') | |
| assert isinstance(agent.prompt_enhancer, CalculatorPromptEnhancer) | |
| logger.info("✅ Agent has prompt enhancer initialized") | |
| def test_exponentiation_question_enhancement(self): | |
| """Test that exponentiation questions are enhanced.""" | |
| # Test questions with exponentiation | |
| test_cases = [ | |
| { | |
| 'question': 'Calculate 2^8', | |
| 'should_enhance': True, | |
| 'description': 'caret notation' | |
| }, | |
| { | |
| 'question': 'What is 3**4?', | |
| 'should_enhance': True, | |
| 'description': 'double asterisk notation' | |
| }, | |
| { | |
| 'question': 'Compute 2 to the power of 8', | |
| 'should_enhance': True, | |
| 'description': 'power of notation' | |
| }, | |
| { | |
| 'question': 'What is 5 squared?', | |
| 'should_enhance': True, | |
| 'description': 'squared notation' | |
| }, | |
| { | |
| 'question': 'Calculate 25 * 17', | |
| 'should_enhance': False, | |
| 'description': 'regular multiplication' | |
| }, | |
| { | |
| 'question': 'What is 144 / 12?', | |
| 'should_enhance': False, | |
| 'description': 'division operation' | |
| } | |
| ] | |
| for case in test_cases: | |
| question = case['question'] | |
| should_enhance = case['should_enhance'] | |
| description = case['description'] | |
| # Test enhancement | |
| enhanced = self.enhancer.enhance_prompt_for_exponentiation(question) | |
| is_enhanced = enhanced != question | |
| assert is_enhanced == should_enhance, f"Enhancement mismatch for {description}: '{question}'" | |
| if should_enhance: | |
| # Verify enhancement contains Python guidance | |
| assert 'python' in enhanced.lower() or 'pow(' in enhanced or '**' in enhanced | |
| logger.info(f"✅ Enhanced {description}: {question}") | |
| else: | |
| logger.info(f"✅ No enhancement needed for {description}: {question}") | |
| def test_agent_uses_enhanced_prompt(self, mock_agent_class, mock_mistral_class): | |
| """Test that agent uses enhanced prompts for exponentiation.""" | |
| # Mock the agent run method | |
| mock_agent_instance = Mock() | |
| mock_agent_instance.run = Mock(return_value=Mock(content="FINAL ANSWER: 256")) | |
| mock_agent_class.return_value = mock_agent_instance | |
| # Mock Mistral | |
| mock_mistral_class.return_value = Mock() | |
| with patch.dict('os.environ', {'MISTRAL_API_KEY': 'test-key'}): | |
| agent = FixedGAIAAgent() | |
| # Test with exponentiation question | |
| question = "Calculate 2^8" | |
| result = agent(question) | |
| # Verify agent.run was called | |
| assert mock_agent_instance.run.called | |
| # Get the actual prompt passed to agent.run | |
| call_args = mock_agent_instance.run.call_args | |
| actual_prompt = call_args[0][0] # First positional argument | |
| # Verify the prompt was enhanced (should be longer and contain guidance) | |
| assert len(actual_prompt) > len(question) | |
| assert 'python' in actual_prompt.lower() or 'pow(' in actual_prompt or '**' in actual_prompt | |
| logger.info(f"✅ Agent used enhanced prompt for exponentiation") | |
| logger.info(f" Original: {question}") | |
| logger.info(f" Enhanced length: {len(actual_prompt)} vs {len(question)}") | |
| def test_agent_no_enhancement_for_regular_math(self, mock_agent_class, mock_mistral_class): | |
| """Test that agent doesn't enhance regular math questions.""" | |
| # Mock the agent run method | |
| mock_agent_instance = Mock() | |
| mock_agent_instance.run = Mock(return_value=Mock(content="FINAL ANSWER: 425")) | |
| mock_agent_class.return_value = mock_agent_instance | |
| # Mock Mistral | |
| mock_mistral_class.return_value = Mock() | |
| with patch.dict('os.environ', {'MISTRAL_API_KEY': 'test-key'}): | |
| agent = FixedGAIAAgent() | |
| # Test with regular math question | |
| question = "Calculate 25 * 17" | |
| result = agent(question) | |
| # Verify agent.run was called | |
| assert mock_agent_instance.run.called | |
| # Get the actual prompt passed to agent.run | |
| call_args = mock_agent_instance.run.call_args | |
| actual_prompt = call_args[0][0] # First positional argument | |
| # Verify the prompt was NOT enhanced (should be the same) | |
| assert actual_prompt == question | |
| logger.info(f"✅ Agent did not enhance regular math question") | |
| logger.info(f" Question: {question}") | |
| def test_enhancement_preserves_file_context(self): | |
| """Test that enhancement works with file context.""" | |
| # Simulate a question with file context | |
| base_question = "Calculate 2^8" | |
| file_context = "File 1: data.csv (CSV format), 1024 bytes\nCSV Data: numbers,values\n1,2\n3,4" | |
| question_with_files = f"{base_question}\n\nFile Context:\n{file_context}" | |
| # Test enhancement | |
| enhanced = self.enhancer.enhance_prompt_for_exponentiation(question_with_files) | |
| # Verify enhancement occurred | |
| assert enhanced != question_with_files | |
| assert len(enhanced) > len(question_with_files) | |
| # Verify file context is preserved | |
| assert file_context in enhanced | |
| # Verify exponentiation guidance is added | |
| assert 'python' in enhanced.lower() or 'pow(' in enhanced or '**' in enhanced | |
| logger.info("✅ Enhancement preserves file context") | |
| logger.info(f" Original length: {len(question_with_files)}") | |
| logger.info(f" Enhanced length: {len(enhanced)}") | |
| if __name__ == "__main__": | |
| # Run tests with verbose output | |
| pytest.main([__file__, "-v", "-s"]) |