diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py index 1f78f8d154c..0b8ebe9792b 100644 --- a/litellm/llms/brave/search/transformation.py +++ b/litellm/llms/brave/search/transformation.py @@ -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 diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index 8ecd5a24c08..40cb7c56a7e 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -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, diff --git a/litellm/search/main.py b/litellm/search/main.py index b2dd51799a1..20fa6046553 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -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 diff --git a/tests/test_litellm/llms/brave/search/test_brave_search_transformation.py b/tests/test_litellm/llms/brave/search/test_brave_search_transformation.py new file mode 100644 index 00000000000..d21b297eb61 --- /dev/null +++ b/tests/test_litellm/llms/brave/search/test_brave_search_transformation.py @@ -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 \ No newline at end of file diff --git a/tests/test_litellm/llms/linkup/search/test_linkup_search_transformation.py b/tests/test_litellm/llms/linkup/search/test_linkup_search_transformation.py new file mode 100644 index 00000000000..b172c5f3deb --- /dev/null +++ b/tests/test_litellm/llms/linkup/search/test_linkup_search_transformation.py @@ -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" \ No newline at end of file diff --git a/tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py b/tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py index d6292c9cf3e..dc1e8ab5625 100644 --- a/tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py +++ b/tests/test_litellm/llms/nimble/search/test_nimble_search_transformation.py @@ -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}) diff --git a/tests/test_litellm/llms/tavily/search/test_tavily_search_transformation.py b/tests/test_litellm/llms/tavily/search/test_tavily_search_transformation.py new file mode 100644 index 00000000000..80ada94d22c --- /dev/null +++ b/tests/test_litellm/llms/tavily/search/test_tavily_search_transformation.py @@ -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"] \ No newline at end of file