test_search.py

#!/usr/bin/env python3
"""
Test script to verify that diverse documents are being returned for different queries
"""
import asyncio
import logging
from rag_service import search_documents_async

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("test_search")

async def test_diverse_responses():
    """Test that different queries return different, relevant documents"""
    
    test_queries = [
        "pension rules",
        "leave policy",
        "salary structure", 
        "training programs",
        "retirement benefits",
        "procurement policy",
        "get out",  # This was giving generic response before
        "what about transfers"
    ]
    
    logger.info("🧪 Testing diverse document responses...")
    
    for query in test_queries:
        logger.info(f"\n🔍 Testing query: '{query}'")
        try:
            docs = await search_documents_async(query, limit=2)
            if docs:
                for i, doc in enumerate(docs):
                    title = doc.get('source_title', 'Unknown')
                    content_preview = doc.get('clause_text', '')[:100] + "..."
                    logger.info(f"  📄 Document {i+1}: {title}")
                    logger.info(f"      Preview: {content_preview}")
            else:
                logger.warning(f"  ⚠️ No documents found for: {query}")
        except Exception as e:
            logger.error(f"  ❌ Error searching for '{query}': {e}")

if __name__ == "__main__":
    asyncio.run(test_diverse_responses())