lijiaoqiao/llm-gateway-competitors/litellm-wheel-src/litellm/llms/tavily/search/transformation.py

"""
Calls Tavily's /search endpoint to search the web.

Tavily API Reference: https://docs.tavily.com/documentation/api-reference/endpoint/search
"""
from typing import Dict, List, Optional, TypedDict, Union

import httpx

from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
    BaseSearchConfig,
    SearchResponse,
    SearchResult,
)
from litellm.secret_managers.main import get_secret_str


class _TavilySearchRequestRequired(TypedDict):
    """Required fields for Tavily Search API request."""

    query: str  # Required - search query


class TavilySearchRequest(_TavilySearchRequestRequired, total=False):
    """
    Tavily Search API request format.
    Based on: https://docs.tavily.com/documentation/api-reference/endpoint/search
    """

    max_results: int  # Optional - maximum number of results (0-20), default 5
    include_domains: List[str]  # Optional - list of domains to include (max 300)
    exclude_domains: List[str]  # Optional - list of domains to exclude (max 150)
    topic: str  # Optional - category of search ('general', 'news', 'finance'), default 'general'
    search_depth: str  # Optional - depth of search ('basic', 'advanced'), default 'basic'
    include_answer: Union[bool, str]  # Optional - include LLM-generated answer
    include_raw_content: Union[bool, str]  # Optional - include raw HTML content
    include_images: bool  # Optional - perform image search
    include_image_descriptions: bool  # Optional - add descriptions for images
    include_favicon: bool  # Optional - include favicon URL
    time_range: str  # Optional - time range filter ('day', 'week', 'month', 'year', 'd', 'w', 'm', 'y')
    start_date: str  # Optional - start date filter (YYYY-MM-DD)
    end_date: str  # Optional - end date filter (YYYY-MM-DD)
    country: str  # Optional - country code filter (e.g., 'US', 'GB', 'DE')


class TavilySearchConfig(BaseSearchConfig):
    TAVILY_API_BASE = "https://api.tavily.com"

    @staticmethod
    def ui_friendly_name() -> str:
        return "Tavily"

    def validate_environment(
        self,
        headers: Dict,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        **kwargs,
    ) -> Dict:
        """
        Validate environment and return headers.
        """
        api_key = api_key or get_secret_str("TAVILY_API_KEY")
        if not api_key:
            raise ValueError(
                "TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable."
            )
        headers["Authorization"] = f"Bearer {api_key}"
        headers["Content-Type"] = "application/json"
        return headers

    def get_complete_url(
        self,
        api_base: Optional[str],
        optional_params: dict,
        data: Optional[Union[Dict, List[Dict]]] = None,
        **kwargs,
    ) -> str:
        """
        Get complete URL for Search endpoint.
        """
        api_base = api_base or get_secret_str("TAVILY_API_BASE") or self.TAVILY_API_BASE

        # Append "/search" to the api base if it's not already there
        if not api_base.endswith("/search"):
            api_base = f"{api_base}/search"

        return api_base

    def transform_search_request(
        self,
        query: Union[str, List[str]],
        optional_params: dict,
        **kwargs,
    ) -> Dict:
        """
        Transform Search request to Tavily API format.

        Args:
            query: Search query (string or list of strings). Tavily only supports single string queries.
            optional_params: Optional parameters for the request
                - max_results: Maximum number of search results (0-20)
                - search_domain_filter: List of domains to include (max 300) -> maps to `include_domains`
                - exclude_domains: List of domains to exclude (max 150)
                - topic: Category of search ('general', 'news', 'finance')
                - search_depth: Depth of search ('basic', 'advanced')
                - include_answer: Include LLM-generated answer (bool or 'basic', 'advanced')
                - include_raw_content: Include raw HTML content (bool or 'markdown', 'text')
                - include_images: Perform image search (bool)
                - include_image_descriptions: Add descriptions for images (bool)
                - include_favicon: Include favicon URL (bool)
                - time_range: Time range filter ('day', 'week', 'month', 'year', 'd', 'w', 'm', 'y')
                - start_date: Start date filter (YYYY-MM-DD)
                - end_date: End date filter (YYYY-MM-DD)
                - country: Country code filter (e.g., 'US', 'GB', 'DE')

        Returns:
            Dict with typed request data following TavilySearchRequest spec
        """
        if isinstance(query, list):
            # Tavily only supports single string queries
            query = " ".join(query)

        request_data: TavilySearchRequest = {
            "query": query,
        }

        # Transform Perplexity unified spec parameters to Tavily format
        if "max_results" in optional_params:
            request_data["max_results"] = optional_params["max_results"]

        if "search_domain_filter" in optional_params:
            request_data["include_domains"] = optional_params["search_domain_filter"]

        if "country" in optional_params:
            # Tavily expects lowercase country names
            request_data["country"] = optional_params["country"].lower()

        # Convert to dict before dynamic key assignments
        result_data = dict(request_data)

        # pass through all other parameters as-is
        for param, value in optional_params.items():
            if (
                param not in self.get_supported_perplexity_optional_params()
                and param not in result_data
            ):
                result_data[param] = value

        return result_data

    def transform_search_response(
        self,
        raw_response: httpx.Response,
        logging_obj: LiteLLMLoggingObj,
        **kwargs,
    ) -> SearchResponse:
        """
        Transform Tavily API response to LiteLLM unified SearchResponse format.

        Tavily → LiteLLM mappings:
        - results[].title → SearchResult.title
        - results[].url → SearchResult.url
        - results[].content → SearchResult.snippet
        - No date/last_updated fields in Tavily response (set to None)

        Args:
            raw_response: Raw httpx response from Tavily API
            logging_obj: Logging object for tracking

        Returns:
            SearchResponse with standardized format
        """
        response_json = raw_response.json()

        # Transform results to SearchResult objects
        results = []
        for result in response_json.get("results", []):
            search_result = SearchResult(
                title=result.get("title", ""),
                url=result.get("url", ""),
                snippet=result.get(
                    "content", ""
                ),  # Tavily uses "content" instead of "snippet"
                date=None,  # Tavily doesn't provide date in response
                last_updated=None,  # Tavily doesn't provide last_updated in response
            )
            results.append(search_result)

        return SearchResponse(
            results=results,
            object="search",
        )