test_real_search.py

#!/usr/bin/env python3
"""
Test script for the new real web search tools in research_team.py
"""

import os
import asyncio
import logging
from dotenv import load_dotenv

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("TestRealSearch")

def setup_environment():
    """Setup environment variables"""
    load_dotenv()
    
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        logger.error("❌ OPENAI_API_KEY not found in environment variables")
        logger.info("💡 Make sure you have a .env file with: OPENAI_API_KEY=your_api_key")
        return False
    
    logger.info("✅ OpenAI API Key configured correctly")
    return True

async def test_search_tools():
    """Test the new search tools"""
    logger.info("🔍 Testing new web search tools...")
    
    try:
        # Import the tools from research_team
        from research_team import google_scholar_search, pubmed_search, arxiv_search
        
        # Test data
        test_query = "machine learning in healthcare"
        test_claim_id = "test_claim_1"
        
        logger.info(f"📝 Test query: {test_query}")
        logger.info(f"🔍 Claim ID: {test_claim_id}")
        
        # Test Google Scholar search
        logger.info("\n" + "="*40)
        logger.info("Testing Google Scholar Search")
        logger.info("="*40)
        
        gs_result = google_scholar_search.invoke({
            "query": test_query,
            "claim_id": test_claim_id
        })
        
        logger.info(f"Google Scholar results: {len(gs_result.get('results', []))} found")
        if gs_result.get('results'):
            logger.info(f"First result title: {gs_result['results'][0].get('title', 'No title')[:50]}...")
        
        # Test PubMed search  
        logger.info("\n" + "="*40)
        logger.info("Testing PubMed Search")
        logger.info("="*40)
        
        pm_result = pubmed_search.invoke({
            "query": test_query,
            "claim_id": test_claim_id
        })
        
        logger.info(f"PubMed results: {len(pm_result.get('results', []))} found")
        if pm_result.get('results'):
            logger.info(f"First result title: {pm_result['results'][0].get('title', 'No title')[:50]}...")
        
        # Test arXiv search
        logger.info("\n" + "="*40)
        logger.info("Testing arXiv Search")
        logger.info("="*40)
        
        ar_result = arxiv_search.invoke({
            "query": test_query,
            "claim_id": test_claim_id
        })
        
        logger.info(f"arXiv results: {len(ar_result.get('results', []))} found")
        if ar_result.get('results'):
            logger.info(f"First result title: {ar_result['results'][0].get('title', 'No title')[:50]}...")
        
        # Summary
        total_results = (
            len(gs_result.get('results', [])) + 
            len(pm_result.get('results', [])) + 
            len(ar_result.get('results', []))
        )
        
        logger.info("\n" + "="*60)
        logger.info("📊 TEST SUMMARY")
        logger.info("="*60)
        logger.info(f"Google Scholar: {len(gs_result.get('results', []))} results")
        logger.info(f"PubMed: {len(pm_result.get('results', []))} results")
        logger.info(f"arXiv: {len(ar_result.get('results', []))} results")
        logger.info(f"Total: {total_results} results")
        
        if total_results > 0:
            logger.info("✅ Web search tools are working correctly!")
            return True
        else:
            logger.warning("⚠️  No results found - check web search functionality")
            return False
            
    except Exception as e:
        logger.error(f"❌ Error testing search tools: {e}")
        logger.error(f"Error type: {type(e).__name__}")
        return False

async def main():
    """Main test function"""
    logger.info("=" * 60)
    logger.info("🚀 TESTING REAL WEB SEARCH TOOLS")
    logger.info("=" * 60)
    
    # 1. Check environment
    if not setup_environment():
        logger.error("❌ Environment setup failed. Aborting tests.")
        return
    
    # 2. Test tools
    success = await test_search_tools()
    
    # 3. Final result
    logger.info("\n" + "=" * 60)
    if success:
        logger.info("🎉 ALL TESTS PASSED! Real web search tools are working")
    else:
        logger.info("❌ TESTS FAILED! Check the implementation")
    logger.info("=" * 60)

if __name__ == "__main__":
    print("🔧 Testing Real Web Search Tools")
    print("=" * 60)
    
    # Run tests
    asyncio.run(main())