Update test_embeddings.py

test_embeddings.py

@@ -4,158 +4,177 @@ from app import LegalTextSearchBot
 import numpy as np
 from dotenv import load_dotenv
 import time
-from tqdm import tqdm
+import gradio as gr
 
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler(),
-        logging.FileHandler('embedding_tests.log')
-    ]
+    format='%(asctime)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-def test_environment():
-    """Test environment variables and connections"""
-    try:
-        load_dotenv()
-        required_vars = [
-            "ASTRA_DB_APPLICATION_TOKEN",
-            "ASTRA_DB_API_ENDPOINT",
-            "ASTRA_DB_COLLECTION",
-            "HUGGINGFACE_API_TOKEN"
-        ]
+class TestResults:
+    def __init__(self):
+        self.results = []
         
-        missing_vars = [var for var in required_vars if not os.getenv(var)]
-        if missing_vars:
-            logger.error(f"Missing environment variables: {missing_vars}")
-            return False
-            
-        logger.info("✅ Environment variables verified")
-        return True
+    def add_result(self, test_name, status, message):
+        self.results.append({
+            'test_name': test_name,
+            'status': status,
+            'message': message,
+            'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
+        })
         
-    except Exception as e:
-        logger.error(f"Environment test failed: {str(e)}")
-        return False
-
-def test_bot_initialization():
-    """Test LegalTextSearchBot initialization"""
-    try:
-        bot = LegalTextSearchBot()
-        logger.info("✅ Bot initialization successful")
-        return bot
-    except Exception as e:
-        logger.error(f"Bot initialization failed: {str(e)}")
-        return None
+    def get_markdown_report(self):
+        report = ["# Test Results\n"]
+        for result in self.results:
+            status_emoji = "✅" if result['status'] else "❌"
+            report.append(f"## {status_emoji} {result['test_name']}")
+            report.append(f"Status: {status_emoji} {'Passed' if result['status'] else 'Failed'}")
+            report.append(f"Time: {result['timestamp']}")
+            report.append(f"Details: {result['message']}\n")
+        return "\n".join(report)
 
-def test_embedding_generation(bot):
-    """Test embedding generation"""
+def run_tests(progress=gr.Progress()):
+    test_results = TestResults()
+    
     try:
-        test_queries = [
-            "What are the penalties for corruption?",
-            "Explain criminal conspiracy",
-            "What constitutes culpable homicide?",
-            "",  # Test empty string
-            "   ",  # Test whitespace
-            "a" * 1000,  # Test long string
-            "Section 123 of IPC",  # Test with numbers
-            "धारा 123",  # Test with non-English
-        ]
+        progress(0, desc="Starting tests...")
         
-        logger.info("Testing embedding generation...")
-        for query in tqdm(test_queries, desc="Testing queries"):
-            embedding = bot.get_embedding(query)
+        # Test 1: Environment Variables
+        progress(0.1, desc="Checking environment variables...")
+        try:
+            load_dotenv()
+            required_vars = [
+                "ASTRA_DB_APPLICATION_TOKEN",
+                "ASTRA_DB_API_ENDPOINT",
+                "ASTRA_DB_COLLECTION",
+                "HUGGINGFACE_API_TOKEN"
+            ]
             
-            # Verify embedding dimension
-            assert len(embedding) == 1024, f"Wrong embedding dimension: {len(embedding)}"
-            
-            # Verify embedding values
-            embedding_array = np.array(embedding)
-            assert not np.isnan(embedding_array).any(), "Embedding contains NaN values"
-            assert not np.isinf(embedding_array).any(), "Embedding contains infinite values"
+            missing_vars = [var for var in required_vars if not os.getenv(var)]
+            if missing_vars:
+                test_results.add_result(
+                    "Environment Check",
+                    False,
+                    f"Missing environment variables: {missing_vars}"
+                )
+            else:
+                test_results.add_result(
+                    "Environment Check",
+                    True,
+                    "All environment variables present"
+                )
+                
+        except Exception as e:
+            test_results.add_result(
+                "Environment Check",
+                False,
+                f"Error checking environment: {str(e)}"
+            )
+        
+        # Test 2: Bot Initialization
+        progress(0.3, desc="Testing bot initialization...")
+        try:
+            bot = LegalTextSearchBot()
+            test_results.add_result(
+                "Bot Initialization",
+                True,
+                "Successfully initialized LegalTextSearchBot"
+            )
             
-            # Log embedding statistics
-            logger.debug(f"Query: {query[:50]}...")
-            logger.debug(f"Embedding stats - Mean: {embedding_array.mean():.4f}, Std: {embedding_array.std():.4f}")
+            # Test 3: Embedding Generation
+            progress(0.5, desc="Testing embedding generation...")
+            test_queries = [
+                "What are the penalties for corruption?",
+                "Explain criminal conspiracy",
+                "What constitutes culpable homicide?"
+            ]
             
-        logger.info("✅ Embedding generation tests passed")
-        return True
-        
-    except Exception as e:
-        logger.error(f"Embedding generation test failed: {str(e)}")
-        return False
-
-def test_search_functionality(bot):
-    """Test search functionality"""
-    try:
-        test_queries = [
-            "What are the penalties for corruption?",
-            "Explain criminal conspiracy",
-            "What constitutes culpable homicide?"
-        ]
-        
-        logger.info("Testing search functionality...")
-        for query in tqdm(test_queries, desc="Testing searches"):
-            start_time = time.time()
+            embedding_results = []
+            for query in test_queries:
+                embedding = bot.get_embedding(query)
+                embedding_array = np.array(embedding)
+                
+                embedding_results.append({
+                    'query': query,
+                    'dimension': len(embedding),
+                    'mean': embedding_array.mean(),
+                    'std': embedding_array.std()
+                })
             
-            # Test vector search
-            results = bot._search_astra(query)
+            test_results.add_result(
+                "Embedding Generation",
+                True,
+                f"Generated embeddings for {len(test_queries)} queries\n" +
+                "\n".join([f"Query: {r['query'][:50]}...\n"
+                          f"Dimension: {r['dimension']}\n"
+                          f"Mean: {r['mean']:.4f}, Std: {r['std']:.4f}\n"
+                          for r in embedding_results])
+            )
             
-            # Log search performance
-            elapsed_time = time.time() - start_time
-            logger.info(f"Search time for '{query[:50]}...': {elapsed_time:.2f}s")
+            # Test 4: Search Functionality
+            progress(0.7, desc="Testing search functionality...")
+            search_results = []
+            for query in test_queries:
+                start_time = time.time()
+                results = bot._search_astra(query)
+                elapsed_time = time.time() - start_time
+                
+                search_results.append({
+                    'query': query,
+                    'num_results': len(results),
+                    'time': elapsed_time
+                })
             
-            # Verify results
-            assert isinstance(results, list), "Search results should be a list"
-            if results:
-                logger.info(f"Found {len(results)} results for '{query[:50]}...'")
-                # Verify result structure
-                first_result = results[0]
-                required_fields = ["section_number", "title", "content"]
-                for field in required_fields:
-                    assert field in first_result, f"Missing required field: {field}"
+            test_results.add_result(
+                "Search Functionality",
+                True,
+                f"Completed searches for {len(test_queries)} queries\n" +
+                "\n".join([f"Query: {r['query'][:50]}...\n"
+                          f"Results found: {r['num_results']}\n"
+                          f"Search time: {r['time']:.2f}s\n"
+                          for r in search_results])
+            )
             
-        logger.info("✅ Search functionality tests passed")
-        return True
+        except Exception as e:
+            test_results.add_result(
+                "Bot Tests",
+                False,
+                f"Error during bot tests: {str(e)}"
+            )
+        
+        progress(1.0, desc="Tests completed!")
+        return test_results.get_markdown_report()
         
     except Exception as e:
-        logger.error(f"Search functionality test failed: {str(e)}")
-        return False
+        return f"# ❌ Test Suite Failed\n\nError: {str(e)}"
 
-def run_all_tests():
-    """Run all tests"""
-    try:
-        logger.info("\n=== Starting Comprehensive Tests ===\n")
+def create_test_interface():
+    with gr.Blocks(title="Legal Search System Tests") as iface:
+        gr.Markdown("""
+        # 🧪 Legal Search System Test Suite
         
-        # Test 1: Environment
-        if not test_environment():
-            return False
-            
-        # Test 2: Bot Initialization
-        bot = test_bot_initialization()
-        if not bot:
-            return False
-            
-        # Test 3: Embedding Generation
-        if not test_embedding_generation(bot):
-            return False
+        This interface runs comprehensive tests on the legal search system components:
+        1. Environment Configuration
+        2. Bot Initialization
+        3. Embedding Generation
+        4. Search Functionality
+        """)
+        
+        with gr.Row():
+            run_button = gr.Button("🚀 Run Tests", variant="primary")
             
-        # Test 4: Search Functionality
-        if not test_search_functionality(bot):
-            return False
+        with gr.Row():
+            output = gr.Markdown("Click 'Run Tests' to start testing...")
             
-        logger.info("\n=== All Tests Completed Successfully ===\n")
-        return True
-        
-    except Exception as e:
-        logger.error(f"Test suite failed: {str(e)}")
-        return False
+        run_button.click(
+            fn=run_tests,
+            outputs=output
+        )
+    
+    return iface
 
 if __name__ == "__main__":
-    success = run_all_tests()
-    if success:
-        print("\n✅ All tests passed successfully!")
-    else:
-        print("\n❌ Some tests failed. Check the logs for details.")
+    demo = create_test_interface()
+    demo.launch()