tests/test_search_tool.py

#!/usr/bin/env python3
"""
Simple test for the web search tool functionality
"""

import os
import asyncio
import sys
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Check if Tavily API key is configured
if not os.getenv("TAVILY_API_KEY"):
    print("❌ TAVILY_API_KEY not found in environment variables")
    print("   Get your API key from: https://tavily.com/")
    print("   Add it to your .env file: TAVILY_API_KEY=tvly-your-key-here")
    sys.exit(1)

async def test_search():
    """Test the web search functionality"""
    from langchain_tavily import TavilySearch
    
    # Initialize Tavily client
    tavily_search = TavilySearch(
        api_key=os.getenv("TAVILY_API_KEY"),
        max_results=5,
        topic="general",
        search_depth="advanced",
        include_answer=True,
        include_raw_content=False
    )
    
    print("🔍 Testing Tavily Web Search Tool")
    print("=" * 80)
    
    test_queries = [
        "GDPR latest updates 2024",
        "NIS2 directive implementation",
        "EU Cyber Resilience Act news"
    ]
    
    for i, query in enumerate(test_queries, 1):
        print(f"\n📝 Test {i}: Searching for '{query}'")
        print("-" * 80)
        
        try:
            result = await tavily_search.ainvoke({"query": query})
            print(str(result)[:500] + "..." if len(str(result)) > 500 else str(result))
        except Exception as e:
            print(f"❌ Error: {e}")
        
        print("\n" + "=" * 80)
    
    print("\n✅ Search tool test completed!")

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