mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge 01e10340d5 into 0c98afa780
This commit is contained in:
commit
2fdba7f7c6
7 changed files with 752 additions and 19 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
Loading…
Add table
Reference in a new issue