src/tools/web_search.py

"""
Web Search Tool - Tavily and Exa implementations
Author: @mangobee
Date: 2026-01-02

Provides web search functionality with:
- Tavily as primary search (free tier: 1000 req/month)
- Exa as fallback (paid tier)
- Retry logic with exponential backoff
- Structured error handling
"""

import logging
from typing import Dict, List, Optional
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
)

from src.config.settings import Settings

# ============================================================================
# CONFIG
# ============================================================================
MAX_RETRIES = 3
RETRY_MIN_WAIT = 1  # seconds
RETRY_MAX_WAIT = 10  # seconds
DEFAULT_MAX_RESULTS = 5

# ============================================================================
# Logging Setup
# ============================================================================
logger = logging.getLogger(__name__)


# ============================================================================
# Tavily Search Implementation
# ============================================================================


@retry(
    stop=stop_after_attempt(MAX_RETRIES),
    wait=wait_exponential(multiplier=1, min=RETRY_MIN_WAIT, max=RETRY_MAX_WAIT),
    retry=retry_if_exception_type((ConnectionError, TimeoutError)),
    reraise=True,
)
def tavily_search(query: str, max_results: int = DEFAULT_MAX_RESULTS) -> Dict:
    """
    Search using Tavily API with retry logic.

    Args:
        query: Search query string
        max_results: Maximum number of results to return (default: 5)

    Returns:
        Dict with structure: {
            "results": [{"title": str, "url": str, "snippet": str}, ...],
            "source": "tavily",
            "query": str,
            "count": int
        }

    Raises:
        ValueError: If API key not configured
        ConnectionError: If API connection fails after retries
        Exception: For other API errors
    """
    try:
        from tavily import TavilyClient

        settings = Settings()
        api_key = settings.tavily_api_key

        if not api_key:
            raise ValueError("TAVILY_API_KEY not configured in settings")

        logger.info(f"Tavily search: query='{query}', max_results={max_results}")

        client = TavilyClient(api_key=api_key)
        response = client.search(query=query, max_results=max_results)

        # Extract and structure results
        results = []
        for item in response.get("results", []):
            results.append(
                {
                    "title": item.get("title", ""),
                    "url": item.get("url", ""),
                    "snippet": item.get("content", ""),
                }
            )

        logger.info(f"Tavily search successful: {len(results)} results")

        return {
            "results": results,
            "source": "tavily",
            "query": query,
            "count": len(results),
        }

    except ValueError as e:
        logger.error(f"Tavily configuration error: {e}")
        raise
    except (ConnectionError, TimeoutError) as e:
        logger.warning(f"Tavily connection error (will retry): {e}")
        raise
    except Exception as e:
        logger.error(f"Tavily search error: {e}")
        raise Exception(f"Tavily search failed: {str(e)}")


# ============================================================================
# Exa Search Implementation
# ============================================================================


@retry(
    stop=stop_after_attempt(MAX_RETRIES),
    wait=wait_exponential(multiplier=1, min=RETRY_MIN_WAIT, max=RETRY_MAX_WAIT),
    retry=retry_if_exception_type((ConnectionError, TimeoutError)),
    reraise=True,
)
def exa_search(query: str, max_results: int = DEFAULT_MAX_RESULTS) -> Dict:
    """
    Search using Exa API with retry logic.

    Args:
        query: Search query string
        max_results: Maximum number of results to return (default: 5)

    Returns:
        Dict with structure: {
            "results": [{"title": str, "url": str, "snippet": str}, ...],
            "source": "exa",
            "query": str,
            "count": int
        }

    Raises:
        ValueError: If API key not configured
        ConnectionError: If API connection fails after retries
        Exception: For other API errors
    """
    try:
        from exa_py import Exa

        settings = Settings()
        api_key = settings.exa_api_key

        if not api_key:
            raise ValueError("EXA_API_KEY not configured in settings")

        logger.info(f"Exa search: query='{query}', max_results={max_results}")

        client = Exa(api_key=api_key)
        response = client.search(
            query=query, num_results=max_results, use_autoprompt=True
        )

        # Extract and structure results
        results = []
        for item in response.results:
            results.append(
                {
                    "title": item.title if hasattr(item, "title") else "",
                    "url": item.url if hasattr(item, "url") else "",
                    "snippet": item.text if hasattr(item, "text") else "",
                }
            )

        logger.info(f"Exa search successful: {len(results)} results")

        return {
            "results": results,
            "source": "exa",
            "query": query,
            "count": len(results),
        }

    except ValueError as e:
        logger.error(f"Exa configuration error: {e}")
        raise
    except (ConnectionError, TimeoutError) as e:
        logger.warning(f"Exa connection error (will retry): {e}")
        raise
    except Exception as e:
        logger.error(f"Exa search error: {e}")
        raise Exception(f"Exa search failed: {str(e)}")


# ============================================================================
# Unified Search with Fallback
# ============================================================================


def search(query: str, max_results: int = DEFAULT_MAX_RESULTS) -> Dict:
    """
    Unified search function with automatic fallback.

    Tries Tavily first (free tier), falls back to Exa if Tavily fails.

    Args:
        query: Search query string
        max_results: Maximum number of results to return (default: 5)

    Returns:
        Dict with search results from either Tavily or Exa

    Raises:
        Exception: If both Tavily and Exa searches fail
    """
    settings = Settings()
    default_tool = settings.default_search_tool

    # Try default tool first
    if default_tool == "tavily":
        try:
            return tavily_search(query, max_results)
        except Exception as e:
            logger.warning(f"Tavily failed, falling back to Exa: {e}")
            try:
                return exa_search(query, max_results)
            except Exception as exa_error:
                logger.error(f"Both Tavily and Exa failed")
                raise Exception(f"Search failed - Tavily: {e}, Exa: {exa_error}")
    else:
        # Default is Exa
        try:
            return exa_search(query, max_results)
        except Exception as e:
            logger.warning(f"Exa failed, falling back to Tavily: {e}")
            try:
                return tavily_search(query, max_results)
            except Exception as tavily_error:
                logger.error(f"Both Exa and Tavily failed")
                raise Exception(f"Search failed - Exa: {e}, Tavily: {tavily_error}")