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:
James Liounis 2026-09-02 00:46:46 -04:00 committed by GitHub
parent 92d453373a
commit e2c3f51c46
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 637 additions and 55 deletions

View 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

View file

@ -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}))

View file

@ -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"
},

View file

@ -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

View file

@ -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)

View file

@ -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"
},

View file

@ -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

View file

@ -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)

View file

@ -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"