test_tavily.py

import os
from dotenv import load_dotenv
from langchain_tavily import TavilySearch

# Load environment variables
load_dotenv()


def test_tavily_search():
    """Test the Tavily search functionality independently"""

    print("🔍 Testing Tavily Search Tool...")
    print("="*50)

    # Check if API key is available
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        print("❌ Error: TAVILY_API_KEY not found in environment variables")
        print("Please add TAVILY_API_KEY to your .env file")
        return

    print(f"✅ API Key found: {api_key[:10]}...{api_key[-4:]}")

    try:
        # Initialize Tavily search
        print("\n🚀 Initializing Tavily search...")
        tavily_search = TavilySearch(
            api_key=api_key,
            max_results=2,
            topic="general",
            search_depth="basic",
            include_answer=True,
            include_raw_content=False,
            include_images=False
        )
        print("✅ Tavily search initialized successfully")

        # Test search query
        test_query = "latest portugal won"
        print(f"\n🔎 Searching for: '{test_query}'")

        # Perform the search
        results = tavily_search.invoke({"query": test_query})

        print(f"\n📊 Results type: {type(results)}")
        print(
            f"📊 Results length: {len(results) if hasattr(results, '__len__') else 'N/A'}")

        # Display results
        print("\n📰 Raw Results:")
        print("-" * 30)
        print(results)
        print("-" * 30)

        # Process and format results
        if isinstance(results, list) and results:
            print(f"\n✅ Found {len(results)} results")
            formatted_news = f"📰 Latest News about '{test_query}':\n\n"

            for i, result in enumerate(results, 1):
                print(f"\n🔍 Processing result {i}:")
                print(f"   Type: {type(result)}")

                if isinstance(result, dict):
                    print(f"   Keys: {list(result.keys())}")
                    title = result.get('title', 'No title')
                    content = result.get('content', 'No content available')
                    url = result.get('url', 'No URL')

                    formatted_news += f"{i}. **{title}**\n"
                    formatted_news += f"   Summary: {content[:200]}...\n"
                    formatted_news += f"   Source: {url}\n\n"
                else:
                    # Handle case where result is a string
                    formatted_news += f"{i}. {str(result)[:300]}...\n\n"

            print(f"\n📝 Formatted Output:")
            print("="*50)
            print(formatted_news)
            print("="*50)

        else:
            print(f"❌ No results found or unexpected result format")
            print(f"Results: {results}")

    except Exception as e:
        print(f"❌ Error during search: {e}")
        print(f"Error type: {type(e)}")
        import traceback
        print(f"Full traceback:\n{traceback.format_exc()}")


def test_tavily_raw():
    """Test the raw Tavily search functionality"""

    print("🔍 Testing Raw Tavily Search...")
    print("="*50)

    # Check if API key is available
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        print("❌ Error: TAVILY_API_KEY not found in environment variables")
        return False

    try:
        # Initialize Tavily search
        tavily_search = TavilySearch(
            api_key=api_key,
            max_results=2,
            topic="general",
            search_depth="basic",
            include_answer=True,
            include_raw_content=False,
            include_images=False
        )

        # Test search query
        test_query = "latest portugal won"
        print(f"🔎 Searching for: '{test_query}'")

        # Perform the search
        results = tavily_search.invoke({"query": test_query})

        print(f"✅ Raw search successful!")
        print(f"📊 Results type: {type(results)}")
        if isinstance(results, dict) and 'results' in results:
            print(f"📊 Number of results: {len(results['results'])}")
            return True
        return False

    except Exception as e:
        print(f"❌ Raw search failed: {e}")
        return False


def test_tools_function():
    """Test our get_latest_news function from tools.py"""

    print("\n🔧 Testing get_latest_news Function...")
    print("="*50)

    try:
        # Import our function
        from tools import get_latest_news

        # Test the function
        test_query = "portugal won"
        print(f"🔎 Testing with query: '{test_query}'")

        result = get_latest_news(test_query)

        print("✅ Function executed successfully!")
        print("\n📰 Function Output:")
        print("-" * 50)
        print(result)
        print("-" * 50)

        return True

    except Exception as e:
        print(f"❌ Function test failed: {e}")
        import traceback
        print(f"Full traceback:\n{traceback.format_exc()}")
        return False


if __name__ == "__main__":
    print("🧪 Running Tavily Tests...\n")

    # Test 1: Raw Tavily
    raw_success = test_tavily_raw()

    # Test 2: Our function
    if raw_success:
        function_success = test_tools_function()

        if function_success:
            print("\n🎉 All tests passed! The news search tool is working correctly.")
        else:
            print("\n⚠️ Raw search works but our function has issues.")
    else:
        print("\n❌ Raw search failed - check your API key and connection.")