From 208f76f8ade735dde06e5038f97f5771291b04b4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 21 Oct 2025 17:00:05 -0700 Subject: [PATCH] [Feat] Add Parallel AI - Search API (#15772) * 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 * add TAVILY to LlmProviders * add TavilySearchConfig * add TavilySearchConfig * TestTavilySearch * add tavily transform * TestParallelAISearch * add LlmProviders * add ParallelAISearchConfig * add ParallelAISearchConfig * ParallelAISearchConfig --- litellm/llms/parallel_ai/search/__init__.py | 7 + .../llms/parallel_ai/search/transformation.py | 198 ++++++++++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 4 + tests/search_tests/test_parallel_ai_search.py | 19 ++ 5 files changed, 229 insertions(+) create mode 100644 litellm/llms/parallel_ai/search/__init__.py create mode 100644 litellm/llms/parallel_ai/search/transformation.py create mode 100644 tests/search_tests/test_parallel_ai_search.py diff --git a/litellm/llms/parallel_ai/search/__init__.py b/litellm/llms/parallel_ai/search/__init__.py new file mode 100644 index 0000000000..cc2ff91ea3 --- /dev/null +++ b/litellm/llms/parallel_ai/search/__init__.py @@ -0,0 +1,7 @@ +""" +Parallel AI Search API module. +""" +from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig + +__all__ = ["ParallelAISearchConfig"] + diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py new file mode 100644 index 0000000000..2f465c7bcb --- /dev/null +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -0,0 +1,198 @@ +""" +Calls Parallel AI's /search endpoint to search the web. + +Parallel AI API Reference: https://docs.parallel.ai/api-reference/search-and-extract-api-beta/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 _ParallelAISourcePolicy(TypedDict, total=False): + """Source policy for Parallel AI search results.""" + allowed_domains: List[str] # Optional - list of allowed domains + disallowed_domains: List[str] # Optional - list of disallowed domains + + +class _ParallelAISearchRequestRequired(TypedDict): + """Required fields for Parallel AI Search API request.""" + # Note: At least one of objective or search_queries must be provided + pass + + +class ParallelAISearchRequest(_ParallelAISearchRequestRequired, total=False): + """ + Parallel AI Search API request format. + Based on: https://docs.parallel.ai/api-reference/search-and-extract-api-beta/search + """ + objective: str # Optional - natural-language description of search goal + search_queries: List[str] # Optional - list of keyword search queries + processor: str # Optional - search processor ('base', 'pro'), default 'base' + max_results: int # Optional - maximum number of results, default 10 + max_chars_per_result: int # Optional - max characters per result excerpt + source_policy: _ParallelAISourcePolicy # Optional - source policy for allowed/disallowed domains + + +class ParallelAISearchConfig(BaseSearchConfig): + PARALLEL_AI_API_BASE = "https://api.parallel.ai" + PARALLEL_HEADER_SEARCH_EXTRACT_VALUE = "search-extract-2025-10-10" + + 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("PARALLEL_AI_API_KEY") or get_secret_str("PARALLEL_API_KEY") + if not api_key: + raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") + headers["x-api-key"] = api_key + headers["Content-Type"] = "application/json" + headers["parallel-beta"] = self.PARALLEL_HEADER_SEARCH_EXTRACT_VALUE + 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("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE + + # Parallel AI search endpoint is at /v1beta/search + if not api_base.endswith("/v1beta/search"): + if api_base.endswith("/"): + api_base = f"{api_base}v1beta/search" + else: + api_base = f"{api_base}/v1beta/search" + + return api_base + + def _transform_query_to_objective(self, query: Union[str, List[str]]) -> str: + """ + Transform query to objective. + """ + if isinstance(query, list): + return " ".join(query) + return query + + + def transform_search_request( + self, + query: Union[str, List[str]], + optional_params: dict, + **kwargs, + ) -> Dict: + """ + Transform Search request to Parallel AI API format. + + Args: + query: Search query (string or list of strings) + - If string: maps to `objective` (natural language) + - If list: maps to `search_queries` (keyword queries) + optional_params: Optional parameters for the request + - max_results: Maximum number of search results (default 10) + - search_domain_filter: List of domains to include -> maps to `source_policy.allowed_domains` + - exclude_domains: List of domains to exclude -> maps to `source_policy.disallowed_domains` + - processor: Search processor ('base', 'pro') + - max_chars_per_result: Max characters per result excerpt + + Returns: + Dict with typed request data following ParallelAISearchRequest spec + """ + request_data: ParallelAISearchRequest = {} + + # Map query to objective (string) or search_queries (list) + if isinstance(query, list): + # List of queries -> search_queries + request_data["objective"] = self._transform_query_to_objective(query) + else: + # Single string -> objective (natural language description) + request_data["objective"] = query + + # Map max_results (same field name) + if "max_results" in optional_params: + request_data["max_results"] = optional_params["max_results"] + + # Map processor (same field name) + if "processor" in optional_params: + request_data["processor"] = optional_params["processor"] + + # Map max_chars_per_result (same field name) + if "max_chars_per_result" in optional_params: + request_data["max_chars_per_result"] = optional_params["max_chars_per_result"] + + # Map domain filters to source_policy + source_policy: _ParallelAISourcePolicy = {} + + if "search_domain_filter" in optional_params: + source_policy["allowed_domains"] = optional_params["search_domain_filter"] + + if "exclude_domains" in optional_params: + source_policy["disallowed_domains"] = optional_params["exclude_domains"] + + if source_policy: + request_data["source_policy"] = source_policy + + return dict(request_data) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs, + ) -> SearchResponse: + """ + Transform Parallel AI API response to LiteLLM unified SearchResponse format. + + Parallel AI → LiteLLM mappings: + - results[].title → SearchResult.title + - results[].url → SearchResult.url + - results[].excerpts (array) → SearchResult.snippet (joined string) + - No date/last_updated fields in Parallel AI response (set to None) + + Args: + raw_response: Raw httpx response from Parallel AI 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", []): + # Join excerpts array into a single snippet string + excerpts = result.get("excerpts", []) + snippet = " ... ".join(excerpts) if excerpts else "" + + search_result = SearchResult( + title=result.get("title", ""), + url=result.get("url", ""), + snippet=snippet, + date=None, # Parallel AI doesn't provide date in response + last_updated=None, # Parallel AI doesn't provide last_updated in response + ) + results.append(search_result) + + return SearchResponse( + results=results, + object="search", + ) + diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1d6a5b547c..34b008bd5a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2514,6 +2514,7 @@ class LlmProviders(str, Enum): HUMANLOOP = "humanloop" TOPAZ = "topaz" TAVILY = "tavily" + PARALLEL_AI = "parallel_ai" ASSEMBLYAI = "assemblyai" GITHUB_COPILOT = "github_copilot" SNOWFLAKE = "snowflake" diff --git a/litellm/utils.py b/litellm/utils.py index 31662b460c..c720a713c5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7627,6 +7627,9 @@ class ProviderConfigManager: """ Get Search configuration for a given provider. """ + from litellm.llms.parallel_ai.search.transformation import ( + ParallelAISearchConfig, + ) from litellm.llms.perplexity.search.transformation import ( PerplexitySearchConfig, ) @@ -7637,6 +7640,7 @@ class ProviderConfigManager: PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.PERPLEXITY: PerplexitySearchConfig, litellm.LlmProviders.TAVILY: TavilySearchConfig, + litellm.LlmProviders.PARALLEL_AI: ParallelAISearchConfig, } config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/tests/search_tests/test_parallel_ai_search.py b/tests/search_tests/test_parallel_ai_search.py new file mode 100644 index 0000000000..6ecff2bdd1 --- /dev/null +++ b/tests/search_tests/test_parallel_ai_search.py @@ -0,0 +1,19 @@ +import pytest +import litellm +from typing import List, Union + +from tests.search_tests.base_search_unit_tests import BaseSearchTest + + +class TestParallelAISearch(BaseSearchTest): + """ + Tests for Parallel AI Search functionality. + """ + + def get_custom_llm_provider(self) -> str: + """ + Return custom_llm_provider for Parallel AI Search. + """ + return "parallel_ai" + +