RimaAlaya/CineRAG
1
1"""2Test suite to evaluate RAG system quality3This helps you understand if your system is actually working!4"""5 6from main import MovieRAGSystem7import json8 9# Test categories with expected behaviors10TEST_CASES = {11 "factual_questions": [12 {13 "query": "Who directed Inception?",14 "expected_movie": "Inception",15 "expected_chunk_type": "crew",16 "expected_answer": "Christopher Nolan"17 },18 {19 "query": "Who stars in Titanic?",20 "expected_movie": "Titanic",21 "expected_chunk_type": "cast",22 "expected_answer": "Leonardo DiCaprio"23 },24 {25 "query": "What year was The Matrix released?",26 "expected_movie": "The Matrix",27 "expected_chunk_type": "metadata",28 "expected_answer": "1999"29 },30 {31 "query": "Who plays Neo in The Matrix?",32 "expected_movie": "The Matrix",33 "expected_chunk_type": "cast",34 "expected_answer": "Keanu Reeves"35 },36 {37 "query": "What is the runtime of Inception?",38 "expected_movie": "Inception",39 "expected_chunk_type": "metadata",40 "expected_answer": "148 minutes"41 }42 ],43 44 "plot_questions": [45 {46 "query": "What is Inception about?",47 "expected_movie": "Inception",48 "expected_chunk_type": "plot",49 },50 {51 "query": "Describe the plot of The Matrix",52 "expected_movie": "The Matrix",53 "expected_chunk_type": "plot",54 },55 {56 "query": "What happens in Titanic?",57 "expected_movie": "Titanic",58 "expected_chunk_type": "plot",59 }60 ],61 62 "genre_questions": [63 {64 "query": "What genre is The Dark Knight?",65 "expected_chunk_type": "metadata",66 },67 {68 "query": "Is Inception a sci-fi movie?",69 "expected_movie": "Inception",70 "expected_chunk_type": "metadata",71 }72 ],73 74 "rating_questions": [75 {76 "query": "What is the rating of The Matrix?",77 "expected_movie": "The Matrix",78 "expected_chunk_type": "metadata",79 },80 {81 "query": "How popular is Inception?",82 "expected_movie": "Inception",83 "expected_chunk_type": "metadata",84 }85 ],86 87 "edge_cases": [88 {89 "query": "Movies about dreams",90 "note": "Should find Inception"91 },92 {93 "query": "Leonardo DiCaprio movies",94 "note": "Should find multiple movies"95 },96 {97 "query": "Christopher Nolan films",98 "note": "Should find Nolan-directed movies"99 }100 ]101}102 103def verify_result(result, expected_movie=None, expected_chunk_type=None, expected_answer=None):104 """Check if a result matches expectations"""105 checks = []106 107 # Check movie match108 if expected_movie:109 movie_match = result['movie_title'] == expected_movie110 checks.append(("Movie Match", movie_match))111 112 # Check chunk type113 if expected_chunk_type:114 chunk_match = result['chunk_type'] == expected_chunk_type115 checks.append(("Chunk Type", chunk_match))116 117 # Check if answer appears in text118 if expected_answer:119 answer_found = expected_answer.lower() in result['text'].lower()120 checks.append(("Answer Found", answer_found))121 122 return checks123 124def run_test_category(rag, category_name, test_cases):125 """Run all tests in a category"""126 print(f"\n{'='*80}")127 print(f"📋 {category_name.upper().replace('_', ' ')}")128 print(f"{'='*80}")129 130 passed = 0131 failed = 0132 133 for i, test in enumerate(test_cases, 1):134 query = test['query']135 print(f"\n{i}. Query: '{query}'")136 print("-" * 80)137 138 # Get results139 results = rag.search(query, top_k=3)140 top_result = results[0]141 142 # Display top result143 print(f" Top Result: {top_result['movie_title']} [{top_result['chunk_type']}]")144 print(f" Score: {top_result['relevance_score']:.4f}")145 print(f" Text: {top_result['text'][:100]}...")146 147 # Check expectations148 if 'expected_movie' in test or 'expected_chunk_type' in test or 'expected_answer' in test:149 checks = verify_result(150 top_result,151 test.get('expected_movie'),152 test.get('expected_chunk_type'),153 test.get('expected_answer')154 )155 156 print("\n Checks:")157 all_passed = True158 for check_name, check_result in checks:159 status = "✅" if check_result else "❌"160 print(f" {status} {check_name}: {check_result}")161 if not check_result:162 all_passed = False163 164 if all_passed:165 passed += 1166 print(" Result: ✅ PASSED")167 else:168 failed += 1169 print(" Result: ❌ FAILED")170 171 if 'note' in test:172 print(f"\n Note: {test['note']}")173 174 # Category summary175 if passed + failed > 0:176 print(f"\n{'='*80}")177 print(f"Category Results: ✅ {passed} passed | ❌ {failed} failed")178 success_rate = (passed / (passed + failed)) * 100179 print(f"Success Rate: {success_rate:.1f}%")180 return passed, failed181 else:182 return 0, 0183 184def run_all_tests():185 """Run complete test suite"""186 print("🎬 RAG SYSTEM TEST SUITE")187 print("="*80)188 print("This will test if your RAG system retrieves correct information\n")189 190 # Initialize RAG system191 rag = MovieRAGSystem()192 193 # Run each category194 total_passed = 0195 total_failed = 0196 197 for category, tests in TEST_CASES.items():198 passed, failed = run_test_category(rag, category, tests)199 total_passed += passed200 total_failed += failed201 202 # Final summary203 print(f"\n\n{'='*80}")204 print("📊 FINAL SUMMARY")205 print(f"{'='*80}")206 207 total_tests = total_passed + total_failed208 if total_tests > 0:209 overall_success = (total_passed / total_tests) * 100210 print(f"\nTotal Tests: {total_tests}")211 print(f"✅ Passed: {total_passed}")212 print(f"❌ Failed: {total_failed}")213 print(f"Success Rate: {overall_success:.1f}%")214 215 print("\nKey Findings:")216 print("✅ Your RAG system successfully retrieves relevant chunks")217 print("✅ Semantic search is working (similar meaning → similar results)")218 if total_failed > 0:219 print("⚠️ Some queries might need better chunking or reranking")220 print("\n💡 Next step: Create evaluation metrics to measure this systematically!")221 222if __name__ == "__main__":223 run_all_tests()