""" Tests for TinyFish Search API integration. """ from unittest.mock import MagicMock, patch import httpx import pytest from litellm.llms.tinyfish.search.transformation import ( TinyfishSearchConfig, _append_domain_filters, _default_missing_result_fields, ) MOCK_TINYFISH_RESPONSE = { "query": "web automation tools", "results": [ { "position": 1, "site_name": "tinyfish.ai", "title": "TinyFish - AI Web Automation", "snippet": "Automate any website with natural language.", "url": "https://tinyfish.ai", }, { "position": 2, "site_name": "github.com", "title": "Top Web Automation Tools", "snippet": "A curated list of browser automation frameworks.", "url": "https://github.com/example/web-automation", }, ], "total_results": 2, "page": 0, } def _make_mock_response( json_data: dict | None = None, status_code: int = 200, request_url: str | None = None, text: str | None = None, headers: dict | None = None, ) -> MagicMock: import json as _json mock = MagicMock() mock.status_code = status_code mock.headers = headers or {} if json_data is not None: mock.json.return_value = json_data mock.text = text if text is not None else _json.dumps(json_data) else: # Force .json() to raise as httpx.Response does for non-JSON bodies. mock.json.side_effect = _json.JSONDecodeError("Expecting value", text or "", 0) mock.text = text or "" if request_url: mock.request = MagicMock() mock.request.url = httpx.URL(request_url) else: mock.request = None return mock class TestTinyfishSearchConfig: def test_ui_friendly_name(self): assert TinyfishSearchConfig.ui_friendly_name() == "TinyFish" def test_get_http_method(self): assert TinyfishSearchConfig().get_http_method() == "GET" def test_validate_environment_with_explicit_key(self): config = TinyfishSearchConfig() headers = config.validate_environment(headers={}, api_key="sk-tinyfish-test") assert headers["X-API-Key"] == "sk-tinyfish-test" assert headers["Accept"] == "application/json" def test_validate_environment_from_env(self, monkeypatch): monkeypatch.setenv("TINYFISH_API_KEY", "sk-from-env") config = TinyfishSearchConfig() headers = config.validate_environment(headers={}) assert headers["X-API-Key"] == "sk-from-env" def test_validate_environment_missing_key(self, monkeypatch): monkeypatch.delenv("TINYFISH_API_KEY", raising=False) config = TinyfishSearchConfig() with pytest.raises(ValueError, match="TINYFISH_API_KEY"): config.validate_environment(headers={}) def test_validate_environment_uses_api_base_kwarg(self): config = TinyfishSearchConfig() headers = config.validate_environment( headers={}, api_key="sk-test", api_base="https://custom.tinyfish.ai", ) assert headers["X-API-Key"] == "sk-test" class TestTransformSearchRequest: def test_basic_query(self): config = TinyfishSearchConfig() result = config.transform_search_request( query="hello world", optional_params={} ) assert result == {"_tinyfish_params": {"query": "hello world"}} def test_list_query_joined(self): config = TinyfishSearchConfig() result = config.transform_search_request( query=["hello", "world"], optional_params={} ) assert result["_tinyfish_params"]["query"] == "hello world" def test_country_maps_to_location(self): config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"country": "US"} ) assert result["_tinyfish_params"]["location"] == "US" def test_max_results_not_sent_on_wire(self): # TinyFish doesn't honor max_results server-side; we apply it client-side # in transform_search_response. The querystring should be free of it. config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"max_results": 5} ) assert "max_results" not in result["_tinyfish_params"] def test_max_results_clamped_upper_stored_on_self(self): config = TinyfishSearchConfig() config.transform_search_request( query="test", optional_params={"max_results": 100} ) assert config._caller_max_results == 10 # TinyFish's natural cap def test_max_results_clamped_lower_stored_on_self(self): config = TinyfishSearchConfig() config.transform_search_request( query="test", optional_params={"max_results": 0} ) assert config._caller_max_results == 1 def test_max_results_normal_stored_on_self(self): config = TinyfishSearchConfig() config.transform_search_request( query="test", optional_params={"max_results": 5} ) assert config._caller_max_results == 5 def test_max_results_non_numeric_string_warns_and_skips(self, caplog): # `int("abc")` would raise ValueError; guard makes the failure visible # via warning and treats the value as if max_results wasn't set. config = TinyfishSearchConfig() with caplog.at_level("WARNING"): result = config.transform_search_request( query="test", optional_params={"max_results": "abc"} ) assert config._caller_max_results is None assert "max_results" not in result["_tinyfish_params"] messages = [r.getMessage() for r in caplog.records] assert any("max_results" in m and "abc" in m for m in messages) def test_max_results_infinity_float_warns_and_skips(self, caplog): # `int(float('inf'))` raises OverflowError, not ValueError/TypeError. # Guard must catch it so a caller passing math.inf gets the same # warn-and-ignore behavior as other malformed values. config = TinyfishSearchConfig() with caplog.at_level("WARNING"): result = config.transform_search_request( query="test", optional_params={"max_results": float("inf")} ) assert config._caller_max_results is None assert "max_results" not in result["_tinyfish_params"] messages = [r.getMessage() for r in caplog.records] assert any("max_results" in m for m in messages) def test_domain_filter_appends_site_operators(self): config = TinyfishSearchConfig() result = config.transform_search_request( query="python tutorials", optional_params={"search_domain_filter": ["arxiv.org", "github.com"]}, ) query_value = result["_tinyfish_params"]["query"] assert "site:arxiv.org" in query_value assert "site:github.com" in query_value assert "(python tutorials) (site:arxiv.org OR site:github.com)" == query_value def test_domain_filter_empty_list_ignored(self): config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"search_domain_filter": []} ) assert result["_tinyfish_params"]["query"] == "test" def test_domain_filter_non_list_ignored(self): config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"search_domain_filter": "not-a-list"} ) assert result["_tinyfish_params"]["query"] == "test" def test_unknown_params_passed_through(self): config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"language": "en", "page": 2} ) params = result["_tinyfish_params"] assert params["language"] == "en" assert params["page"] == 2 def test_perplexity_params_not_passed_through(self): config = TinyfishSearchConfig() supported = config.get_supported_perplexity_optional_params() if supported: param = next(p for p in supported if p != "max_results" and p != "country") result = config.transform_search_request( query="test", optional_params={param: "value"} ) assert param not in result["_tinyfish_params"] def test_arbitrary_param_passed_through(self): # `fetch` is a TinyFish-specific param (JSON-encoded tf-fetch config). # The passthrough loop should forward it verbatim without LiteLLM needing # to know about it. config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"fetch": "{}"} ) assert result["_tinyfish_params"]["fetch"] == "{}" def test_dict_param_auto_json_encoded(self): # Callers naturally pass dict-shaped params; we serialize so the # downstream urlencode step (which only accepts str|int|bool) doesn't reject. config = TinyfishSearchConfig() result = config.transform_search_request( query="test", optional_params={"fetch": {"format": "html", "fetch_path": "fast"}}, ) assert ( result["_tinyfish_params"]["fetch"] == '{"format":"html","fetch_path":"fast"}' ) def test_bool_param_serialized_as_lowercase(self): # urlencode renders Python bool as capitalized "True"/"False"; ux-labs # rejects those (e.g. include_thumbnail must be literal "true"/"false"). # Normalize before passing through. config = TinyfishSearchConfig() true_result = config.transform_search_request( query="test", optional_params={"include_thumbnail": True} ) false_result = config.transform_search_request( query="test", optional_params={"include_thumbnail": False} ) assert true_result["_tinyfish_params"]["include_thumbnail"] == "true" assert false_result["_tinyfish_params"]["include_thumbnail"] == "false" def test_pre_stringified_param_passed_unchanged(self): # If the caller already JSON-encoded, don't re-encode. config = TinyfishSearchConfig() already = '{"format":"html"}' result = config.transform_search_request( query="test", optional_params={"fetch": already} ) assert result["_tinyfish_params"]["fetch"] == already class TestGetCompleteUrl: def test_default_api_base(self): config = TinyfishSearchConfig() with patch( "litellm.llms.tinyfish.search.transformation.get_secret_str", return_value=None, ): url = config.get_complete_url(api_base=None, optional_params={}) assert url == "https://api.search.tinyfish.ai" def test_custom_api_base(self): config = TinyfishSearchConfig() url = config.get_complete_url( api_base="https://custom.api.tinyfish.ai", optional_params={} ) assert url == "https://custom.api.tinyfish.ai" def test_env_api_base(self): config = TinyfishSearchConfig() with patch( "litellm.llms.tinyfish.search.transformation.get_secret_str", return_value="https://env.tinyfish.ai", ): url = config.get_complete_url(api_base=None, optional_params={}) assert url == "https://env.tinyfish.ai" def test_with_tinyfish_params(self): config = TinyfishSearchConfig() with patch( "litellm.llms.tinyfish.search.transformation.get_secret_str", return_value=None, ): url = config.get_complete_url( api_base=None, optional_params={}, data={"_tinyfish_params": {"query": "hello", "max_results": 5}}, ) assert "query=hello" in url assert "max_results=5" in url assert url.startswith("https://api.search.tinyfish.ai?") def test_without_tinyfish_params_key(self): config = TinyfishSearchConfig() with patch( "litellm.llms.tinyfish.search.transformation.get_secret_str", return_value=None, ): url = config.get_complete_url( api_base=None, optional_params={}, data={"other": "value"} ) assert url == "https://api.search.tinyfish.ai" def test_data_none(self): config = TinyfishSearchConfig() with patch( "litellm.llms.tinyfish.search.transformation.get_secret_str", return_value=None, ): url = config.get_complete_url(api_base=None, optional_params={}, data=None) assert url == "https://api.search.tinyfish.ai" class TestTransformSearchResponse: def test_basic_response(self): config = TinyfishSearchConfig() mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert result.object == "search" assert len(result.results) == 2 assert result.results[0].title == "TinyFish - AI Web Automation" assert result.results[0].url == "https://tinyfish.ai" assert ( result.results[0].snippet == "Automate any website with natural language." ) def test_empty_results(self): config = TinyfishSearchConfig() mock_response = _make_mock_response({"results": []}) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert result.object == "search" assert len(result.results) == 0 def test_max_results_truncates_from_self_state(self): config = TinyfishSearchConfig() # Simulate transform_search_request having set the threaded value. config._caller_max_results = 3 many_results = { "results": [ { "title": f"Result {i}", "url": f"https://example.com/{i}", "snippet": f"Snippet {i}", } for i in range(10) ] } mock_response = _make_mock_response(many_results) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert len(result.results) == 3 assert result.results[0].title == "Result 0" assert result.results[2].title == "Result 2" def test_max_results_default_is_tinyfish_cap(self): # No caller value → fall back to TinyFish's natural ceiling (10). config = TinyfishSearchConfig() many_results = { "results": [ { "title": f"Result {i}", "url": f"https://example.com/{i}", "snippet": f"Snippet {i}", } for i in range(15) ] } mock_response = _make_mock_response(many_results) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert len(result.results) == 10 def test_missing_required_fields_default_to_empty_string(self): # title/url/snippet are required by LiteLLM's SearchResult schema. # We default missing/null values to "" so a degraded TinyFish result # flows through instead of failing the whole call. config = TinyfishSearchConfig() mock_response = _make_mock_response( {"results": [{}, {"title": None, "url": None, "snippet": None}]} ) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert len(result.results) == 2 for r in result.results: assert r.title == "" assert r.url == "" assert r.snippet == "" def test_extra_per_result_fields_surface_as_attributes(self): config = TinyfishSearchConfig() mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) first = result.results[0] assert getattr(first, "position", None) == 1 assert getattr(first, "site_name", None) == "tinyfish.ai" def test_fetch_field_rides_through_to_search_result(self): # Mirrors browser-search's per-result `fetch` nested object (see # api/src/parser.rs SearchResult.fetch). Confirms `fetch=...` requests # surface their content to LiteLLM callers without provider changes. config = TinyfishSearchConfig() fetched = { "results": [ { "title": "TinyFish", "url": "https://tinyfish.ai", "snippet": "Web automation.", "fetch": { "url": "https://tinyfish.ai", "title": "TinyFish", "text": "Body text", "cached": False, }, } ] } mock_response = _make_mock_response(fetched) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) first = result.results[0] fetch_field = getattr(first, "fetch", None) assert isinstance(fetch_field, dict) assert fetch_field["text"] == "Body text" def test_no_request_uses_default_max_results(self): config = TinyfishSearchConfig() mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) result = config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert len(result.results) == 2 def test_parameter_warnings_reader_emits_log_lines(self, caplog): # When TinyFish responds with a top-level `parameter_warnings` array # (post-rollout of that contract), each entry is re-fired as a # verbose_logger.warning so callers see what was ignored. config = TinyfishSearchConfig() body = { "results": [ {"title": "x", "url": "https://x", "snippet": "x"}, ], "parameter_warnings": [ { "type": "unsupported", "parameter": "max_tokens_per_page", "message": "Parameter not supported by TinyFish Search.", "docs_url": "https://docs.tinyfish.ai/search-api", }, ], } mock_response = _make_mock_response(body) with caplog.at_level("WARNING"): config.transform_search_response( raw_response=mock_response, logging_obj=None ) messages = [r.getMessage() for r in caplog.records] assert any("max_tokens_per_page" in m for m in messages) # The type is included in the message so agents can branch on it. assert any("unsupported" in m for m in messages) def test_parameter_warnings_absent_no_log(self, caplog): # Absence of the field is silent — most responses won't carry it. config = TinyfishSearchConfig() mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) with caplog.at_level("WARNING"): config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert not any( "TinyFish Search ignored" in r.getMessage() for r in caplog.records ) def test_parameter_warnings_malformed_shapes_never_throw(self): # Every shape that doesn't match {parameter: str, message: str} should # silently no-op. None of these should raise an exception. config = TinyfishSearchConfig() good_results = [{"title": "x", "url": "https://x", "snippet": "x"}] malformed_field_values = [ "not a list", # string 42, # int {"parameter": "x", "message": "y"}, # dict instead of list True, # bool ] for bad_value in malformed_field_values: body = {"results": good_results, "parameter_warnings": bad_value} config.transform_search_response( raw_response=_make_mock_response(body), logging_obj=None ) # must not raise malformed_entries = [ "string in list", # non-dict 42, # int {}, # missing all {"type": "unsupported", "parameter": "x"}, # missing message {"type": "unsupported", "message": "y"}, # missing parameter {"parameter": "x", "message": "y"}, # missing type { "type": "unsupported", "parameter": None, "message": "y", }, # null parameter {"type": "unsupported", "parameter": "x", "message": ""}, # empty message { "type": "unsupported", "parameter": 42, "message": "y", }, # non-string parameter {"type": 1, "parameter": "x", "message": "y"}, # non-string type ] body = {"results": good_results, "parameter_warnings": malformed_entries} config.transform_search_response( raw_response=_make_mock_response(body), logging_obj=None ) # must not raise def test_parameter_warnings_malformed_entries_emit_nothing(self, caplog): config = TinyfishSearchConfig() body = { "results": [{"title": "x", "url": "https://x", "snippet": "x"}], "parameter_warnings": [ {"type": "unsupported", "parameter": "x"}, # missing message — skipped { "type": "unsupported", "parameter": "valid_one", "message": "actual msg", }, # ok — emitted {"parameter": "x", "message": "y"}, # missing type — skipped ], } with caplog.at_level("WARNING"): config.transform_search_response( raw_response=_make_mock_response(body), logging_obj=None ) messages = [r.getMessage() for r in caplog.records] assert sum("parameter_warning" in m for m in messages) == 1 assert any("valid_one" in m for m in messages) class TestErrorHandling: def test_4xx_response_raises_with_attribution_and_unwrapped_message(self): # Reproduces ux-labs' error envelope shape for an INVALID_INPUT response. config = TinyfishSearchConfig() body = { "error": { "code": "INVALID_INPUT", "message": "query is required", "details": [{"field": "query"}], } } mock_response = _make_mock_response(body, status_code=400) with pytest.raises(Exception) as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) msg = str(exc_info.value) assert "TinyFish Search:" in msg assert "query is required" in msg assert "docs.tinyfish.ai/search-api" in msg assert getattr(exc_info.value, "status_code", None) == 400 def test_429_preserves_status_code_and_headers(self): config = TinyfishSearchConfig() body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "60 rpm"}} mock_response = _make_mock_response( body, status_code=429, headers={"Retry-After": "60"} ) with pytest.raises(Exception) as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) assert getattr(exc_info.value, "status_code", None) == 429 headers = getattr(exc_info.value, "headers", {}) or {} assert headers.get("Retry-After") == "60" def test_5xx_with_non_ux_labs_body_falls_back_to_raw_text(self): # Cloudflare-style JSON or any other envelope: unwrap fails, fall back to raw. config = TinyfishSearchConfig() body = {"errors": [{"code": "10000", "message": "Internal"}]} mock_response = _make_mock_response(body, status_code=502) with pytest.raises(Exception) as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) msg = str(exc_info.value) assert "TinyFish Search:" in msg # The raw JSON body string should appear in the message verbatim. assert "10000" in msg def test_non_json_4xx_body_uses_raw_text(self): config = TinyfishSearchConfig() mock_response = _make_mock_response( json_data=None, status_code=502, text="Bad Gateway" ) with pytest.raises(Exception) as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) msg = str(exc_info.value) assert "TinyFish Search:" in msg assert "Bad Gateway" in msg def test_non_json_200_body_routes_through_get_error_class(self): # 200 but the body isn't JSON (degraded backend, CDN-injected page, etc.) config = TinyfishSearchConfig() mock_response = _make_mock_response( json_data=None, status_code=200, text="not json" ) with pytest.raises(Exception) as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) msg = str(exc_info.value) assert "TinyFish Search:" in msg assert "Expected JSON response" in msg def test_wrap_error_returns_attributed_baselm_exception_directly(self): # Direct unit test of the private _wrap_error helper used by # transform_search_response. Network failures don't go through this; # they hit BaseSearchConfig.get_error_class via LiteLLM core. config = TinyfishSearchConfig() body = '{"error": {"code": "UNAUTHORIZED", "message": "bad key"}}' exc = config._wrap_error( error_message=body, status_code=401, headers={"x": "y"} ) msg = str(exc) assert "TinyFish Search:" in msg assert "bad key" in msg assert exc.status_code == 401 def test_schema_mismatch_wraps_with_attribution(self): # When TinyFish returns a 200 with a body shape that doesn't match # LiteLLM's SearchResponse contract (e.g. missing top-level `results`), # raise with TinyFish attribution + docs link so the caller knows to # check TinyFish's schema, not their own input. config = TinyfishSearchConfig() mock_response = _make_mock_response({"query": "x"}) # no `results` key with pytest.raises(Exception) as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) msg = str(exc_info.value) assert "TinyFish Search:" in msg assert "schema" in msg.lower() assert "docs.tinyfish.ai/search-api" in msg class TestAppendDomainFilters: def test_single_domain(self): result = _append_domain_filters("test", ["example.com"]) assert result == "(test) (site:example.com)" def test_multiple_domains(self): result = _append_domain_filters("query", ["a.com", "b.com", "c.com"]) assert result == "(query) (site:a.com OR site:b.com OR site:c.com)" class TestDefaultMissingResultFields: def test_non_dict_raw_json_is_noop(self): # raw_json could be a string/list/None if TinyFish ever returns a # non-envelope shape; the helper just returns without mutating. for payload in ("not a dict", ["list"], None, 42): _default_missing_result_fields(payload) # must not raise def test_non_dict_results_item_skipped(self): # If `results` contains a non-dict entry (string, int, etc.), the helper # skips it; SearchResponse.model_validate will reject it later. raw_json = {"results": ["string item", 42, {"title": "ok"}]} _default_missing_result_fields(raw_json) # Only the dict item gets defaulted; the others are unchanged. assert raw_json["results"][0] == "string item" assert raw_json["results"][1] == 42 assert raw_json["results"][2] == {"title": "ok", "url": "", "snippet": ""}