From e1cb92862ebf0dd50ce68eda194cd250b03dbcff Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 21 Oct 2025 16:58:51 -0700 Subject: [PATCH] [Feat] Add def search() APIs for Web Search - Perplexity API (#15769) * add BaseSearchConfig * add BaseSearchConfig * validate_environment * fix handlers * add PerplexitySearchConfig * add PerplexitySearchConfig * add LiteLLM Search API module. * add BaseSearchConfig * add _build_search_optional_params * add search_testing * add BaseSearchTest * add TestPerplexitySearch * fix BASE * fix handler * add search API * add to init * fix: working perplexity search API * add _hidden_params to search --- .circleci/config.yml | 53 ++- litellm/__init__.py | 1 + litellm/llms/base_llm/search/__init__.py | 15 + .../llms/base_llm/search/transformation.py | 120 +++++++ litellm/llms/custom_httpx/llm_http_handler.py | 165 +++++++++ .../llms/perplexity/search/transformation.py | 154 +++++++++ litellm/search/__init__.py | 7 + litellm/search/main.py | 313 ++++++++++++++++++ litellm/utils.py | 20 ++ tests/search_tests/__init__.py | 4 + tests/search_tests/base_search_unit_tests.py | 139 ++++++++ tests/search_tests/test_perplexity_search.py | 25 ++ 12 files changed, 1015 insertions(+), 1 deletion(-) create mode 100644 litellm/llms/base_llm/search/__init__.py create mode 100644 litellm/llms/base_llm/search/transformation.py create mode 100644 litellm/llms/perplexity/search/transformation.py create mode 100644 litellm/search/__init__.py create mode 100644 litellm/search/main.py create mode 100644 tests/search_tests/__init__.py create mode 100644 tests/search_tests/base_search_unit_tests.py create mode 100644 tests/search_tests/test_perplexity_search.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 6ffc09cf6e..5449016013 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1084,6 +1084,49 @@ jobs: paths: - ocr_coverage.xml - ocr_coverage + search_testing: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + + steps: + - checkout + - setup_google_dns + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-cov==5.0.0" + pip install "pytest-asyncio==0.21.1" + pip install "respx==0.22.0" + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/search_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml search_coverage.xml + mv .coverage search_coverage + + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - search_coverage.xml + - search_coverage litellm_mapped_tests: docker: - image: cimg/python:3.11 @@ -2827,7 +2870,7 @@ jobs: python -m venv venv . venv/bin/activate pip install coverage - coverage combine llm_translation_coverage llm_responses_api_coverage ocr_coverage mcp_coverage logging_coverage audio_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_security_tests_coverage guardrails_coverage + coverage combine llm_translation_coverage llm_responses_api_coverage ocr_coverage search_coverage mcp_coverage logging_coverage audio_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_security_tests_coverage guardrails_coverage coverage xml - codecov/upload: file: ./coverage.xml @@ -3381,6 +3424,12 @@ workflows: only: - main - /litellm_.*/ + - search_testing: + filters: + branches: + only: + - main + - /litellm_.*/ - litellm_mapped_enterprise_tests: filters: branches: @@ -3437,6 +3486,7 @@ workflows: - guardrails_testing - llm_responses_api_testing - ocr_testing + - search_testing - litellm_mapped_tests - litellm_mapped_enterprise_tests - batches_testing @@ -3501,6 +3551,7 @@ workflows: - google_generate_content_endpoint_testing - llm_responses_api_testing - ocr_testing + - search_testing - litellm_mapped_tests - litellm_mapped_enterprise_tests - batches_testing diff --git a/litellm/__init__.py b/litellm/__init__.py index a871279d37..21d2da2329 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1334,6 +1334,7 @@ from .rerank_api.main import * from .llms.anthropic.experimental_pass_through.messages.handler import * from .responses.main import * from .ocr.main import * +from .search.main import * from .realtime_api.main import _arealtime from .fine_tuning.main import * from .files.main import * diff --git a/litellm/llms/base_llm/search/__init__.py b/litellm/llms/base_llm/search/__init__.py new file mode 100644 index 0000000000..5a46482ed4 --- /dev/null +++ b/litellm/llms/base_llm/search/__init__.py @@ -0,0 +1,15 @@ +""" +Base Search API module. +""" +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) + +__all__ = [ + "BaseSearchConfig", + "SearchResponse", + "SearchResult", +] + diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py new file mode 100644 index 0000000000..a1a04bc0f7 --- /dev/null +++ b/litellm/llms/base_llm/search/transformation.py @@ -0,0 +1,120 @@ +""" +Base Search transformation configuration. +""" +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import httpx +from pydantic import PrivateAttr + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.base import LiteLLMPydanticObjectBase + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class SearchResult(LiteLLMPydanticObjectBase): + """Single search result.""" + title: str + url: str + snippet: str + date: Optional[str] = None + last_updated: Optional[str] = None + + model_config = {"extra": "allow"} + + +class SearchResponse(LiteLLMPydanticObjectBase): + """ + Standard Search response format. + Standardized to Perplexity Search format - other providers should transform to this format. + """ + results: List[SearchResult] + object: str = "search" + + model_config = {"extra": "allow"} + + # Define private attributes using PrivateAttr + _hidden_params: dict = PrivateAttr(default_factory=dict) + + +class BaseSearchConfig: + """ + Base configuration for Search transformations. + Handles provider-agnostic Search operations. + """ + + def __init__(self) -> None: + pass + + def validate_environment( + self, + headers: Dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + **kwargs, + ) -> Dict: + """ + Validate environment and return headers. + Override in provider-specific implementations. + """ + return headers + + def get_complete_url( + self, + api_base: Optional[str], + optional_params: dict, + **kwargs, + ) -> str: + """ + Get complete URL for Search endpoint. + Override in provider-specific implementations. + """ + raise NotImplementedError("get_complete_url must be implemented by provider") + + def transform_search_request( + self, + query: Union[str, List[str]], + optional_params: dict, + **kwargs, + ) -> Dict: + """ + Transform Search request to provider-specific format. + Override in provider-specific implementations. + + Args: + query: Search query (string or list of strings) + optional_params: Optional parameters for the request + + Returns: + Dict with request data + """ + raise NotImplementedError("transform_search_request must be implemented by provider") + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs, + ) -> SearchResponse: + """ + Transform provider-specific Search response to standard format. + Override in provider-specific implementations. + """ + raise NotImplementedError("transform_search_response must be implemented by provider") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict, + ) -> Exception: + """Get appropriate error class for the provider.""" + return BaseLLMException( + status_code=status_code, + message=error_message, + headers=headers, + ) + diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 28fb5f0269..9c0f35c739 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -44,6 +44,10 @@ from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, +) from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, ) @@ -1545,6 +1549,166 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + + def search( + self, + query: Union[str, List[str]], + optional_params: dict, + timeout: Union[float, httpx.Timeout], + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + api_base: Optional[str], + custom_llm_provider: str, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + asearch: bool = False, + headers: Optional[Dict[str, Any]] = None, + provider_config: Optional[BaseSearchConfig] = None, + ) -> Union[SearchResponse, Coroutine[Any, Any, SearchResponse]]: + """ + Sync Search handler. + """ + if provider_config is None: + raise ValueError( + f"No provider config found for provider: {custom_llm_provider}" + ) + + if asearch is True: + return self.async_search( + query=query, + optional_params=optional_params, + timeout=timeout, + logging_obj=logging_obj, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + client=client, + headers=headers, + provider_config=provider_config, + ) + + # Validate environment and get headers + headers = provider_config.validate_environment( + api_key=api_key, + api_base=api_base, + headers=headers or {}, + ) + + # Get complete URL + complete_url = provider_config.get_complete_url( + api_base=api_base, + optional_params=optional_params, + ) + + # Transform the request + data = provider_config.transform_search_request( + query=query, + optional_params=optional_params, + ) + + ## LOGGING + logging_obj.pre_call( + input=query if isinstance(query, str) else str(query), + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + if client is None or not isinstance(client, HTTPHandler): + client = _get_httpx_client() + + try: + # Make the POST request with JSON data + response = client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=provider_config) + + return provider_config.transform_search_response( + raw_response=response, + logging_obj=logging_obj, + ) + + async def async_search( + self, + query: Union[str, List[str]], + optional_params: dict, + timeout: Union[float, httpx.Timeout], + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + api_base: Optional[str], + custom_llm_provider: str, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + headers: Optional[Dict[str, Any]] = None, + provider_config: Optional[BaseSearchConfig] = None, + ) -> SearchResponse: + """ + Async Search handler. + """ + if provider_config is None: + raise ValueError( + f"No provider config found for provider: {custom_llm_provider}" + ) + + # Validate environment and get headers + headers = provider_config.validate_environment( + api_key=api_key, + api_base=api_base, + headers=headers or {}, + ) + + # Get complete URL + complete_url = provider_config.get_complete_url( + api_base=api_base, + optional_params=optional_params, + ) + + # Transform the request + data = provider_config.transform_search_request( + query=query, + optional_params=optional_params, + ) + + ## LOGGING + logging_obj.pre_call( + input=query if isinstance(query, str) else str(query), + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + ) + else: + async_httpx_client = client + + try: + # Make the async POST request with JSON data + response = await async_httpx_client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=provider_config) + + return provider_config.transform_search_response( + raw_response=response, + logging_obj=logging_obj, + ) + async def async_anthropic_messages_handler( self, model: str, @@ -3285,6 +3449,7 @@ class BaseLLMHTTPHandler: BaseAnthropicMessagesConfig, BaseBatchesConfig, BaseOCRConfig, + BaseSearchConfig, BaseTextToSpeechConfig, "BasePassthroughConfig", ], diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py new file mode 100644 index 0000000000..ef1ab9daad --- /dev/null +++ b/litellm/llms/perplexity/search/transformation.py @@ -0,0 +1,154 @@ +""" +Calls Perplexity's /search endpoint to search the web. +""" +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 _PerplexitySearchRequestRequired(TypedDict): + """Required fields for Perplexity Search API request.""" + query: Union[str, List[str]] # Required - search query or queries + + +class PerplexitySearchRequest(_PerplexitySearchRequestRequired, total=False): + """ + Perplexity Search API request format. + Based on: https://docs.perplexity.ai/api-reference/search-post + """ + max_results: int # Optional - maximum number of results (1-20), default 10 + search_domain_filter: List[str] # Optional - list of domains to filter (max 20) + max_tokens_per_page: int # Optional - max tokens per page, default 1024 + country: str # Optional - country code filter (e.g., 'US', 'GB', 'DE') + + +class PerplexitySearchConfig(BaseSearchConfig): + PERPLEXITY_API_BASE = "https://api.perplexity.ai" + + 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("PERPLEXITYAI_API_KEY") + if not api_key: + raise ValueError("PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_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, + **kwargs, + ) -> str: + """ + Get complete URL for Search endpoint. + """ + api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_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 Perplexity API format. + + Note: LiteLLM's native spec is the perplexity search spec. + + There's no transformation needed for the request data. + + https://docs.perplexity.ai/api-reference/search-post + + Args: + query: Search query (string or list of strings) + optional_params: Optional parameters for the request + - max_results: Maximum number of search results (1-20) + - search_domain_filter: List of domains to filter (max 20) + - max_tokens_per_page: Max tokens per page (default 1024) + - country: Country code filter (e.g., 'US', 'GB', 'DE') + + Returns: + Dict with typed request data following PerplexitySearchRequest spec + """ + request_data: PerplexitySearchRequest = { + "query": query, + } + + # Add optional parameters following Perplexity API spec (only if not None) + max_results = optional_params.get("max_results") + if max_results is not None: + request_data["max_results"] = max_results + + search_domain_filter = optional_params.get("search_domain_filter") + if search_domain_filter is not None: + request_data["search_domain_filter"] = search_domain_filter + + max_tokens_per_page = optional_params.get("max_tokens_per_page") + if max_tokens_per_page is not None: + request_data["max_tokens_per_page"] = max_tokens_per_page + + country = optional_params.get("country") + if country is not None: + request_data["country"] = country + + return dict(request_data) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs, + ) -> SearchResponse: + """ + Transform Perplexity API response to standard SearchResponse format. + + Args: + raw_response: Raw httpx response from Perplexity 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("snippet", ""), + date=result.get("date"), + last_updated=result.get("last_updated"), + ) + results.append(search_result) + + return SearchResponse( + results=results, + object="search", + ) + diff --git a/litellm/search/__init__.py b/litellm/search/__init__.py new file mode 100644 index 0000000000..b654fc66a3 --- /dev/null +++ b/litellm/search/__init__.py @@ -0,0 +1,7 @@ +""" +LiteLLM Search API module. +""" +from litellm.search.main import asearch, search + +__all__ = ["search", "asearch"] + diff --git a/litellm/search/main.py b/litellm/search/main.py new file mode 100644 index 0000000000..d90861fadc --- /dev/null +++ b/litellm/search/main.py @@ -0,0 +1,313 @@ +""" +Main Search function for LiteLLM. +""" +import asyncio +import contextvars +from functools import partial +from typing import Any, Coroutine, Dict, List, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.constants import request_timeout +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.utils import ProviderConfigManager, client + +####### ENVIRONMENT VARIABLES ################### +base_llm_http_handler = BaseLLMHTTPHandler() +################################################# + + +def _build_search_optional_params( + max_results: Optional[int] = None, + search_domain_filter: Optional[List[str]] = None, + max_tokens_per_page: Optional[int] = None, + country: Optional[str] = None, +) -> Dict[str, Any]: + """ + Helper function to build optional_params dict from Perplexity Search API parameters. + + Args: + max_results: Maximum number of results (1-20) + search_domain_filter: List of domains to filter (max 20) + max_tokens_per_page: Max tokens per page + country: Country code filter + + Returns: + Dict with non-None optional parameters + """ + optional_params: Dict[str, Any] = {} + + if max_results is not None: + optional_params["max_results"] = max_results + if search_domain_filter is not None: + optional_params["search_domain_filter"] = search_domain_filter + if max_tokens_per_page is not None: + optional_params["max_tokens_per_page"] = max_tokens_per_page + if country is not None: + optional_params["country"] = country + + return optional_params + + +@client +async def asearch( + query: Union[str, List[str]], + custom_llm_provider: str, + max_results: Optional[int] = None, + search_domain_filter: Optional[List[str]] = None, + max_tokens_per_page: Optional[int] = None, + country: Optional[str] = None, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: Optional[Dict[str, Any]] = None, + **kwargs, +) -> SearchResponse: + """ + Async Search function. + + Args: + query: Search query (string or list of strings) + custom_llm_provider: Provider name (e.g., "perplexity") + max_results: Optional maximum number of results (1-20), default 10 + search_domain_filter: Optional list of domains to filter (max 20) + max_tokens_per_page: Optional max tokens per page, default 1024 + country: Optional country code filter (e.g., 'US', 'GB', 'DE') + api_key: Optional API key + api_base: Optional API base URL + timeout: Optional timeout + extra_headers: Optional extra headers + **kwargs: Additional parameters + + Returns: + SearchResponse with results list following Perplexity format + + Example: + ```python + import litellm + + # Basic search + response = await litellm.asearch( + query="latest AI developments 2024", + custom_llm_provider="perplexity" + ) + + # Search with options + response = await litellm.asearch( + query="AI developments", + custom_llm_provider="perplexity", + max_results=10, + search_domain_filter=["arxiv.org", "nature.com"], + max_tokens_per_page=1024, + country="US" + ) + + # Access results + for result in response.results: + print(f"{result.title}: {result.url}") + print(f"Snippet: {result.snippet}") + ``` + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["asearch"] = True + + func = partial( + search, + query=query, + custom_llm_provider=custom_llm_provider, + max_results=max_results, + search_domain_filter=search_domain_filter, + max_tokens_per_page=max_tokens_per_page, + country=country, + api_key=api_key, + api_base=api_base, + timeout=timeout, + extra_headers=extra_headers, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + if response is None: + raise ValueError( + f"Got an unexpected None response from the Search API: {response}" + ) + + return response + except Exception as e: + raise litellm.exception_type( + model="", + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def search( + query: Union[str, List[str]], + custom_llm_provider: str, + max_results: Optional[int] = None, + search_domain_filter: Optional[List[str]] = None, + max_tokens_per_page: Optional[int] = None, + country: Optional[str] = None, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: Optional[Dict[str, Any]] = None, + **kwargs, +) -> Union[SearchResponse, Coroutine[Any, Any, SearchResponse]]: + """ + Synchronous Search function. + + Args: + query: Search query (string or list of strings) + custom_llm_provider: Provider name (e.g., "perplexity") + max_results: Optional maximum number of results (1-20), default 10 + search_domain_filter: Optional list of domains to filter (max 20) + max_tokens_per_page: Optional max tokens per page, default 1024 + country: Optional country code filter (e.g., 'US', 'GB', 'DE') + api_key: Optional API key + api_base: Optional API base URL + timeout: Optional timeout + extra_headers: Optional extra headers + **kwargs: Additional parameters + + Returns: + SearchResponse with results list following Perplexity format + + Example: + ```python + import litellm + + # Basic search + response = litellm.search( + query="latest AI developments 2024", + custom_llm_provider="perplexity" + ) + + # Search with options + response = litellm.search( + query="AI developments", + custom_llm_provider="perplexity", + max_results=10, + search_domain_filter=["arxiv.org", "nature.com"], + max_tokens_per_page=1024, + country="US" + ) + + # Multi-query search + response = litellm.search( + query=["AI developments", "machine learning trends"], + custom_llm_provider="perplexity" + ) + + # Access results + for result in response.results: + print(f"{result.title}: {result.url}") + print(f"Snippet: {result.snippet}") + if result.date: + print(f"Date: {result.date}") + ``` + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("asearch", False) is True + + # Validate query parameter + if not isinstance(query, (str, list)): + raise ValueError(f"query must be a string or list of strings, got {type(query)}") + + if isinstance(query, list) and not all(isinstance(q, str) for q in query): + raise ValueError("All items in query list must be strings") + + # Get provider config + search_provider_config: Optional[BaseSearchConfig] = ( + ProviderConfigManager.get_provider_search_config( + provider=litellm.LlmProviders(custom_llm_provider), + ) + ) + + if search_provider_config is None: + raise ValueError( + f"Search is not supported for provider: {custom_llm_provider}" + ) + + verbose_logger.debug( + f"Search call - provider: {custom_llm_provider}" + ) + + # Build optional_params from explicit parameters + optional_params = _build_search_optional_params( + max_results=max_results, + search_domain_filter=search_domain_filter, + max_tokens_per_page=max_tokens_per_page, + country=country, + ) + + verbose_logger.debug(f"Search optional_params: {optional_params}") + + # Validate environment and get headers + headers = search_provider_config.validate_environment( + api_key=api_key, + api_base=api_base, + headers=extra_headers or {}, + ) + + # Get complete URL + complete_url = search_provider_config.get_complete_url( + api_base=api_base, + optional_params=optional_params, + ) + + # Pre Call logging + litellm_logging_obj.update_environment_variables( + model="", + optional_params=optional_params, + litellm_params={ + "litellm_call_id": litellm_call_id, + "api_base": complete_url, + }, + custom_llm_provider=custom_llm_provider, + ) + + # Call the handler + response = base_llm_http_handler.search( + query=query, + optional_params=optional_params, + timeout=timeout or request_timeout, + logging_obj=litellm_logging_obj, + api_key=api_key, + api_base=complete_url, + custom_llm_provider=custom_llm_provider, + asearch=_is_async, + headers=headers, + provider_config=search_provider_config, + ) + + return response + except Exception as e: + raise litellm.exception_type( + model="", + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + diff --git a/litellm/utils.py b/litellm/utils.py index c2e245acb5..1de29596f9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -144,6 +144,7 @@ from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig +from litellm.llms.base_llm.search.transformation import BaseSearchConfig from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, ) @@ -7619,6 +7620,25 @@ class ProviderConfigManager: return None return config_class() + @staticmethod + def get_provider_search_config( + provider: LlmProviders, + ) -> Optional["BaseSearchConfig"]: + """ + Get Search configuration for a given provider. + """ + from litellm.llms.perplexity.search.transformation import ( + PerplexitySearchConfig, + ) + + PROVIDER_TO_CONFIG_MAP = { + litellm.LlmProviders.PERPLEXITY: PerplexitySearchConfig, + } + config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None) + if config_class is None: + return None + return config_class() + @staticmethod def get_provider_text_to_speech_config( model: str, diff --git a/tests/search_tests/__init__.py b/tests/search_tests/__init__.py new file mode 100644 index 0000000000..9e0d0be6af --- /dev/null +++ b/tests/search_tests/__init__.py @@ -0,0 +1,4 @@ +""" +Search API tests. +""" + diff --git a/tests/search_tests/base_search_unit_tests.py b/tests/search_tests/base_search_unit_tests.py new file mode 100644 index 0000000000..321692350e --- /dev/null +++ b/tests/search_tests/base_search_unit_tests.py @@ -0,0 +1,139 @@ +""" +Base test class for Search functionality across different providers. + +This follows the same pattern as BaseOCRTest in tests/ocr_tests/base_ocr_unit_tests.py +""" +import pytest +import litellm +from abc import ABC, abstractmethod +import json + + +class BaseSearchTest(ABC): + """ + Abstract base test class that enforces common Search tests across all providers. + + Each provider-specific test class should inherit from this and implement + get_custom_llm_provider() to return provider name. + """ + + @abstractmethod + def get_custom_llm_provider(self) -> str: + """Must return the custom_llm_provider for the specific provider""" + pass + + @pytest.fixture(autouse=True) + def _handle_rate_limits(self): + """Fixture to handle rate limit errors for all test methods""" + try: + yield + except litellm.RateLimitError: + pytest.skip("Rate limit exceeded") + except litellm.InternalServerError: + pytest.skip("Model is overloaded") + + @pytest.mark.asyncio + async def test_basic_search(self): + """ + Test basic search functionality with a simple query. + """ + litellm._turn_on_debug() + custom_llm_provider = self.get_custom_llm_provider() + print("Custom LLM Provider=", custom_llm_provider) + + try: + response = await litellm.asearch( + query="latest developments in AI", + custom_llm_provider=custom_llm_provider, + ) + print("Search response=", response.model_dump_json(indent=4)) + + print(f"\n{'='*80}") + print(f"Response type: {type(response)}") + print(f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}") + + # Check if response has expected Search format + assert hasattr(response, "results"), "Response should have 'results' attribute" + assert hasattr(response, "object"), "Response should have 'object' attribute" + assert response.object == "search", f"Expected object='search', got '{response.object}'" + + # Validate results structure + assert isinstance(response.results, list), "results should be a list" + assert len(response.results) > 0, "Should have at least one result" + + # Check first result structure + first_result = response.results[0] + assert hasattr(first_result, "title"), "Result should have 'title' attribute" + assert hasattr(first_result, "url"), "Result should have 'url' attribute" + assert hasattr(first_result, "snippet"), "Result should have 'snippet' attribute" + + print(f"Total results: {len(response.results)}") + print(f"First result title: {first_result.title}") + print(f"First result URL: {first_result.url}") + print(f"First result snippet: {first_result.snippet[:100]}...") + print(f"{'='*80}\n") + + assert len(first_result.title) > 0, "Title should not be empty" + assert len(first_result.url) > 0, "URL should not be empty" + assert len(first_result.snippet) > 0, "Snippet should not be empty" + + except Exception as e: + pytest.fail(f"Search call failed: {str(e)}") + + def test_search_response_structure(self): + """ + Test that the Search response has the correct structure. + """ + litellm.set_verbose = True + custom_llm_provider = self.get_custom_llm_provider() + + response = litellm.search( + query="artificial intelligence recent news", + custom_llm_provider=custom_llm_provider, + ) + + # Validate response structure + assert hasattr(response, "results"), "Response should have 'results' attribute" + assert hasattr(response, "object"), "Response should have 'object' attribute" + + assert isinstance(response.results, list), "results should be a list" + assert len(response.results) > 0, "Should have at least one result" + assert response.object == "search", "object should be 'search'" + + # Validate first result structure + first_result = response.results[0] + assert hasattr(first_result, "title"), "Result should have 'title' attribute" + assert hasattr(first_result, "url"), "Result should have 'url' attribute" + assert hasattr(first_result, "snippet"), "Result should have 'snippet' attribute" + assert isinstance(first_result.title, str), "title should be a string" + assert isinstance(first_result.url, str), "url should be a string" + assert isinstance(first_result.snippet, str), "snippet should be a string" + + print(f"\nResponse structure validated:") + print(f" - object: {response.object}") + print(f" - results: {len(response.results)}") + print(f" - first result has all required fields") + + def test_search_with_optional_params(self): + """ + Test search with optional parameters. + """ + litellm.set_verbose = True + custom_llm_provider = self.get_custom_llm_provider() + + response = litellm.search( + query="machine learning", + custom_llm_provider=custom_llm_provider, + max_results=5, + ) + + # Validate response + assert hasattr(response, "results"), "Response should have 'results' attribute" + assert isinstance(response.results, list), "results should be a list" + assert len(response.results) > 0, "Should have at least one result" + assert len(response.results) <= 5, "Should have at most 5 results as requested" + + print(f"\nSearch with optional params validated:") + print(f" - Requested max_results: 5") + print(f" - Received results: {len(response.results)}") + diff --git a/tests/search_tests/test_perplexity_search.py b/tests/search_tests/test_perplexity_search.py new file mode 100644 index 0000000000..e68e054c76 --- /dev/null +++ b/tests/search_tests/test_perplexity_search.py @@ -0,0 +1,25 @@ +""" +Tests for Perplexity Search API integration. +""" +import os +import sys +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) + +from tests.search_tests.base_search_unit_tests import BaseSearchTest + + +class TestPerplexitySearch(BaseSearchTest): + """ + Tests for Perplexity Search functionality. + """ + + def get_custom_llm_provider(self) -> str: + """ + Return custom_llm_provider for Perplexity Search. + """ + return "perplexity" +