diff --git a/litellm/llms/parallel_ai/search/cost_calculator.py b/litellm/llms/parallel_ai/search/cost_calculator.py new file mode 100644 index 00000000000..809cd280cc8 --- /dev/null +++ b/litellm/llms/parallel_ai/search/cost_calculator.py @@ -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 diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index ea21d1153fe..bde7b7b86db 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -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})) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 55618d9f772..a3cfb300ea6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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" }, diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index d96defbbcd6..ab5ef5853c9 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -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 diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 84461115e8e..21f27075e0f 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -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) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 55618d9f772..a3cfb300ea6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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" }, diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index dcd2e9edf7b..7dbac243d55 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -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 diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index 8a9ae4dae6d..62b4d003b45 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -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) diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py new file mode 100644 index 00000000000..72c69fc622c --- /dev/null +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py @@ -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"