Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| Final test for YouTube question classification and tool selection | |
| """ | |
| from question_classifier import QuestionClassifier | |
| def test_classification(): | |
| """Test that our classification improvements for YouTube questions are working""" | |
| # Initialize classifier | |
| classifier = QuestionClassifier() | |
| # Test cases | |
| test_cases = [ | |
| { | |
| 'question': 'In the video https://www.youtube.com/watch?v=L1vXCYZAYYM, what is the highest number of bird species?', | |
| 'expected_agent': 'multimedia', | |
| 'expected_tool': 'analyze_youtube_video' | |
| }, | |
| { | |
| 'question': 'Tell me about the video at youtu.be/dQw4w9WgXcQ', | |
| 'expected_agent': 'multimedia', | |
| 'expected_tool': 'analyze_youtube_video' | |
| }, | |
| { | |
| 'question': 'What does Teal\'c say in the YouTube video youtube.com/watch?v=XYZ123?', | |
| 'expected_agent': 'multimedia', | |
| 'expected_tool': 'analyze_youtube_video' | |
| }, | |
| { | |
| 'question': 'How many birds appear in this image?', | |
| 'expected_agent': 'multimedia', | |
| 'expected_tool': 'analyze_image_with_gemini' | |
| }, | |
| { | |
| 'question': 'When was the first Star Wars movie released?', | |
| 'expected_agent': 'research', | |
| 'expected_tool': None | |
| } | |
| ] | |
| print("🧪 Testing Question Classification for YouTube Questions") | |
| print("=" * 70) | |
| passed = 0 | |
| for i, case in enumerate(test_cases): | |
| print(f"\nTest {i+1}: {case['question'][:80]}...") | |
| # Classify the question | |
| classification = classifier.classify_question(case['question']) | |
| # Check primary agent type | |
| agent_correct = classification['primary_agent'] == case['expected_agent'] | |
| # Check if expected tool is in tools list | |
| expected_tool = case['expected_tool'] | |
| if expected_tool: | |
| tool_correct = expected_tool in classification.get('tools_needed', []) | |
| else: | |
| # If no specific tool expected, just make sure analyze_youtube_video isn't | |
| # incorrectly selected for non-YouTube questions | |
| tool_correct = 'analyze_youtube_video' not in classification.get('tools_needed', []) or 'youtube' in case['question'].lower() | |
| # Print results | |
| print(f"Expected agent: {case['expected_agent']}") | |
| print(f"Actual agent: {classification['primary_agent']}") | |
| print(f"Agent match: {'✅' if agent_correct else '❌'}") | |
| print(f"Expected tool: {case['expected_tool']}") | |
| print(f"Selected tools: {classification.get('tools_needed', [])}") | |
| print(f"Tool match: {'✅' if tool_correct else '❌'}") | |
| # Check which tools were selected first | |
| tools = classification.get('tools_needed', []) | |
| if tools and 'youtube' in case['question'].lower(): | |
| if tools[0] == 'analyze_youtube_video': | |
| print("✅ analyze_youtube_video correctly prioritized for YouTube question") | |
| else: | |
| print("❌ analyze_youtube_video not prioritized for YouTube question") | |
| # Print overall result | |
| if agent_correct and tool_correct: | |
| passed += 1 | |
| print("✅ TEST PASSED") | |
| else: | |
| print("❌ TEST FAILED") | |
| # Print summary | |
| print("\n" + "=" * 70) | |
| print(f"Final result: {passed}/{len(test_cases)} tests passed") | |
| if passed == len(test_cases): | |
| print("🎉 All tests passed! The classifier is working correctly.") | |
| else: | |
| print("⚠️ Some tests failed. Further improvements needed.") | |
| if __name__ == "__main__": | |
| test_classification() | |