chore: initial snapshot for gitea/github upload
This commit is contained in:
Your Name
@@ -0,0 +1,6 @@
|
||||
"""
|
||||
Parallel AI Search API module.
|
||||
"""
|
||||
from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig
|
||||
|
||||
__all__ = ["ParallelAISearchConfig"]
|
||||
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
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"
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Parallel 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("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,
|
||||
data: Optional[Union[Dict, List[Dict]]] = None,
|
||||
**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 list both become objective)
|
||||
if isinstance(query, list):
|
||||
request_data["objective"] = self._transform_query_to_objective(query)
|
||||
else:
|
||||
request_data["objective"] = query
|
||||
|
||||
# Transform Perplexity unified spec parameters to Parallel AI format
|
||||
if "max_results" in optional_params:
|
||||
request_data["max_results"] = optional_params["max_results"]
|
||||
|
||||
# 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
|
||||
|
||||
# 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 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",
|
||||
)
|
||||
Reference in New Issue
Block a user