This commit is contained in:
HenryHYeary 2026-09-12 14:54:46 -04:00 committed by GitHub
commit 2fdba7f7c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 752 additions and 19 deletions

View file

@ -175,6 +175,8 @@ class BraveSearchConfig(BaseSearchConfig):
- max_results count
- search_domain_filter q (append domain filters)
- country country
- start_date + end_date freshness (combined as "YYYY-MM-DDtoYYYY-MM-DD";
only sent when both are provided Brave's custom range requires both)
- max_tokens_per_page (not applicable, ignored)
All other Brave Search API-specific parameters are passed through as-is.
@ -190,36 +192,48 @@ class BraveSearchConfig(BaseSearchConfig):
# Brave Search API only supports single string queries
query = " ".join(query)
remaining = dict(optional_params)
request_data: Final[BraveSearchRequest] = {
"q": query,
}
# Only include "include_fetch_metadata" if it is not explicitly set to False
# This parameter results (more often than not) in a timestamp which we can use for last_updated
if "include_fetch_metadata" in optional_params and optional_params["include_fetch_metadata"] is False:
if remaining.pop("include_fetch_metadata", None) is False:
request_data["include_fetch_metadata"] = False
else:
request_data["include_fetch_metadata"] = True
# Transform unified spec parameters to Brave Search API format
if "max_results" in optional_params:
if "max_results" in remaining:
# Brave Search API supports 1-20 results per /web/search request
num_results: Final = min(optional_params["max_results"], 20)
num_results: Final = min(remaining.pop("max_results"), 20)
request_data["count"] = num_results
if "search_domain_filter" in optional_params:
# Convert to multiple "site:domain" clauses, joined by OR
domains: Final = optional_params["search_domain_filter"]
domains: Final = remaining.pop("search_domain_filter")
if isinstance(domains, list) and len(domains) > 0:
request_data["q"] = self._append_domain_filters(request_data["q"], domains)
start_date = remaining.pop("start_date", None)
end_date = remaining.pop("end_date", None)
if start_date and end_date:
request_data["freshness"] = f"{start_date}to{end_date}"
if "country" in remaining:
request_data["country"] = remaining.pop("country")
# Not applicable therefore popped and ignored
if "max_tokens_per_page" in remaining:
remaining.pop("max_tokens_per_page")
# Convert to dict before dynamic key assignments
result_data: Final = dict(request_data)
# Pass through all other parameters as-is
for param, value in optional_params.items():
if param not in self.get_supported_perplexity_optional_params() and param not in result_data:
result_data[param] = value
# Pass through anything the function did not explicitly consume
result_data.update(remaining)
# Store params in special key for URL building (Brave Search API uses GET not POST)
# Return a wrapper dict that stores params for get_complete_url to use

View file

@ -104,6 +104,8 @@ class LinkupSearchConfig(BaseSearchConfig):
- max_results -> maxResults
- search_domain_filter -> includeDomains
- country -> (not directly supported)
- start_date -> fromDate
- end_date -> toDate
- max_tokens_per_page -> (not applicable)
All other Linkup-specific parameters are passed through as-is.
@ -119,26 +121,36 @@ class LinkupSearchConfig(BaseSearchConfig):
# Linkup only supports single string queries, join with spaces
query = " ".join(query)
# Copy for passthrough data (Done this way to avoid having to change / add to Perplexity unified spec parameters)
remaining = dict(optional_params)
request_data: Final[LinkupSearchRequest] = {
"q": query,
"depth": optional_params.get("depth", "standard"),
"outputType": optional_params.get("outputType", "searchResults"),
"depth": remaining.pop("depth", "standard"),
"outputType": remaining.pop("outputType", "searchResults"),
}
# Transform Perplexity unified spec parameters to Linkup format
if "max_results" in optional_params:
request_data["maxResults"] = optional_params["max_results"]
if "max_results" in remaining:
request_data["maxResults"] = remaining.pop("max_results")
if "search_domain_filter" in optional_params:
request_data["includeDomains"] = optional_params["search_domain_filter"]
if "search_domain_filter" in remaining:
request_data["includeDomains"] = remaining.pop("search_domain_filter")
if "start_date" in remaining:
request_data["fromDate"] = remaining.pop("start_date")
if "end_date" in remaining:
request_data["toDate"] = remaining.pop("end_date")
if "max_tokens_per_page" in remaining:
remaining.pop("max_tokens_per_page")
# Convert to dict before dynamic key assignments
result_data: Final = dict(request_data)
# pass through all other parameters as-is
for param, value in optional_params.items():
if param not in self.get_supported_perplexity_optional_params() and param not in result_data:
result_data[param] = value
# Pass through any unhandled data
result_data.update(remaining)
return result_data
@ -189,7 +201,7 @@ class LinkupSearchConfig(BaseSearchConfig):
elif result_type == "image":
# For image results, use the URL as both title and snippet if name not provided
search_result = SearchResult(
title=result.get("name", result.get("url", "")),
title=result.get("name") or result.get("url", ""),
url=result.get("url", ""),
snippet=result.get("content", ""),
date=None,

View file

@ -29,6 +29,8 @@ def _build_search_optional_params(
search_domain_filter: list[str] | None = None,
max_tokens_per_page: int | None = None,
country: str | None = None,
start_date: str | None = None,
end_date: str | None = None,
) -> dict[str, Any]:
"""
Helper function to build optional_params dict from Perplexity Search API parameters.
@ -38,6 +40,8 @@ def _build_search_optional_params(
search_domain_filter: List of domains to filter (max 20)
max_tokens_per_page: Max tokens per page
country: Country code filter
start_date: Start date for results (YYYY-MM-DD)
end_date: End date for results (YYYY-MM-DD)
Returns:
Dict with non-None optional parameters
@ -52,6 +56,10 @@ def _build_search_optional_params(
optional_params["max_tokens_per_page"] = max_tokens_per_page
if country is not None:
optional_params["country"] = country
if start_date is not None:
optional_params["start_date"] = start_date
if end_date is not None:
optional_params["end_date"] = end_date
return optional_params
@ -64,6 +72,8 @@ async def asearch(
search_domain_filter: list[str] | None = None,
max_tokens_per_page: int | None = None,
country: str | None = None,
start_date: str | None = None,
end_date: str | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
@ -80,6 +90,8 @@ async def asearch(
search_domain_filter: Optional list of domains to filter (max 20)
max_tokens_per_page: Optional max tokens per page, default 1024
country: Optional country code filter (e.g., 'US', 'GB', 'DE')
start_date: Optional start date for results (YYYY-MM-DD)
end_date: Optional end date for results (YYYY-MM-DD)
api_key: Optional API key
api_base: Optional API base URL
timeout: Optional timeout
@ -128,6 +140,8 @@ async def asearch(
search_domain_filter=search_domain_filter,
max_tokens_per_page=max_tokens_per_page,
country=country,
start_date=start_date,
end_date=end_date,
api_key=api_key,
api_base=api_base,
timeout=timeout,
@ -167,6 +181,8 @@ def search(
search_domain_filter: list[str] | None = None,
max_tokens_per_page: int | None = None,
country: str | None = None,
start_date: str | None = None,
end_date: str | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
@ -183,6 +199,8 @@ def search(
search_domain_filter: Optional list of domains to filter (max 20)
max_tokens_per_page: Optional max tokens per page, default 1024
country: Optional country code filter (e.g., 'US', 'GB', 'DE')
start_date: Optional start date for results (YYYY-MM-DD)
end_date: Optional end date for results (YYYY-MM-DD)
api_key: Optional API key
api_base: Optional API base URL
timeout: Optional timeout
@ -255,6 +273,8 @@ def search(
search_domain_filter=search_domain_filter,
max_tokens_per_page=max_tokens_per_page,
country=country,
start_date=start_date,
end_date=end_date,
)
# Filter out internal LiteLLM parameters from kwargs

View file

@ -0,0 +1,279 @@
from unittest.mock import Mock
import pytest
from litellm.llms.brave.search.transformation import BraveSearchConfig, to_yyyy_mm_dd
def _config() -> BraveSearchConfig:
return BraveSearchConfig()
def _resp(payload: dict, params: dict | None = None) -> Mock:
"""Mock httpx.Response with the .json() and .request.url.params shape
Brave's transform_search_response actually reads."""
r = Mock()
r.json.return_value = payload
req = Mock()
req.url = Mock()
req.url.params = params or {}
r.request = req
return r
def _result(**overrides):
base = {
"title": "Test Title",
"url": "https://example.com",
"description": "Test description",
"page_age": None,
"age": None,
"fetched_content_timestamp": None,
}
return {**base, **overrides}
# --- ui_friendly_name / get_http_method ---
def test_ui_friendly_name():
assert _config().ui_friendly_name() == "Brave Search"
def test_get_http_method_is_get():
"""Brave's /web/search endpoint takes query params, not a JSON body."""
assert _config().get_http_method() == "GET"
# --- validate_environment ---
def test_validate_environment_with_explicit_key():
headers = _config().validate_environment({}, api_key="explicit-key")
assert headers["X-Subscription-Token"] == "explicit-key"
assert headers["Accept"] == "application/json"
assert headers["Accept-Encoding"] == "gzip"
assert headers["Content-Type"] == "application/json"
def test_validate_environment_reads_env_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("BRAVE_API_KEY", "env-key")
assert _config().validate_environment({})["X-Subscription-Token"] == "env-key"
def test_validate_environment_missing_key_raises(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("BRAVE_API_KEY", raising=False)
with pytest.raises(ValueError, match="BRAVE_API_KEY"):
_config().validate_environment({})
# --- get_complete_url ---
def test_get_complete_url_default_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("BRAVE_API_BASE", raising=False)
url = _config().get_complete_url(None, {}, data={"_brave_params": {"q": "test"}})
assert url.startswith("https://api.search.brave.com/res/v1/web/search?")
assert "q=test" in url
def test_get_complete_url_reads_env_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("BRAVE_API_BASE", "https://env-base.local/search")
url = _config().get_complete_url(None, {}, data={"_brave_params": {"q": "test"}})
assert url.startswith("https://env-base.local/search?")
def test_get_complete_url_without_brave_params_returns_bare_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("BRAVE_API_BASE", raising=False)
assert _config().get_complete_url(None, {}, data=None) == "https://api.search.brave.com/res/v1/web/search"
# --- transform_search_request ---
def test_transform_search_request_joins_list_query():
data = _config().transform_search_request(["foo", "bar"], {})
assert data["_brave_params"]["q"] == "foo bar"
def test_transform_search_request_max_results_clamped_to_20():
"""Unlike Nimble, Brave's /web/search hard-caps at 20 results per request."""
data = _config().transform_search_request("q", {"max_results": 500})
assert data["_brave_params"]["count"] == 20
def test_transform_search_request_max_results_under_cap_passes_through():
data = _config().transform_search_request("q", {"max_results": 5})
assert data["_brave_params"]["count"] == 5
def test_transform_search_request_appends_domain_filters_as_site_clauses():
data = _config().transform_search_request("q", {"search_domain_filter": ["arxiv.org", "nature.com"]})
assert data["_brave_params"]["q"] == "(q) AND (site:arxiv.org OR site:nature.com)"
def test_transform_search_request_empty_domain_filter_leaves_query_unchanged():
data = _config().transform_search_request("q", {"search_domain_filter": []})
assert data["_brave_params"]["q"] == "q"
def test_transform_search_request_drops_max_tokens_per_page():
"""No Brave equivalent — must not leak through as an unrecognized param."""
assert "max_tokens_per_page" not in _config().transform_search_request("q", {"max_tokens_per_page": 1024})[
"_brave_params"
]
def test_transform_search_request_country_is_forwarded():
"""Regression test: the docstring always claimed country -> country, but
there was previously no code path that actually implemented it it
relied on shared-list passthrough, which excluded it."""
data = _config().transform_search_request("q", {"country": "US"})
assert data["_brave_params"]["country"] == "US"
def test_transform_search_request_passes_through_unhandled_kwargs():
data = _config().transform_search_request("q", {"safesearch": "strict"})
assert data["_brave_params"]["safesearch"] == "strict"
@pytest.mark.parametrize(
"optional_params",
[
{}, # absent entirely
{"include_fetch_metadata": True}, # explicitly True
],
)
def test_transform_search_request_include_fetch_metadata_defaults_true(optional_params):
data = _config().transform_search_request("q", optional_params)
assert data["_brave_params"]["include_fetch_metadata"] is True
def test_transform_search_request_include_fetch_metadata_explicit_false_is_respected():
data = _config().transform_search_request("q", {"include_fetch_metadata": False})
assert data["_brave_params"]["include_fetch_metadata"] is False
# --- transform_search_request: date filtering ---
def test_transform_search_request_date_range_maps_to_freshness():
data = _config().transform_search_request("q", {"start_date": "2022-04-01", "end_date": "2022-07-30"})
assert data["_brave_params"]["freshness"] == "2022-04-01to2022-07-30"
def test_transform_search_request_only_start_date_is_dropped():
"""Brave's custom range needs both bounds; there's no correct open-ended equivalent."""
data = _config().transform_search_request("q", {"start_date": "2022-04-01"})
params = data["_brave_params"]
assert "freshness" not in params
assert "start_date" not in params # must not leak through raw either
def test_transform_search_request_only_end_date_is_dropped():
data = _config().transform_search_request("q", {"end_date": "2022-07-30"})
params = data["_brave_params"]
assert "freshness" not in params
assert "end_date" not in params
def test_transform_search_request_no_dates_omits_freshness():
data = _config().transform_search_request("q", {"max_results": 5})
assert "freshness" not in data["_brave_params"]
def test_transform_search_request_explicit_freshness_not_clobbered_by_absent_dates():
"""A caller-supplied native `freshness` preset (e.g. 'pw') should survive
when start_date/end_date aren't given at all."""
data = _config().transform_search_request("q", {"freshness": "pw"})
assert data["_brave_params"]["freshness"] == "pw"
# --- transform_search_response ---
def test_transform_search_response_extracts_basic_fields():
resp = _resp({"web": {"results": [_result()]}})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.title == "Test Title"
assert result.url == "https://example.com"
assert result.snippet == "Test description"
def test_transform_search_response_date_from_page_age():
resp = _resp({"web": {"results": [_result(page_age="2024-01-15")]}})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.date == "2024-01-15"
def test_transform_search_response_date_falls_back_to_age_when_page_age_absent():
resp = _resp({"web": {"results": [_result(page_age=None, age="2024-01-15")]}})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.date == "2024-01-15"
def test_transform_search_response_last_updated_from_fetched_content_timestamp():
resp = _resp({"web": {"results": [_result(fetched_content_timestamp="1705334400")]}})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.last_updated == "2024-01-15"
def test_transform_search_response_zero_hits():
resp = _resp({"web": {"results": []}})
assert _config().transform_search_response(resp, logging_obj=Mock()).results == []
def test_transform_search_response_result_filter_limits_sections():
resp = _resp(
{
"web": {"results": [_result(title="web result")]},
"news": {"results": [_result(title="news result")]},
},
params={"result_filter": "news"},
)
titles = [r.title for r in _config().transform_search_response(resp, logging_obj=Mock()).results]
assert titles == ["news result"]
def test_transform_search_response_no_result_filter_includes_all_sections():
resp = _resp(
{
"web": {"results": [_result(title="web result")]},
"news": {"results": [_result(title="news result")]},
},
)
titles = {r.title for r in _config().transform_search_response(resp, logging_obj=Mock()).results}
assert titles == {"web result", "news result"}
def test_transform_search_response_max_results_limits_across_combined_sections():
"""count doesn't limit Brave's own per-section results server-side, so
LiteLLM has to truncate across sections manually."""
resp = _resp(
{
"web": {"results": [_result(title=f"web-{i}") for i in range(5)]},
"news": {"results": [_result(title=f"news-{i}") for i in range(5)]},
},
params={"count": "3"},
)
results = _config().transform_search_response(resp, logging_obj=Mock()).results
assert len(results) == 3
# --- to_yyyy_mm_dd helper ---
@pytest.mark.parametrize(
"value,expected",
[
(None, None),
("", None),
("not-a-date", None),
("2024-01-15", "2024-01-15"),
("2024/01/15", "2024-01-15"),
(1705334400, "2024-01-15"), # unix seconds
(1705334400000, "2024-01-15"), # unix milliseconds
],
)
def test_to_yyyy_mm_dd(value, expected):
assert to_yyyy_mm_dd(value) == expected

View file

@ -0,0 +1,220 @@
from unittest.mock import Mock
import pytest
from litellm.llms.linkup.search.transformation import LinkupSearchConfig
def _config() -> LinkupSearchConfig:
return LinkupSearchConfig()
def _resp(payload: dict) -> Mock:
r = Mock()
r.json.return_value = payload
return r
def _result(**overrides):
base = {
"type": "text",
"name": "Test Title",
"url": "https://example.com",
"content": "Test content",
}
return {**base, **overrides}
# --- ui_friendly_name ---
def test_ui_friendly_name():
assert _config().ui_friendly_name() == "Linkup"
# --- validate_environment ---
def test_validate_environment_with_explicit_key():
headers = _config().validate_environment({}, api_key="explicit-key")
assert headers["Authorization"] == "Bearer explicit-key"
assert headers["Content-Type"] == "application/json"
def test_validate_environment_reads_env_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("LINKUP_API_KEY", "env-key")
assert _config().validate_environment({})["Authorization"] == "Bearer env-key"
def test_validate_environment_missing_key_raises(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("LINKUP_API_KEY", raising=False)
with pytest.raises(ValueError, match="LINKUP_API_KEY"):
_config().validate_environment({})
# --- get_complete_url ---
def test_get_complete_url_default_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("LINKUP_API_BASE", raising=False)
assert _config().get_complete_url(None, {}) == "https://api.linkup.so/v1/search"
def test_get_complete_url_reads_env_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("LINKUP_API_BASE", "https://env-base.local/v1")
assert _config().get_complete_url(None, {}) == "https://env-base.local/v1/search"
def test_get_complete_url_does_not_duplicate_search_suffix():
assert _config().get_complete_url("https://self-hosted.local/v1/search", {}) == "https://self-hosted.local/v1/search"
def test_get_complete_url_appends_search_when_missing():
assert _config().get_complete_url("https://self-hosted.local/v1", {}) == "https://self-hosted.local/v1/search"
# --- transform_search_request: query / defaults ---
def test_transform_search_request_joins_list_query():
assert _config().transform_search_request(["foo", "bar"], {})["q"] == "foo bar"
def test_transform_search_request_string_query_unchanged():
assert _config().transform_search_request("foo bar", {})["q"] == "foo bar"
def test_transform_search_request_defaults_depth_and_output_type():
data = _config().transform_search_request("q", {})
assert data["depth"] == "standard"
assert data["outputType"] == "searchResults"
def test_transform_search_request_respects_explicit_depth_and_output_type():
data = _config().transform_search_request("q", {"depth": "deep", "outputType": "sourcedAnswer"})
assert data["depth"] == "deep"
assert data["outputType"] == "sourcedAnswer"
# --- transform_search_request: unified param mapping ---
def test_transform_search_request_max_results_maps_to_maxResults():
data = _config().transform_search_request("q", {"max_results": 10})
assert data["maxResults"] == 10
assert "max_results" not in data
def test_transform_search_request_search_domain_filter_maps_to_includeDomains():
data = _config().transform_search_request("q", {"search_domain_filter": ["arxiv.org", "nature.com"]})
assert data["includeDomains"] == ["arxiv.org", "nature.com"]
assert "search_domain_filter" not in data
def test_transform_search_request_drops_max_tokens_per_page():
"""No Linkup equivalent — must not leak through as an unrecognized param."""
assert "max_tokens_per_page" not in _config().transform_search_request("q", {"max_tokens_per_page": 1024})
def test_transform_search_request_country_has_no_native_equivalent():
"""Linkup has no native country filter; per the docstring this is intentionally unmapped —
it should pass through raw rather than silently vanish, since Linkup's real API will just
ignore an unrecognized field."""
data = _config().transform_search_request("q", {"country": "US"})
assert data.get("country") == "US"
# --- transform_search_request: date filtering ---
def test_transform_search_request_maps_start_date_to_from_date():
data = _config().transform_search_request("q", {"start_date": "1999-03-20"})
assert data["fromDate"] == "1999-03-20"
assert "start_date" not in data
def test_transform_search_request_maps_end_date_to_to_date():
data = _config().transform_search_request("q", {"end_date": "1999-04-20"})
assert data["toDate"] == "1999-04-20"
assert "end_date" not in data
def test_transform_search_request_date_range_together():
data = _config().transform_search_request(
"q", {"start_date": "1999-03-20", "end_date": "1999-04-20"}
)
assert data["fromDate"] == "1999-03-20"
assert data["toDate"] == "1999-04-20"
def test_transform_search_request_without_dates_omits_both():
data = _config().transform_search_request("q", {"max_results": 5})
assert "fromDate" not in data
assert "toDate" not in data
# --- transform_search_request: passthrough ---
def test_transform_search_request_passes_through_unhandled_kwargs():
"""A param this function doesn't explicitly handle should still reach Linkup unchanged."""
data = _config().transform_search_request("q", {"includeImages": True})
assert data["includeImages"] is True
def test_transform_search_request_does_not_duplicate_consumed_params():
"""Regression test: a translated param (start_date -> fromDate) must not
also leak through raw via the passthrough path."""
data = _config().transform_search_request(
"q", {"start_date": "1999-03-20", "max_results": 5, "includeImages": True}
)
assert "start_date" not in data
assert "max_results" not in data
assert data["fromDate"] == "1999-03-20"
assert data["maxResults"] == 5
assert data["includeImages"] is True
# --- transform_search_response ---
def test_transform_search_response_text_result_fields():
resp = _resp({"results": [_result()]})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.title == "Test Title"
assert result.url == "https://example.com"
assert result.snippet == "Test content"
assert result.date is None
assert result.last_updated is None
def test_transform_search_response_image_result_falls_back_to_url_for_title():
resp = _resp({"results": [_result(type="image", name="", url="https://example.com/img.png")]})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.title == "https://example.com/img.png"
def test_transform_search_response_image_result_uses_name_when_present():
resp = _resp({"results": [_result(type="image", name="Alt text")]})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.title == "Alt text"
def test_transform_search_response_preserves_order():
resp = _resp({"results": [_result(name=n) for n in ("first", "second", "third")]})
titles = [r.title for r in _config().transform_search_response(resp, logging_obj=Mock()).results]
assert titles == ["first", "second", "third"]
def test_transform_search_response_zero_hits():
resp = _resp({"results": []})
assert _config().transform_search_response(resp, logging_obj=Mock()).results == []
def test_transform_search_response_unknown_result_type_is_skipped():
"""A result type that's neither "text" nor "image" (e.g. a future Linkup
addition) shouldn't crash the call; it's just silently omitted."""
resp = _resp({"results": [_result(type="video"), _result()]})
results = _config().transform_search_response(resp, logging_obj=Mock()).results
assert len(results) == 1
assert results[0].title == "Test Title"

View file

@ -100,6 +100,14 @@ def test_transform_search_request_max_results_is_not_clamped():
def test_transform_search_request_uppercases_country():
assert _config().transform_search_request("q", {"country": "us"})["country"] == "US"
def test_transform_search_request_date_range_passes_through_native_names():
"""Nimble already uses the unified spec's own field names for date filtering."""
data = _config().transform_search_request(
"q", {"start_date": "1999-03-20", "end_date": "1999-04-20"}
)
assert data["start_date"] == "1999-03-20"
assert data["end_date"] == "1999-04-20"
def test_transform_search_request_drops_max_tokens_per_page():
assert "max_tokens_per_page" not in _config().transform_search_request("q", {"max_tokens_per_page": 1024})

View file

@ -0,0 +1,180 @@
from unittest.mock import Mock
import pytest
from litellm.llms.tavily.search.transformation import TavilySearchConfig
def _config() -> TavilySearchConfig:
return TavilySearchConfig()
def _resp(payload: dict) -> Mock:
r = Mock()
r.json.return_value = payload
return r
def _result(**overrides):
base = {
"title": "Test Title",
"url": "https://example.com",
"content": "Test content",
}
return {**base, **overrides}
# --- ui_friendly_name ---
def test_ui_friendly_name():
assert _config().ui_friendly_name() == "Tavily"
# --- validate_environment ---
def test_validate_environment_with_explicit_key():
headers = _config().validate_environment({}, api_key="explicit-key")
assert headers["Authorization"] == "Bearer explicit-key"
assert headers["Content-Type"] == "application/json"
def test_validate_environment_reads_env_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("TAVILY_API_KEY", "env-key")
assert _config().validate_environment({})["Authorization"] == "Bearer env-key"
def test_validate_environment_missing_key_raises(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("TAVILY_API_KEY", raising=False)
with pytest.raises(ValueError, match="TAVILY_API_KEY"):
_config().validate_environment({})
# --- get_complete_url ---
def test_get_complete_url_default_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("TAVILY_API_BASE", raising=False)
assert _config().get_complete_url(None, {}) == "https://api.tavily.com/search"
def test_get_complete_url_reads_env_base(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("TAVILY_API_BASE", "https://env-base.local")
assert _config().get_complete_url(None, {}) == "https://env-base.local/search"
def test_get_complete_url_does_not_duplicate_search_suffix():
assert _config().get_complete_url("https://self-hosted.local/search", {}) == "https://self-hosted.local/search"
# --- transform_search_request: query ---
def test_transform_search_request_joins_list_query():
assert _config().transform_search_request(["foo", "bar"], {})["query"] == "foo bar"
def test_transform_search_request_string_query_unchanged():
assert _config().transform_search_request("foo bar", {})["query"] == "foo bar"
# --- transform_search_request: unified param mapping ---
def test_transform_search_request_max_results_passes_through():
data = _config().transform_search_request("q", {"max_results": 10})
assert data["max_results"] == 10
def test_transform_search_request_search_domain_filter_maps_to_include_domains():
data = _config().transform_search_request("q", {"search_domain_filter": ["arxiv.org", "nature.com"]})
assert data["include_domains"] == ["arxiv.org", "nature.com"]
assert "search_domain_filter" not in data
def test_transform_search_request_lowercases_country():
data = _config().transform_search_request("q", {"country": "US"})
assert data["country"] == "us"
def test_transform_search_request_drops_max_tokens_per_page():
"""No Tavily equivalent — must not leak through as an unrecognized param."""
assert "max_tokens_per_page" not in _config().transform_search_request("q", {"max_tokens_per_page": 1024})
# --- transform_search_request: date filtering (native passthrough) ---
def test_transform_search_request_date_range_passes_through_native_names():
"""Tavily already uses the unified spec's own field names for date filtering."""
data = _config().transform_search_request(
"q", {"start_date": "1999-03-20", "end_date": "1999-04-20"}
)
assert data["start_date"] == "1999-03-20"
assert data["end_date"] == "1999-04-20"
def test_transform_search_request_single_date_passes_through():
"""Unlike Brave, Tavily has no combined-range requirement — a lone bound is valid."""
data = _config().transform_search_request("q", {"start_date": "1999-03-20"})
assert data["start_date"] == "1999-03-20"
assert "end_date" not in data
def test_transform_search_request_without_dates_omits_both():
data = _config().transform_search_request("q", {"max_results": 5})
assert "start_date" not in data
assert "end_date" not in data
# --- transform_search_request: passthrough ---
def test_transform_search_request_passes_through_unhandled_kwargs():
"""Native Tavily-specific params (topic, search_depth, time_range, etc.)
aren't explicitly mapped, but should still reach the request unchanged."""
data = _config().transform_search_request("q", {"topic": "news", "search_depth": "advanced"})
assert data["topic"] == "news"
assert data["search_depth"] == "advanced"
def test_transform_search_request_does_not_duplicate_consumed_params():
"""Regression test: a translated param (search_domain_filter -> include_domains)
must not also leak through raw via the passthrough path."""
data = _config().transform_search_request(
"q", {"search_domain_filter": ["arxiv.org"], "max_results": 5}
)
assert "search_domain_filter" not in data
assert data["include_domains"] == ["arxiv.org"]
assert data["max_results"] == 5
# --- transform_search_response ---
def test_transform_search_response_extracts_basic_fields():
resp = _resp({"results": [_result()]})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.title == "Test Title"
assert result.url == "https://example.com"
assert result.snippet == "Test content"
def test_transform_search_response_date_and_last_updated_are_none():
"""Tavily's response has no date/last_updated equivalent."""
resp = _resp({"results": [_result()]})
result = _config().transform_search_response(resp, logging_obj=Mock()).results[0]
assert result.date is None
assert result.last_updated is None
def test_transform_search_response_zero_hits():
resp = _resp({"results": []})
assert _config().transform_search_response(resp, logging_obj=Mock()).results == []
def test_transform_search_response_preserves_order():
resp = _resp({"results": [_result(title=t) for t in ("first", "second", "third")]})
titles = [r.title for r in _config().transform_search_response(resp, logging_obj=Mock()).results]
assert titles == ["first", "second", "third"]