mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
fix(search): forward search-tool params through the router, complete Parallel AI v1 param mapping (#37883)
* fix(search): forward search-tool params through the router, complete Parallel AI v1 param mapping
SearchAPIRouter dropped every parameter configured on a search tool, forwarding
only per-request kwargs. Any tool-level setting (mode, max_results, ...) was
silently lost on the way to the adapter, for every search provider.
Also completes the Parallel AI v1 search surface: after_date, fetch_policy,
location and include_domains now nest under advanced_settings instead of being
sent as unknown top-level fields, responses preserve search_id / session_id /
warnings / raw excerpts, and search cost is derived from the request mode and
the provider's reported usage rather than a single flat rate.
* fix(parallel_ai): stop a caller from pricing its own search request
`_parallel_ai_usage` carries the provider's reported usage into cost
calculation. It was only written when the response contained a usage block, so
a caller could pass `_parallel_ai_usage=[{"name": "sku_search", "count": 0}]`
and, whenever the provider omitted usage, bill $0.00 instead of $0.005 — the
value also reached the upstream request body as an unknown field.
The key is now stripped from inbound params and written unconditionally from
the parsed response, so only the provider can populate it.
* fix(parallel_ai): price fast search mode correctly
* test(parallel_ai): fake search at HTTP boundary
* fix(parallel_ai): tolerate null search result fields
---------
Co-authored-by: khushishelat <shelatkhushi@gmail.com>
This commit is contained in:
parent
92d453373a
commit
e2c3f51c46
9 changed files with 637 additions and 55 deletions
90
litellm/llms/parallel_ai/search/cost_calculator.py
Normal file
90
litellm/llms/parallel_ai/search/cost_calculator.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
PARALLEL_AI_DEFAULT_RESULTS: Final = 10
|
||||
PARALLEL_AI_ADDITIONAL_RESULT_COST: Final = 0.001
|
||||
PARALLEL_AI_USAGE_PARAM: Final = "_parallel_ai_usage"
|
||||
PARALLEL_AI_STANDARD_SEARCH_MODEL: Final = "parallel_ai/search"
|
||||
PARALLEL_AI_FAST_SEARCH_MODEL: Final = "parallel_ai/search-fast"
|
||||
PARALLEL_AI_TURBO_SEARCH_MODEL: Final = "parallel_ai/search-turbo"
|
||||
PARALLEL_AI_PRICING_MODEL_BY_MODE: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"fast": PARALLEL_AI_FAST_SEARCH_MODEL,
|
||||
"turbo": PARALLEL_AI_TURBO_SEARCH_MODEL,
|
||||
}
|
||||
)
|
||||
ADVANCED_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _non_negative_int(value: object) -> int | None:
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def _usage_count(usage: Sequence[Mapping[str, object]], sku: str) -> int | None:
|
||||
counts: Final = tuple(
|
||||
count
|
||||
for item in usage
|
||||
if item.get("name") == sku
|
||||
if (count := _non_negative_int(item.get("count"))) is not None
|
||||
)
|
||||
return sum(counts) if counts else None
|
||||
|
||||
|
||||
def _effective_mode(optional_params: Mapping[str, object]) -> str:
|
||||
mode: Final = optional_params.get("mode")
|
||||
if isinstance(mode, str):
|
||||
return mode
|
||||
|
||||
processor: Final = optional_params.get("processor")
|
||||
if processor == "pro":
|
||||
return "advanced"
|
||||
return "basic"
|
||||
|
||||
|
||||
def _effective_max_results(optional_params: Mapping[str, object]) -> int:
|
||||
try:
|
||||
advanced_settings: Final = ADVANCED_SETTINGS_ADAPTER.validate_python(optional_params.get("advanced_settings"))
|
||||
advanced_max_results: Final = _non_negative_int(advanced_settings.get("max_results"))
|
||||
if advanced_max_results is not None:
|
||||
return advanced_max_results
|
||||
except ValidationError:
|
||||
pass
|
||||
|
||||
max_results: Final = _non_negative_int(optional_params.get("max_results"))
|
||||
return max_results if max_results is not None else PARALLEL_AI_DEFAULT_RESULTS
|
||||
|
||||
|
||||
def _request_cost(mode: str) -> float:
|
||||
pricing_model: Final = PARALLEL_AI_PRICING_MODEL_BY_MODE.get(mode, PARALLEL_AI_STANDARD_SEARCH_MODEL)
|
||||
model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="parallel_ai")
|
||||
return float(model_info.get("input_cost_per_query") or 0.0)
|
||||
|
||||
|
||||
def _additional_results(
|
||||
optional_params: Mapping[str, object],
|
||||
usage: Sequence[Mapping[str, object]] | None,
|
||||
) -> int:
|
||||
usage_count: Final = _usage_count(usage, "sku_search_additional_results") if usage is not None else None
|
||||
if usage_count is not None:
|
||||
return usage_count
|
||||
if usage is not None:
|
||||
return 0
|
||||
return max(_effective_max_results(optional_params) - PARALLEL_AI_DEFAULT_RESULTS, 0)
|
||||
|
||||
|
||||
def parallel_ai_search_cost(
|
||||
optional_params: Mapping[str, object],
|
||||
usage: Sequence[Mapping[str, object]] | None,
|
||||
) -> float:
|
||||
request_cost: Final = _request_cost(_effective_mode(optional_params))
|
||||
request_count_from_usage: Final = _usage_count(usage, "sku_search") if usage is not None else None
|
||||
request_count: Final = request_count_from_usage if request_count_from_usage is not None else 1
|
||||
additional_results: Final = _additional_results(optional_params, usage)
|
||||
return request_count * request_cost + additional_results * PARALLEL_AI_ADDITIONAL_RESULT_COST
|
||||
|
|
@ -4,9 +4,13 @@ Calls Parallel AI's /v1/search endpoint to search the web.
|
|||
Parallel AI API Reference: https://docs.parallel.ai/api-reference/search/search
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypedDict
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
|
|
@ -14,9 +18,29 @@ from litellm.llms.base_llm.search.transformation import (
|
|||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from litellm.llms.parallel_ai.search.cost_calculator import PARALLEL_AI_USAGE_PARAM
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
|
||||
class _ParallelAIV1SearchResult(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
url: str | None = None
|
||||
title: str | None = None
|
||||
publish_date: str | None = None
|
||||
excerpts: Sequence[str] | None = None
|
||||
|
||||
|
||||
class _ParallelAIV1SearchResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
search_id: str | None = None
|
||||
session_id: str | None = None
|
||||
results: Sequence[_ParallelAIV1SearchResult] = ()
|
||||
usage: Sequence[Mapping[str, object]] | None = None
|
||||
warnings: Sequence[Mapping[str, object]] | None = None
|
||||
|
||||
|
||||
class _ParallelAISourcePolicy(TypedDict, total=False):
|
||||
include_domains: list[str]
|
||||
exclude_domains: list[str]
|
||||
|
|
@ -27,10 +51,16 @@ class _ParallelAIExcerptSettings(TypedDict, total=False):
|
|||
max_chars_per_result: int
|
||||
|
||||
|
||||
class _ParallelAIFetchPolicy(TypedDict, total=False):
|
||||
max_age_seconds: ReadOnly[int]
|
||||
timeout_seconds: ReadOnly[float]
|
||||
disable_cache_fallback: ReadOnly[bool]
|
||||
|
||||
|
||||
class _ParallelAIAdvancedSettings(TypedDict, total=False):
|
||||
source_policy: _ParallelAISourcePolicy
|
||||
excerpt_settings: _ParallelAIExcerptSettings
|
||||
fetch_policy: dict
|
||||
fetch_policy: _ParallelAIFetchPolicy
|
||||
location: str
|
||||
max_results: int
|
||||
|
||||
|
|
@ -43,14 +73,14 @@ class ParallelAISearchRequest(TypedDict, total=False):
|
|||
|
||||
search_queries: list[str] # Required - at least one keyword search query
|
||||
objective: str # Optional - natural-language description of search goal
|
||||
mode: str # Optional - 'turbo', 'basic', or 'advanced' (default 'advanced')
|
||||
mode: str # Optional - 'turbo', 'fast', 'basic', or 'advanced' (default 'advanced')
|
||||
max_chars_total: int # Optional - upper bound on total excerpt characters
|
||||
session_id: str # Optional - tracks calls across search/extract requests
|
||||
client_model: str # Optional - model consuming the results
|
||||
advanced_settings: _ParallelAIAdvancedSettings
|
||||
|
||||
|
||||
LEGACY_PROCESSOR_TO_MODE: Final = {"base": "basic", "pro": "advanced"}
|
||||
LEGACY_PROCESSOR_TO_MODE: Final = MappingProxyType({"base": "basic", "pro": "advanced"})
|
||||
|
||||
|
||||
class ParallelAISearchConfig(BaseSearchConfig):
|
||||
|
|
@ -67,16 +97,16 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
api_key = self.resolve_server_api_key(
|
||||
resolved_api_key: Final = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"),
|
||||
base_env_var="PARALLEL_AI_API_BASE",
|
||||
default_api_base=self.PARALLEL_AI_API_BASE,
|
||||
)
|
||||
if not api_key:
|
||||
if not resolved_api_key:
|
||||
raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.")
|
||||
headers["x-api-key"] = api_key
|
||||
headers["x-api-key"] = resolved_api_key
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
|
|
@ -87,13 +117,12 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
data: dict | list[dict] | None = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE
|
||||
resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE
|
||||
|
||||
api_base = api_base.rstrip("/")
|
||||
if not api_base.endswith("/v1/search"):
|
||||
api_base = f"{api_base.removesuffix('/v1')}/v1/search"
|
||||
|
||||
return api_base
|
||||
trimmed: Final = resolved_api_base.rstrip("/")
|
||||
if trimmed.endswith("/v1/search"):
|
||||
return trimmed
|
||||
return f"{trimmed.removesuffix('/v1')}/v1/search"
|
||||
|
||||
def transform_search_request(
|
||||
self,
|
||||
|
|
@ -109,14 +138,17 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
- If string: maps to `search_queries` (single item) and `objective`
|
||||
- If list: maps to `search_queries` (keyword queries)
|
||||
optional_params: Optional parameters for the request
|
||||
- mode: Search mode ('turbo', 'basic', 'advanced'); defaults to 'basic'
|
||||
- mode: Search mode ('turbo', 'fast', 'basic', 'advanced'); defaults to 'basic'
|
||||
- processor: Legacy v1beta param; 'base' maps to mode 'basic', 'pro' to 'advanced'
|
||||
- max_results: Maximum number of search results -> `advanced_settings.max_results`
|
||||
- search_domain_filter: Domains to include -> `advanced_settings.source_policy.include_domains`
|
||||
- search_domain_filter / include_domains: Domains to include -> `advanced_settings.source_policy.include_domains`
|
||||
- exclude_domains: Domains to exclude -> `advanced_settings.source_policy.exclude_domains`
|
||||
- country: ISO 3166-1 alpha-2 code -> `advanced_settings.location`
|
||||
- after_date: RFC 3339 date (YYYY-MM-DD) -> `advanced_settings.source_policy.after_date`
|
||||
- country / location: ISO 3166-1 alpha-2 code -> `advanced_settings.location`
|
||||
- max_chars_per_result: -> `advanced_settings.excerpt_settings.max_chars_per_result`
|
||||
- Any other params are passed through to the request body as-is
|
||||
- fetch_policy: Cache vs live-fetch policy -> `advanced_settings.fetch_policy`
|
||||
- Any other params (objective, max_chars_total, session_id, client_model, ...)
|
||||
are passed through to the request body as-is
|
||||
|
||||
Returns:
|
||||
Dict with request data following the v1 search request spec
|
||||
|
|
@ -137,7 +169,7 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
mode = LEGACY_PROCESSOR_TO_MODE.get(processor, processor)
|
||||
# the v1 API defaults to 'advanced' when mode is omitted; default to 'basic'
|
||||
# instead to keep v1beta's default tier (processor 'base') and litellm's
|
||||
# $0.004/query cost map entry for `parallel_ai/search` accurate
|
||||
# cost map entry for `parallel_ai/search` accurate
|
||||
request_data["mode"] = mode or "basic"
|
||||
|
||||
advanced_settings: Final[_ParallelAIAdvancedSettings] = {}
|
||||
|
|
@ -148,17 +180,29 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
if "country" in params:
|
||||
advanced_settings["location"] = params.pop("country")
|
||||
|
||||
if "location" in params:
|
||||
advanced_settings["location"] = params.pop("location")
|
||||
|
||||
if "max_chars_per_result" in params:
|
||||
advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")}
|
||||
|
||||
if "fetch_policy" in params:
|
||||
advanced_settings["fetch_policy"] = params.pop("fetch_policy")
|
||||
|
||||
source_policy: Final[_ParallelAISourcePolicy] = {}
|
||||
|
||||
if "search_domain_filter" in params:
|
||||
source_policy["include_domains"] = params.pop("search_domain_filter")
|
||||
|
||||
if "include_domains" in params:
|
||||
source_policy["include_domains"] = params.pop("include_domains")
|
||||
|
||||
if "exclude_domains" in params:
|
||||
source_policy["exclude_domains"] = params.pop("exclude_domains")
|
||||
|
||||
if "after_date" in params:
|
||||
source_policy["after_date"] = params.pop("after_date")
|
||||
|
||||
if source_policy:
|
||||
advanced_settings["source_policy"] = source_policy
|
||||
|
||||
|
|
@ -170,9 +214,11 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
# unified-spec param with no v1 equivalent
|
||||
params.pop("max_tokens_per_page", None)
|
||||
|
||||
result_data: Final[dict] = dict(request_data)
|
||||
result_data.update(params)
|
||||
return result_data
|
||||
# reserved for the provider's own reported usage, which prices the request;
|
||||
# a caller-supplied value would otherwise set its own cost
|
||||
params.pop(PARALLEL_AI_USAGE_PARAM, None)
|
||||
|
||||
return {**request_data, **params}
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
|
|
@ -186,26 +232,49 @@ class ParallelAISearchConfig(BaseSearchConfig):
|
|||
Parallel AI -> LiteLLM mappings:
|
||||
- results[].title -> SearchResult.title
|
||||
- results[].url -> SearchResult.url
|
||||
- results[].excerpts (array) -> SearchResult.snippet (joined string)
|
||||
- results[].excerpts (array) -> SearchResult.snippet (joined string); the raw
|
||||
array is preserved as an extra `excerpts` field on each result
|
||||
- results[].publish_date -> SearchResult.date
|
||||
- search_id / session_id / warnings are preserved as extra fields on the
|
||||
response; usage is preserved as `parallel_usage` (the `usage` name is
|
||||
reserved for LiteLLM's token-usage object)
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
parsed: Final = _ParallelAIV1SearchResponse.model_validate(raw_response.json())
|
||||
|
||||
results: Final = []
|
||||
for result in response_json.get("results", []):
|
||||
excerpts = result.get("excerpts") or []
|
||||
snippet = " ... ".join(excerpts) if excerpts else ""
|
||||
# written unconditionally: leaving a caller-supplied value in place when the
|
||||
# provider reports no usage would let the caller price its own request
|
||||
logging_obj.optional_params = {
|
||||
**logging_obj.optional_params,
|
||||
PARALLEL_AI_USAGE_PARAM: parsed.usage,
|
||||
}
|
||||
|
||||
search_result = SearchResult(
|
||||
title=result.get("title") or "",
|
||||
url=result.get("url") or "",
|
||||
snippet=snippet,
|
||||
date=result.get("publish_date"),
|
||||
last_updated=None,
|
||||
results: Final = tuple(
|
||||
SearchResult.model_validate(
|
||||
MappingProxyType(
|
||||
{
|
||||
"title": result.title or "",
|
||||
"url": result.url or "",
|
||||
"snippet": " ... ".join(result.excerpts or ()),
|
||||
"date": result.publish_date,
|
||||
"last_updated": None,
|
||||
"excerpts": result.excerpts or (),
|
||||
}
|
||||
)
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
object="search",
|
||||
for result in parsed.results
|
||||
)
|
||||
|
||||
extra_fields: Final = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("search_id", parsed.search_id),
|
||||
("session_id", parsed.session_id),
|
||||
("parallel_usage", parsed.usage),
|
||||
("warnings", parsed.warnings),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
|
||||
return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields}))
|
||||
|
|
|
|||
|
|
@ -38556,12 +38556,22 @@
|
|||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
|
||||
},
|
||||
"parallel_ai/search": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-fast": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-pro": {
|
||||
"input_cost_per_query": 0.009,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-turbo": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import random
|
|||
import traceback
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
|
@ -214,6 +215,15 @@ class SearchAPIRouter:
|
|||
api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials(
|
||||
tool_litellm_params=litellm_params,
|
||||
)
|
||||
protected_params: Final = frozenset(("search_provider", "api_key", "api_base"))
|
||||
search_params: Final = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for params in (litellm_params, kwargs)
|
||||
for key, value in params.items()
|
||||
if key not in protected_params and value is not None
|
||||
}
|
||||
)
|
||||
|
||||
verbose_router_logger.debug("Selected search tool with provider: %s", search_provider)
|
||||
|
||||
|
|
@ -222,7 +232,7 @@ class SearchAPIRouter:
|
|||
search_provider=search_provider,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
**kwargs,
|
||||
**search_params,
|
||||
)
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -2,16 +2,37 @@
|
|||
Cost calculation for search providers.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
|
||||
tuple[Mapping[str, object], ...]
|
||||
)
|
||||
EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _provider_usage(
|
||||
optional_params: Mapping[str, object] | None,
|
||||
usage_param: str,
|
||||
) -> tuple[Mapping[str, object], ...] | None:
|
||||
params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS
|
||||
raw_usage: Final[object] = params.get(usage_param)
|
||||
try:
|
||||
return PROVIDER_USAGE_ADAPTER.validate_python(raw_usage)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def search_provider_cost_per_query(
|
||||
model: str,
|
||||
custom_llm_provider: str | None = None,
|
||||
number_of_queries: int = 1,
|
||||
optional_params: dict | None = None,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculate cost for search-only providers.
|
||||
|
|
@ -28,6 +49,18 @@ def search_provider_cost_per_query(
|
|||
Returns:
|
||||
Tuple of (input_cost, output_cost) where output_cost is always 0.0
|
||||
"""
|
||||
if custom_llm_provider == "parallel_ai":
|
||||
from litellm.llms.parallel_ai.search.cost_calculator import (
|
||||
PARALLEL_AI_USAGE_PARAM,
|
||||
parallel_ai_search_cost,
|
||||
)
|
||||
|
||||
input_cost: Final = parallel_ai_search_cost(
|
||||
optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS,
|
||||
usage=_provider_usage(optional_params, PARALLEL_AI_USAGE_PARAM),
|
||||
)
|
||||
return (input_cost, 0.0)
|
||||
|
||||
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# Check for tiered pricing (e.g., Exa AI based on max_results)
|
||||
|
|
|
|||
|
|
@ -38556,12 +38556,22 @@
|
|||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
|
||||
},
|
||||
"parallel_ai/search": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-fast": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-pro": {
|
||||
"input_cost_per_query": 0.009,
|
||||
"input_cost_per_query": 0.005,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
"parallel_ai/search-turbo": {
|
||||
"input_cost_per_query": 0.001,
|
||||
"litellm_provider": "parallel_ai",
|
||||
"mode": "search"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -2294,6 +2294,7 @@ def search_tools():
|
|||
"search_provider": "perplexity",
|
||||
"api_key": "test-api-key",
|
||||
"api_base": "https://api.perplexity.ai",
|
||||
"mode": "turbo",
|
||||
},
|
||||
},
|
||||
{
|
||||
|
|
@ -2302,6 +2303,7 @@ def search_tools():
|
|||
"search_provider": "perplexity",
|
||||
"api_key": "test-api-key-2",
|
||||
"api_base": "https://api.perplexity.ai",
|
||||
"mode": "turbo",
|
||||
},
|
||||
},
|
||||
]
|
||||
|
|
@ -2393,6 +2395,7 @@ async def test_asearch_with_fallbacks_helper(search_tools):
|
|||
assert "search_provider" in kwargs
|
||||
assert kwargs["search_provider"] == "perplexity"
|
||||
assert "api_key" in kwargs
|
||||
assert kwargs["mode"] == "turbo"
|
||||
assert kwargs["query"] == "helper test query"
|
||||
return mock_response
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Tests for Parallel AI Search API integration (v1 endpoint).
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -30,13 +31,41 @@ MOCK_V1_RESPONSE = {
|
|||
}
|
||||
|
||||
|
||||
def _mock_response():
|
||||
def _mock_response(payload=None):
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = MOCK_V1_RESPONSE
|
||||
mock_response.json.return_value = payload if payload is not None else MOCK_V1_RESPONSE
|
||||
return mock_response
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def httpx_transport(monkeypatch):
|
||||
monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary.
|
||||
litellm,
|
||||
"disable_aiohttp_transport",
|
||||
True,
|
||||
)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
yield
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bundled_cost_map(monkeypatch):
|
||||
"""Price lookups against the bundled cost map.
|
||||
|
||||
litellm caches model-info lookups, so swapping ``model_cost`` only takes
|
||||
effect once those caches are invalidated -- on the way in and back out.
|
||||
"""
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
yield
|
||||
monkeypatch.undo()
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
class TestParallelAISearch:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _set_api_key(self, monkeypatch):
|
||||
|
|
@ -135,9 +164,7 @@ class TestParallelAISearch:
|
|||
json_data = mock_post.call_args.kwargs.get("json")
|
||||
assert json_data["mode"] == "basic"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"processor,expected_mode", [("base", "basic"), ("pro", "advanced")]
|
||||
)
|
||||
@pytest.mark.parametrize("processor,expected_mode", [("base", "basic"), ("pro", "advanced")])
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_processor_maps_to_mode(self, processor, expected_mode):
|
||||
with patch(
|
||||
|
|
@ -222,9 +249,7 @@ class TestParallelAISearch:
|
|||
"arxiv.org",
|
||||
"nature.com",
|
||||
]
|
||||
assert advanced_settings["source_policy"]["exclude_domains"] == [
|
||||
"reddit.com"
|
||||
]
|
||||
assert advanced_settings["source_policy"]["exclude_domains"] == ["reddit.com"]
|
||||
assert advanced_settings["excerpt_settings"]["max_chars_per_result"] == 1500
|
||||
|
||||
assert "max_results" not in json_data
|
||||
|
|
@ -306,10 +331,7 @@ class TestParallelAISearch:
|
|||
)
|
||||
|
||||
call_args = mock_post.call_args
|
||||
assert (
|
||||
call_args.kwargs["url"]
|
||||
== "https://proxy.internal.example.com/v1/search"
|
||||
)
|
||||
assert call_args.kwargs["url"] == "https://proxy.internal.example.com/v1/search"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_api_base_without_key_is_refused(self, monkeypatch):
|
||||
|
|
@ -338,3 +360,147 @@ class TestParallelAISearch:
|
|||
query="AI developments",
|
||||
search_provider="parallel_ai",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self, respx_mock, httpx_transport):
|
||||
route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE)
|
||||
|
||||
await litellm.asearch(
|
||||
query="AI developments",
|
||||
search_provider="parallel_ai",
|
||||
objective="find peer-reviewed AI research",
|
||||
include_domains=["arxiv.org"],
|
||||
after_date="2026-01-01",
|
||||
location="gb",
|
||||
fetch_policy={"max_age_seconds": 600, "disable_cache_fallback": True},
|
||||
client_model="claude-fable-5",
|
||||
)
|
||||
|
||||
json_data = json.loads(route.calls[0].request.content)
|
||||
assert json_data["objective"] == "find peer-reviewed AI research"
|
||||
assert json_data["client_model"] == "claude-fable-5"
|
||||
|
||||
advanced_settings = json_data["advanced_settings"]
|
||||
assert advanced_settings["location"] == "gb"
|
||||
assert advanced_settings["fetch_policy"] == {
|
||||
"max_age_seconds": 600,
|
||||
"disable_cache_fallback": True,
|
||||
}
|
||||
assert advanced_settings["source_policy"]["include_domains"] == ["arxiv.org"]
|
||||
assert advanced_settings["source_policy"]["after_date"] == "2026-01-01"
|
||||
|
||||
assert "include_domains" not in json_data
|
||||
assert "after_date" not in json_data
|
||||
assert "location" not in json_data
|
||||
assert "fetch_policy" not in json_data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_preserves_raw_parallel_fields(self, respx_mock, httpx_transport):
|
||||
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE)
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="AI developments",
|
||||
search_provider="parallel_ai",
|
||||
)
|
||||
|
||||
dumped = response.model_dump()
|
||||
assert dumped["search_id"] == "search_abc123"
|
||||
assert dumped["session_id"] == "session_xyz"
|
||||
assert dumped["parallel_usage"] == [{"name": "search_advanced", "count": 1}]
|
||||
|
||||
first = response.results[0].model_dump()
|
||||
assert first["excerpts"] == ["First excerpt.", "Second excerpt."]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_normalizes_null_result_fields(self, respx_mock, httpx_transport):
|
||||
response_payload = {
|
||||
**MOCK_V1_RESPONSE,
|
||||
"results": [{"url": None, "title": None, "publish_date": None, "excerpts": None}],
|
||||
}
|
||||
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="AI developments",
|
||||
search_provider="parallel_ai",
|
||||
)
|
||||
|
||||
assert len(response.results) == 1
|
||||
result = response.results[0]
|
||||
assert result.url == ""
|
||||
assert result.title == ""
|
||||
assert result.snippet == ""
|
||||
assert result.date is None
|
||||
assert result.model_dump()["excerpts"] == ()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode,usage,max_results,expected_cost",
|
||||
[
|
||||
("turbo", [{"name": "sku_search", "count": 1}], None, 0.001),
|
||||
("fast", [{"name": "sku_search", "count": 1}], None, 0.001),
|
||||
("basic", [{"name": "sku_search", "count": 1}], None, 0.005),
|
||||
("advanced", [{"name": "sku_search", "count": 1}], None, 0.005),
|
||||
(
|
||||
"basic",
|
||||
[
|
||||
{"name": "sku_search", "count": 1},
|
||||
{"name": "sku_search_additional_results", "count": 2},
|
||||
],
|
||||
20,
|
||||
0.007,
|
||||
),
|
||||
("basic", None, 20, 0.015),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_cost_uses_mode_and_provider_usage(
|
||||
self, mode, usage, max_results, expected_cost, bundled_cost_map, respx_mock, httpx_transport
|
||||
):
|
||||
response_payload = {**MOCK_V1_RESPONSE, "usage": usage}
|
||||
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="AI developments",
|
||||
search_provider="parallel_ai",
|
||||
mode=mode,
|
||||
max_results=max_results,
|
||||
)
|
||||
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(expected_cost)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_cost_treats_keyword_queries_as_one_request(
|
||||
self, bundled_cost_map, respx_mock, httpx_transport
|
||||
):
|
||||
response_payload = {
|
||||
**MOCK_V1_RESPONSE,
|
||||
"usage": [{"name": "sku_search", "count": 1}],
|
||||
}
|
||||
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
|
||||
|
||||
response = await litellm.asearch(
|
||||
query=["AI developments", "machine learning trends"],
|
||||
search_provider="parallel_ai",
|
||||
mode="basic",
|
||||
)
|
||||
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map, respx_mock, httpx_transport):
|
||||
"""`_parallel_ai_usage` prices the request, so a caller must not be able to set it.
|
||||
|
||||
The provider reports no usage here, which is the case where a caller-supplied
|
||||
value would otherwise survive into the cost calculation.
|
||||
"""
|
||||
response_payload = {k: v for k, v in MOCK_V1_RESPONSE.items() if k != "usage"}
|
||||
route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="AI developments",
|
||||
search_provider="parallel_ai",
|
||||
mode="basic",
|
||||
_parallel_ai_usage=[{"name": "sku_search", "count": 0}],
|
||||
)
|
||||
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
|
||||
assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,191 @@
|
|||
"""Gateway coverage for Parallel AI Search."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.integrations.websearch_interception.handler import (
|
||||
WebSearchInterceptionLogger,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
PARALLEL_SEARCH_URL: Final = "https://api.parallel.ai/v1/search"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> TestClient:
|
||||
return TestClient(proxy_server.app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def auth_as() -> Iterator[None]:
|
||||
async def _authorized_request() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed-sk-test",
|
||||
user_id="parallel-test-user",
|
||||
)
|
||||
|
||||
previous: Final = proxy_server.app.dependency_overrides.get(user_api_key_auth)
|
||||
proxy_server.app.dependency_overrides[user_api_key_auth] = _authorized_request
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if previous is None:
|
||||
proxy_server.app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
else:
|
||||
proxy_server.app.dependency_overrides[user_api_key_auth] = previous
|
||||
|
||||
|
||||
def _parallel_search_body() -> dict[str, object]:
|
||||
return {
|
||||
"search_id": "search_parallel_gateway",
|
||||
"results": [
|
||||
{
|
||||
"url": "https://example.com/parallel",
|
||||
"title": "Parallel result",
|
||||
"publish_date": "2026-08-13",
|
||||
"excerpts": ["First excerpt", "Second excerpt"],
|
||||
}
|
||||
],
|
||||
"usage": [{"name": "sku_search", "count": 1}],
|
||||
}
|
||||
|
||||
|
||||
def _parallel_router(mode: str = "turbo") -> Router:
|
||||
return Router(
|
||||
model_list=[],
|
||||
search_tools=[
|
||||
{
|
||||
"search_tool_name": "parallel-search",
|
||||
"litellm_params": {
|
||||
"search_provider": "parallel_ai",
|
||||
"api_key": "parallel-search-key",
|
||||
"mode": mode,
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
|
||||
def _mock_async_post(
|
||||
monkeypatch,
|
||||
*,
|
||||
url: str,
|
||||
response_body: dict[str, object],
|
||||
) -> AsyncMock:
|
||||
response = httpx.Response(
|
||||
status_code=200,
|
||||
json=response_body,
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
mock_post = AsyncMock(return_value=response)
|
||||
monkeypatch.setattr(AsyncHTTPHandler, "post", mock_post)
|
||||
return mock_post
|
||||
|
||||
|
||||
def test_parallel_search_gateway_route(client, auth_as, monkeypatch):
|
||||
"""The named search route selects its configured Parallel Search tool.
|
||||
|
||||
The tool-level `mode` must survive the router hop, so the upstream request
|
||||
is sent as `turbo` rather than falling back to the adapter default.
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _parallel_router())
|
||||
mock_post = _mock_async_post(
|
||||
monkeypatch,
|
||||
url=PARALLEL_SEARCH_URL,
|
||||
response_body=_parallel_search_body(),
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/v1/search/parallel-search",
|
||||
json={"query": "Parallel AI news", "max_results": 3},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["results"] == [
|
||||
{
|
||||
"title": "Parallel result",
|
||||
"url": "https://example.com/parallel",
|
||||
"snippet": "First excerpt ... Second excerpt",
|
||||
"date": "2026-08-13",
|
||||
"last_updated": None,
|
||||
"excerpts": ["First excerpt", "Second excerpt"],
|
||||
}
|
||||
]
|
||||
|
||||
request_kwargs = mock_post.await_args.kwargs
|
||||
assert request_kwargs["url"] == PARALLEL_SEARCH_URL
|
||||
assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key"
|
||||
assert request_kwargs["json"] == {
|
||||
"objective": "Parallel AI news",
|
||||
"search_queries": ["Parallel AI news"],
|
||||
"mode": "turbo",
|
||||
"advanced_settings": {"max_results": 3},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_web_search_interception_executes_parallel_search(monkeypatch):
|
||||
"""An intercepted web-search call uses the configured Parallel Search tool."""
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _parallel_router(mode="fast"))
|
||||
mock_post = _mock_async_post(
|
||||
monkeypatch,
|
||||
url=PARALLEL_SEARCH_URL,
|
||||
response_body=_parallel_search_body(),
|
||||
)
|
||||
logger = WebSearchInterceptionLogger(
|
||||
enabled_providers=[LlmProviders.OPENAI],
|
||||
search_tool_name="parallel-search",
|
||||
)
|
||||
|
||||
plan = await logger.async_build_responses_agentic_loop_plan(
|
||||
tools={
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "fc_parallel",
|
||||
"call_id": "fc_parallel",
|
||||
"type": "function_call",
|
||||
"name": "litellm_web_search",
|
||||
"arguments": '{"query":"Parallel AI news"}',
|
||||
"input": {"query": "Parallel AI news"},
|
||||
}
|
||||
]
|
||||
},
|
||||
model="gpt-5",
|
||||
messages=[{"role": "user", "content": "Research Parallel"}],
|
||||
response=None,
|
||||
optional_params={"tools": [{"type": "function", "name": "litellm_web_search"}]},
|
||||
logging_obj=None,
|
||||
stream=False,
|
||||
kwargs={"custom_llm_provider": "openai"},
|
||||
)
|
||||
|
||||
assert plan.run_agentic_loop is True
|
||||
assert plan.request_patch is not None
|
||||
assert plan.request_patch.messages[-1] == {
|
||||
"type": "function_call_output",
|
||||
"call_id": "fc_parallel",
|
||||
"output": (
|
||||
"Title: Parallel result\nURL: https://example.com/parallel\nSnippet: First excerpt ... Second excerpt"
|
||||
),
|
||||
}
|
||||
|
||||
request_kwargs = mock_post.await_args.kwargs
|
||||
assert request_kwargs["url"] == PARALLEL_SEARCH_URL
|
||||
assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key"
|
||||
assert request_kwargs["json"]["mode"] == "fast"
|
||||
Loading…
Add table
Reference in a new issue