From d9b85ab276cb1237fc574833e633f9b5732a49e7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 21 Oct 2025 17:42:18 -0700 Subject: [PATCH] fix: rename search_provider --- litellm/search/main.py | 39 +++++++++++-------- tests/search_tests/base_search_unit_tests.py | 20 +++++----- tests/search_tests/test_exa_ai_search.py | 4 +- tests/search_tests/test_parallel_ai_search.py | 4 +- tests/search_tests/test_perplexity_search.py | 4 +- tests/search_tests/test_tavily_search.py | 4 +- 6 files changed, 40 insertions(+), 35 deletions(-) diff --git a/litellm/search/main.py b/litellm/search/main.py index d90861fadc..e0d5f5841f 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -56,7 +56,7 @@ def _build_search_optional_params( @client async def asearch( query: Union[str, List[str]], - custom_llm_provider: str, + search_provider: str, max_results: Optional[int] = None, search_domain_filter: Optional[List[str]] = None, max_tokens_per_page: Optional[int] = None, @@ -72,7 +72,7 @@ async def asearch( Args: query: Search query (string or list of strings) - custom_llm_provider: Provider name (e.g., "perplexity") + search_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 @@ -93,13 +93,13 @@ async def asearch( # Basic search response = await litellm.asearch( query="latest AI developments 2024", - custom_llm_provider="perplexity" + search_provider="perplexity" ) # Search with options response = await litellm.asearch( query="AI developments", - custom_llm_provider="perplexity", + search_provider="perplexity", max_results=10, search_domain_filter=["arxiv.org", "nature.com"], max_tokens_per_page=1024, @@ -120,7 +120,7 @@ async def asearch( func = partial( search, query=query, - custom_llm_provider=custom_llm_provider, + search_provider=search_provider, max_results=max_results, search_domain_filter=search_domain_filter, max_tokens_per_page=max_tokens_per_page, @@ -150,7 +150,7 @@ async def asearch( except Exception as e: raise litellm.exception_type( model="", - custom_llm_provider=custom_llm_provider, + custom_llm_provider=search_provider, original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, @@ -160,7 +160,7 @@ async def asearch( @client def search( query: Union[str, List[str]], - custom_llm_provider: str, + search_provider: str, max_results: Optional[int] = None, search_domain_filter: Optional[List[str]] = None, max_tokens_per_page: Optional[int] = None, @@ -176,7 +176,7 @@ def search( Args: query: Search query (string or list of strings) - custom_llm_provider: Provider name (e.g., "perplexity") + search_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 @@ -197,13 +197,13 @@ def search( # Basic search response = litellm.search( query="latest AI developments 2024", - custom_llm_provider="perplexity" + search_provider="perplexity" ) # Search with options response = litellm.search( query="AI developments", - custom_llm_provider="perplexity", + search_provider="perplexity", max_results=10, search_domain_filter=["arxiv.org", "nature.com"], max_tokens_per_page=1024, @@ -213,7 +213,7 @@ def search( # Multi-query search response = litellm.search( query=["AI developments", "machine learning trends"], - custom_llm_provider="perplexity" + search_provider="perplexity" ) # Access results @@ -240,17 +240,17 @@ def search( # Get provider config search_provider_config: Optional[BaseSearchConfig] = ( ProviderConfigManager.get_provider_search_config( - provider=litellm.LlmProviders(custom_llm_provider), + provider=litellm.LlmProviders(search_provider), ) ) if search_provider_config is None: raise ValueError( - f"Search is not supported for provider: {custom_llm_provider}" + f"Search is not supported for provider: {search_provider}" ) verbose_logger.debug( - f"Search call - provider: {custom_llm_provider}" + f"Search call - provider: {search_provider}" ) # Build optional_params from explicit parameters @@ -261,6 +261,11 @@ def search( country=country, ) + # Add remaining kwargs to optional_params (for provider-specific params) + for key, value in kwargs.items(): + if key not in optional_params: + optional_params[key] = value + verbose_logger.debug(f"Search optional_params: {optional_params}") # Validate environment and get headers @@ -284,7 +289,7 @@ def search( "litellm_call_id": litellm_call_id, "api_base": complete_url, }, - custom_llm_provider=custom_llm_provider, + custom_llm_provider=search_provider, ) # Call the handler @@ -295,7 +300,7 @@ def search( logging_obj=litellm_logging_obj, api_key=api_key, api_base=complete_url, - custom_llm_provider=custom_llm_provider, + custom_llm_provider=search_provider, asearch=_is_async, headers=headers, provider_config=search_provider_config, @@ -305,7 +310,7 @@ def search( except Exception as e: raise litellm.exception_type( model="", - custom_llm_provider=custom_llm_provider, + custom_llm_provider=search_provider, original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, diff --git a/tests/search_tests/base_search_unit_tests.py b/tests/search_tests/base_search_unit_tests.py index 321692350e..619d1fff9a 100644 --- a/tests/search_tests/base_search_unit_tests.py +++ b/tests/search_tests/base_search_unit_tests.py @@ -14,12 +14,12 @@ 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. + get_search_provider() to return provider name. """ @abstractmethod - def get_custom_llm_provider(self) -> str: - """Must return the custom_llm_provider for the specific provider""" + def get_search_provider(self) -> str: + """Must return the search_provider for the specific provider""" pass @pytest.fixture(autouse=True) @@ -38,13 +38,13 @@ class BaseSearchTest(ABC): 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) + search_provider = self.get_search_provider() + print("Search Provider=", search_provider) try: response = await litellm.asearch( query="latest developments in AI", - custom_llm_provider=custom_llm_provider, + search_provider=search_provider, ) print("Search response=", response.model_dump_json(indent=4)) @@ -85,11 +85,11 @@ class BaseSearchTest(ABC): Test that the Search response has the correct structure. """ litellm.set_verbose = True - custom_llm_provider = self.get_custom_llm_provider() + search_provider = self.get_search_provider() response = litellm.search( query="artificial intelligence recent news", - custom_llm_provider=custom_llm_provider, + search_provider=search_provider, ) # Validate response structure @@ -119,11 +119,11 @@ class BaseSearchTest(ABC): Test search with optional parameters. """ litellm.set_verbose = True - custom_llm_provider = self.get_custom_llm_provider() + search_provider = self.get_search_provider() response = litellm.search( query="machine learning", - custom_llm_provider=custom_llm_provider, + search_provider=search_provider, max_results=5, ) diff --git a/tests/search_tests/test_exa_ai_search.py b/tests/search_tests/test_exa_ai_search.py index f2455e5d4b..60b4eb0389 100644 --- a/tests/search_tests/test_exa_ai_search.py +++ b/tests/search_tests/test_exa_ai_search.py @@ -10,9 +10,9 @@ class TestExaAISearch(BaseSearchTest): Tests for Exa AI Search functionality. """ - def get_custom_llm_provider(self) -> str: + def get_search_provider(self) -> str: """ - Return custom_llm_provider for Exa AI Search. + Return search_provider for Exa AI Search. """ return "exa_ai" diff --git a/tests/search_tests/test_parallel_ai_search.py b/tests/search_tests/test_parallel_ai_search.py index 6ecff2bdd1..1dc3b7c9d8 100644 --- a/tests/search_tests/test_parallel_ai_search.py +++ b/tests/search_tests/test_parallel_ai_search.py @@ -10,9 +10,9 @@ class TestParallelAISearch(BaseSearchTest): Tests for Parallel AI Search functionality. """ - def get_custom_llm_provider(self) -> str: + def get_search_provider(self) -> str: """ - Return custom_llm_provider for Parallel AI Search. + Return search_provider for Parallel AI Search. """ return "parallel_ai" diff --git a/tests/search_tests/test_perplexity_search.py b/tests/search_tests/test_perplexity_search.py index e68e054c76..3db2be59e8 100644 --- a/tests/search_tests/test_perplexity_search.py +++ b/tests/search_tests/test_perplexity_search.py @@ -17,9 +17,9 @@ class TestPerplexitySearch(BaseSearchTest): Tests for Perplexity Search functionality. """ - def get_custom_llm_provider(self) -> str: + def get_search_provider(self) -> str: """ - Return custom_llm_provider for Perplexity Search. + Return search_provider for Perplexity Search. """ return "perplexity" diff --git a/tests/search_tests/test_tavily_search.py b/tests/search_tests/test_tavily_search.py index 604ecf914e..c7e923d06b 100644 --- a/tests/search_tests/test_tavily_search.py +++ b/tests/search_tests/test_tavily_search.py @@ -17,9 +17,9 @@ class TestTavilySearch(BaseSearchTest): Tests for Tavily Search functionality. """ - def get_custom_llm_provider(self) -> str: + def get_search_provider(self) -> str: """ - Return custom_llm_provider for Tavily Search. + Return search_provider for Tavily Search. """ return "tavily"