mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
test(parallel_ai): fake search at HTTP boundary
This commit is contained in:
khushishelat
parent
d551886beb
commit
0e03bc8b8b
1 changed files with 85 additions and 96 deletions
|
|
@ -2,6 +2,7 @@
|
|||
Tests for Parallel AI Search API integration (v1 endpoint).
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -37,6 +38,17 @@ def _mock_response(payload=None):
|
|||
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):
|
||||
|
|
@ -152,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(
|
||||
|
|
@ -239,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
|
||||
|
|
@ -323,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):
|
||||
|
|
@ -357,62 +362,54 @@ class TestParallelAISearch:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self):
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_response()
|
||||
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",
|
||||
)
|
||||
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 = mock_post.call_args.kwargs.get("json")
|
||||
assert json_data["objective"] == "find peer-reviewed AI research"
|
||||
assert json_data["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"
|
||||
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
|
||||
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):
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_response()
|
||||
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",
|
||||
)
|
||||
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}]
|
||||
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."]
|
||||
first = response.results[0].model_dump()
|
||||
assert first["excerpts"] == ["First excerpt.", "Second excerpt."]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode,usage,max_results,expected_cost",
|
||||
|
|
@ -434,63 +431,55 @@ class TestParallelAISearch:
|
|||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_cost_uses_mode_and_provider_usage(self, mode, usage, max_results, expected_cost, bundled_cost_map):
|
||||
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}
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_response(response_payload)
|
||||
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,
|
||||
)
|
||||
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)
|
||||
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):
|
||||
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}],
|
||||
}
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_response(response_payload)
|
||||
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",
|
||||
)
|
||||
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)
|
||||
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):
|
||||
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"}
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
mock_post.return_value = _mock_response(response_payload)
|
||||
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}],
|
||||
)
|
||||
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 mock_post.call_args.kwargs["json"]
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
|
||||
assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue