Update test_embeddings.py

test_embeddings.py

@@ -1,58 +1,161 @@
 import logging
+import os
 from app import LegalTextSearchBot
 import numpy as np
+from dotenv import load_dotenv
+import time
+from tqdm import tqdm
 
-logging.basicConfig(level=logging.INFO)
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler('embedding_tests.log')
+    ]
+)
 logger = logging.getLogger(__name__)
 
-def test_embeddings():
+def test_environment():
+    """Test environment variables and connections"""
     try:
-        logger.info("Initializing LegalTextSearchBot...")
-        bot = LegalTextSearchBot()
+        load_dotenv()
+        required_vars = [
+            "ASTRA_DB_APPLICATION_TOKEN",
+            "ASTRA_DB_API_ENDPOINT",
+            "ASTRA_DB_COLLECTION",
+            "HUGGINGFACE_API_TOKEN"
+        ]
+        
+        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
         
-        # Test queries
+    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 test_embedding_generation(bot):
+    """Test embedding generation"""
+    try:
         test_queries = [
             "What are the penalties for corruption?",
             "Explain criminal conspiracy",
-            "What constitutes culpable homicide?"
+            "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
         ]
         
-        for query in test_queries:
-            logger.info(f"\nTesting query: {query}")
-            
-            # Generate embedding
-            logger.info("Generating embedding...")
+        logger.info("Testing embedding generation...")
+        for query in tqdm(test_queries, desc="Testing queries"):
             embedding = bot.get_embedding(query)
             
-            # Verify embedding
-            logger.info(f"Embedding dimension: {len(embedding)}")
-            assert len(embedding) == 1024, f"Embedding dimension should be 1024, got {len(embedding)}"
+            # Verify embedding dimension
+            assert len(embedding) == 1024, f"Wrong embedding dimension: {len(embedding)}"
             
             # Verify embedding values
             embedding_array = np.array(embedding)
-            logger.info(f"Embedding stats - Mean: {embedding_array.mean():.4f}, Std: {embedding_array.std():.4f}")
+            assert not np.isnan(embedding_array).any(), "Embedding contains NaN values"
+            assert not np.isinf(embedding_array).any(), "Embedding contains infinite values"
+            
+            # 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 search
-            logger.info("Testing vector search...")
+        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()
+            
+            # Test vector search
             results = bot._search_astra(query)
             
+            # Log search performance
+            elapsed_time = time.time() - start_time
+            logger.info(f"Search time for '{query[:50]}...': {elapsed_time:.2f}s")
+            
+            # Verify results
+            assert isinstance(results, list), "Search results should be a list"
             if results:
-                logger.info(f"Successfully retrieved {len(results)} results")
-                # Print first result title
-                logger.info(f"First result: {results[0].get('title', 'No title')}")
-            else:
-                logger.warning("No results found")
-        
+                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}"
+            
+        logger.info("✅ Search functionality tests passed")
         return True
+        
+    except Exception as e:
+        logger.error(f"Search functionality test failed: {str(e)}")
+        return False
+
+def run_all_tests():
+    """Run all tests"""
+    try:
+        logger.info("\n=== Starting Comprehensive Tests ===\n")
+        
+        # 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
+            
+        # Test 4: Search Functionality
+        if not test_search_functionality(bot):
+            return False
+            
+        logger.info("\n=== All Tests Completed Successfully ===\n")
+        return True
+        
     except Exception as e:
-        logger.error(f"Test failed: {str(e)}")
+        logger.error(f"Test suite failed: {str(e)}")
         return False
 
 if __name__ == "__main__":
-    print("\n=== Starting Embedding Tests ===\n")
-    success = test_embeddings()
+    success = run_all_tests()
     if success:
-        print("\n✅ All embedding tests passed!")
+        print("\n✅ All tests passed successfully!")
     else:
-        print("\n❌ Embedding tests failed!")
+        print("\n❌ Some tests failed. Check the logs for details.")