Merge pull request #32448 from ChenluJi/feat/tinyfish-search-headers-and-extras

feat(tinyfish): surface response headers + top-level response extras

litellm/llms/tinyfish/search/transformation.py

@@ -14,6 +14,7 @@ import httpx
 from pydantic import TypeAdapter, ValidationError
 
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.core_helpers import process_response_headers
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.search.transformation import (
@@ -22,7 +23,7 @@ from litellm.llms.base_llm.search.transformation import (
 )
 from litellm.secret_managers.main import get_secret_str
 
-_UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | bool])
+_UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | float | bool])
 _StrList: Final = TypeAdapter(list[str])
 _StrFrozenSet: Final = TypeAdapter(frozenset[str])
 
@@ -94,16 +95,16 @@ class TinyfishSearchConfig(BaseSearchConfig):
         TinyFish equivalents:
         - ``query`` (str or list[str]) → ``query`` (list joined by spaces)
         - ``country`` → ``location``
-        - ``search_domain_filter`` (list[str]) → folded into the query as
-          ``(<query>) (site:a OR site:b ...)`` (TinyFish has no first-class
-          field today; see ML-2084 for the planned ``include_domains``)
+        - ``search_domain_filter`` (list[str]) → folded into the query using
+          search operators
         - ``max_results`` → not sent on the wire; stashed on
           ``self._caller_max_results`` for client-side response truncation
           (TinyFish doesn't honor it server-side)
         - ``max_tokens_per_page`` → silently dropped (no TinyFish equivalent)
 
         Any other ``optional_params`` keys are forwarded to TinyFish as-is.
-        dict/list values are JSON-encoded so they survive ``urlencode``.
+        dict and list values are JSON-encoded so structured payloads survive
+        ``urlencode``.
 
         Returns:
             ``{_TINYFISH_PARAMS_KEY: <dict of querystring entries>}``.
@@ -144,14 +145,12 @@ class TinyfishSearchConfig(BaseSearchConfig):
         supported_perplexity: Final = _StrFrozenSet.validate_python(raw_supported)
         for param, value in optional_params.items():
             if param not in supported_perplexity and param not in request_data:
-                # `fetch` expects a JSON-encoded object on the wire; accept the
-                # natural Python dict form and serialize here so callers don't
-                # have to pre-stringify.
-                if isinstance(value, dict):
+                # Serialize dicts/lists as JSON so structured params survive urlencode.
+                if isinstance(value, (dict, list)):
                     value = json.dumps(value, separators=(",", ":"))
                 # `urlencode` would render Python bool as "True"/"False"
-                # (capitalized). ux-labs validators require lowercase
-                # "true"/"false" (e.g. `include_thumbnail`); normalize here.
+                # (capitalized). TinyFish Search's bool params require lowercase
+                # "true"/"false" strings on the wire; normalize here.
                 elif isinstance(value, bool):
                     value = "true" if value else "false"
                 request_data[param] = value
@@ -167,17 +166,35 @@ class TinyfishSearchConfig(BaseSearchConfig):
         """
         Transform a TinyFish response to LiteLLM's unified ``SearchResponse``.
 
-        Mappings (per-result):
-        - ``title`` → ``SearchResult.title`` (defaults to ``""`` if missing/null)
-        - ``url`` → ``SearchResult.url`` (defaults to ``""``)
-        - ``snippet`` → ``SearchResult.snippet`` (defaults to ``""``)
-        - all other per-result fields (``position``, ``site_name``,
-          ``thumbnail_url``, ``fetch``, ``fetch_error``, ...) ride through as
-          extras on ``SearchResult`` via its ``extra="allow"`` config.
+        Per-result field handling:
+        - ``title``, ``url``, ``snippet`` are declared on ``SearchResult`` and
+          populated by ``SearchResponse.model_validate`` when present. Missing
+          or ``None`` values are defaulted to ``""`` beforehand by
+          ``_default_missing_result_fields`` so a degraded result flows through
+          instead of failing the whole call.
+        - All undeclared per-result fields (``position``, ``site_name``, and
+          any others TinyFish returns) ride through as extras via
+          ``SearchResult``'s ``extra="allow"`` config — accessible as
+          attributes on the result object or enumerable via
+          ``result.model_extra``.
 
-        Top-level ``parameter_warnings`` (see ML-2085) is read when present and
-        each entry is re-fired via ``verbose_logger.warning``. Absent or
-        malformed entries are silently skipped — never throws.
+        Top-level ``parameter_warnings`` is read when present and each entry
+        is re-fired via ``verbose_logger.warning``. Absent or malformed
+        entries are silently skipped — never throws.
+
+        Top-level extras (``query``, ``total_results``, ``page``, and any
+        future TinyFish additions) ride through via
+        ``SearchResponse.extra="allow"``. The validated response is returned
+        in place after truncating ``results`` to the caller's ``max_results``,
+        so every field pydantic populated survives regardless of which
+        storage bucket (declared attribute or ``__pydantic_extra__``) holds it.
+
+        TinyFish response headers (e.g. ``x-request-id``, ``retry-after``,
+        ``x-ratelimit-limit`` — httpx normalizes header names to lowercase)
+        are stashed on ``response._hidden_params["headers"]`` (raw) and
+        ``response._hidden_params["additional_headers"]`` (sanitized via
+        ``process_response_headers``) so callers can correlate a search with
+        server-side logs.
 
         Error paths routed through ``self._wrap_error`` for uniform
         ``"TinyFish Search: <msg>. See <docs> for details."`` wrapping:
@@ -223,7 +240,12 @@ class TinyfishSearchConfig(BaseSearchConfig):
         _emit_parameter_warnings(parsed)
 
         max_results: Final = self._caller_max_results or _TINYFISH_RESULT_CAP
-        return SearchResponse(results=list(parsed.results[:max_results]))
+        parsed.results = parsed.results[:max_results]
+        raw_headers: Final = dict(raw_response.headers)
+        hidden: Final = parsed._hidden_params  # pyright: ignore[reportPrivateUsage]  # sole hidden-params channel
+        hidden["headers"] = raw_headers
+        hidden["additional_headers"] = process_response_headers(raw_headers)
+        return parsed
 
     def _wrap_error(
         self,
@@ -243,9 +265,9 @@ class TinyfishSearchConfig(BaseSearchConfig):
         carry the ``TinyFish Search:`` prefix — the bare error already names
         the host in the URL, so attribution is implicit there.
         """
-        # ux-labs frontend wraps every error body as {"error": {"code", "message", "details"?}}.
+        # TinyFish Search wraps every error body as {"error": {"code", "message", "details"?}}.
         # Best-effort unwrap to surface the inner message; fall back to the raw body
-        # for non-ux-labs responses (CDN HTML pages, other JSON envelopes, plain text).
+        # for other envelope shapes (CDN HTML pages, other JSON envelopes, plain text).
         inner_message = error_message
         try:
             body: Final[object] = json.loads(error_message)  # any-ok: json.loads -> Any
@@ -290,7 +312,7 @@ def _default_missing_result_fields(raw_json: object) -> None:
 
 
 def _emit_parameter_warnings(parsed: SearchResponse) -> None:
-    """Re-fire TinyFish-side ``parameter_warnings`` (see ML-2085) as warnings.
+    """Re-fire TinyFish-side ``parameter_warnings`` as warnings.
 
     Defensive: skip silently on any shape we don't recognize so a malformed
     entry (or an early/partial rollout of the field) never throws.

tests/search_tests/test_tinyfish_search.py

@@ -35,11 +35,16 @@ MOCK_TINYFISH_RESPONSE = {
 
 
 def _make_mock_response(
-    json_data: dict, status_code: int = 200, request_url: str | None = None
+    json_data: dict,
+    status_code: int = 200,
+    request_url: str | None = None,
+    headers: dict | None = None,
 ) -> MagicMock:
     mock = MagicMock()
     mock.status_code = status_code
     mock.json.return_value = json_data
+    # httpx.Headers normalizes keys to lowercase — mirror production behavior.
+    mock.headers = httpx.Headers(headers or {})
     if request_url:
         mock.request = MagicMock()
         mock.request.url = httpx.URL(request_url)
@@ -163,7 +168,7 @@ class TestTinyfishSearch:
 
     @pytest.mark.asyncio
     async def test_fetch_param_round_trip(self):
-        # End-to-end check: caller passes `fetch=...` (JSON-encoded tf-fetch
+        # End-to-end check: caller passes `fetch=...` (JSON-encoded fetch
         # config); param reaches TinyFish on the request side and the nested
         # `fetch` object on each result surfaces back to the SearchResult on the
         # response side. No LiteLLM-side support code is required.
@@ -235,6 +240,58 @@ class TestTinyfishSearch:
         assert result.results[0].title == "Result 0"
         assert result.results[2].title == "Result 2"
 
+    @pytest.mark.asyncio
+    async def test_top_level_extras_surface_end_to_end(self):
+        # Envelope extras (`query`, `total_results`, `page`) must survive the
+        # full asearch dispatch — proves LiteLLM's entry-point plumbing outside
+        # our transformer doesn't accidentally strip them.
+        os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test"
+
+        mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE)
+
+        with patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get",
+            new_callable=AsyncMock,
+        ) as mock_get:
+            mock_get.return_value = mock_response
+
+            response = await litellm.asearch(
+                query="web automation tools",
+                search_provider="tinyfish",
+            )
+
+            assert getattr(response, "query", None) == "web automation tools"
+            assert getattr(response, "total_results", None) == 2
+            assert getattr(response, "page", None) == 0
+
+    @pytest.mark.asyncio
+    async def test_response_headers_surface_end_to_end(self):
+        # Response headers must land on `_hidden_params` after the full
+        # asearch dispatch (both raw and sanitized channels).
+        os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test"
+
+        mock_response = _make_mock_response(
+            MOCK_TINYFISH_RESPONSE,
+            headers={"X-Request-ID": "req-e2e-1"},
+        )
+
+        with patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get",
+            new_callable=AsyncMock,
+        ) as mock_get:
+            mock_get.return_value = mock_response
+
+            response = await litellm.asearch(
+                query="test",
+                search_provider="tinyfish",
+            )
+
+            raw = response._hidden_params["headers"]
+            add = response._hidden_params["additional_headers"]
+            # httpx lowercases; both channels agree on the value.
+            assert raw["x-request-id"] == "req-e2e-1"
+            assert add["llm_provider-x-request-id"] == "req-e2e-1"
+
     @pytest.mark.asyncio
     async def test_empty_results(self):
         os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test"

tests/test_litellm/llms/tinyfish/test_tinyfish_search.py

@@ -47,7 +47,9 @@ def _make_mock_response(
 
     mock = MagicMock()
     mock.status_code = status_code
-    mock.headers = headers or {}
+    # httpx.Headers normalizes keys to lowercase — mirror production so tests
+    # assert what callers actually see.
+    mock.headers = httpx.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)
@@ -222,7 +224,7 @@ class TestTransformSearchRequest:
             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).
+        # `fetch` is a TinyFish-specific param (JSON-encoded fetch config).
         # The passthrough loop should forward it verbatim without LiteLLM needing
         # to know about it.
         config = TinyfishSearchConfig()
@@ -237,26 +239,49 @@ class TestTransformSearchRequest:
         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"}'
+            optional_params={"fetch": {"format": "html"}},
         )
+        assert result["_tinyfish_params"]["fetch"] == '{"format":"html"}'
 
     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.
+        # urlencode renders Python bool as capitalized "True"/"False"; TinyFish
+        # Search's bool params require lowercase "true"/"false" strings on the
+        # wire. Normalize before passing through.
         config = TinyfishSearchConfig()
         true_result = config.transform_search_request(
-            query="test", optional_params={"include_thumbnail": True}
+            query="test", optional_params={"some_bool_param": True}
         )
         false_result = config.transform_search_request(
-            query="test", optional_params={"include_thumbnail": False}
+            query="test", optional_params={"some_bool_param": False}
         )
-        assert true_result["_tinyfish_params"]["include_thumbnail"] == "true"
-        assert false_result["_tinyfish_params"]["include_thumbnail"] == "false"
+        assert true_result["_tinyfish_params"]["some_bool_param"] == "true"
+        assert false_result["_tinyfish_params"]["some_bool_param"] == "false"
+
+    def test_float_param_passes_through(self):
+        # Float values pass the urlencode adapter and land on the wire as
+        # their decimal string form. If TinyFish's server rejects a float
+        # for a param it expects as int, the server's 400 response is
+        # attributed via _wrap_error (`TinyFish Search: ...`) — better than
+        # a client-side pydantic ValidationError with no context.
+        config = TinyfishSearchConfig()
+        result = config.transform_search_request(
+            query="test",
+            optional_params={"some_float_param": 0.5},
+        )
+        assert result["_tinyfish_params"]["some_float_param"] == 0.5
+
+    def test_list_param_auto_json_encoded(self):
+        # TinyFish Search's JSON-array params arrive on the wire as JSON-
+        # encoded strings. Accept the natural Python list form and serialize
+        # so the caller doesn't have to pre-stringify. Params whose wire
+        # format is a plain comma-separated string are the caller's
+        # responsibility to pass as a Python str.
+        config = TinyfishSearchConfig()
+        result = config.transform_search_request(
+            query="test",
+            optional_params={"some_list_param": ["a.example", "b.example"]},
+        )
+        assert result["_tinyfish_params"]["some_list_param"] == '["a.example","b.example"]'
 
     def test_pre_stringified_param_passed_unchanged(self):
         # If the caller already JSON-encoded, don't re-encode.
@@ -422,10 +447,104 @@ class TestTransformSearchResponse:
         assert getattr(first, "position", None) == 1
         assert getattr(first, "site_name", None) == "tinyfish.ai"
 
+    def test_top_level_extras_flow_through(self):
+        # TinyFish returns `query`, `total_results`, `page` at the envelope
+        # level. These must ride through to the caller via SearchResponse's
+        # extra="allow" so pagination logic, echo checks, etc. work.
+        config = TinyfishSearchConfig()
+        mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE)
+        result = config.transform_search_response(
+            raw_response=mock_response, logging_obj=None
+        )
+        assert getattr(result, "query", None) == "web automation tools"
+        assert getattr(result, "total_results", None) == 2
+        assert getattr(result, "page", None) == 0
+
+    def test_top_level_future_extras_flow_through(self):
+        # Any future TinyFish top-level field must ride through unchanged
+        # (design contract: no LiteLLM code change needed for new fields).
+        config = TinyfishSearchConfig()
+        body = {
+            "results": [
+                {"title": "x", "url": "https://x", "snippet": "x"},
+            ],
+            "query": "test",
+            "example_int_extra": 123,  # hypothetical future field
+            "example_str_extra": "value",  # hypothetical future field
+            "example_id_extra": "abc-def",  # hypothetical future field
+        }
+        result = config.transform_search_response(
+            raw_response=_make_mock_response(body), logging_obj=None
+        )
+        assert getattr(result, "example_int_extra", None) == 123
+        assert getattr(result, "example_str_extra", None) == "value"
+        assert getattr(result, "example_id_extra", None) == "abc-def"
+
+    def test_response_headers_stashed_on_hidden_params(self):
+        # TinyFish Search sets X-Request-ID on every success response. Confirm it
+        # lands on both `_hidden_params["headers"]` (raw) and
+        # `_hidden_params["additional_headers"]` (sanitized/prefixed).
+        # httpx.Headers lowercases every key, so assertions use lowercase.
+        config = TinyfishSearchConfig()
+        mock_response = _make_mock_response(
+            MOCK_TINYFISH_RESPONSE,
+            headers={"X-Request-ID": "req-abc-123", "Content-Type": "application/json"},
+        )
+        result = config.transform_search_response(
+            raw_response=mock_response, logging_obj=None
+        )
+        # Raw copy — httpx has normalized keys to lowercase.
+        assert result._hidden_params["headers"]["x-request-id"] == "req-abc-123"
+        # process_response_headers prefixes non-OpenAI-standard keys with "llm_provider-".
+        assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-abc-123"
+
+    def test_response_headers_future_headers_flow_through(self):
+        # "Accept extra": any header TinyFish Search adds later must ride
+        # through without a LiteLLM code change.
+        config = TinyfishSearchConfig()
+        mock_response = _make_mock_response(
+            MOCK_TINYFISH_RESPONSE,
+            headers={
+                "X-Request-ID": "req-1",
+                "X-Example-Header-A": "value-a",  # hypothetical future header
+                "X-Example-Header-B": "value-b",  # hypothetical future header
+            },
+        )
+        result = config.transform_search_response(
+            raw_response=mock_response, logging_obj=None
+        )
+        raw = result._hidden_params["headers"]
+        # httpx lowercases header names on read.
+        assert raw["x-example-header-a"] == "value-a"
+        assert raw["x-example-header-b"] == "value-b"
+
+    def test_response_headers_strips_x_litellm_spoof(self):
+        # A provider setting `x-litellm-*` in its response must not be able to
+        # spoof LiteLLM-internal markers via _hidden_params["additional_headers"].
+        # The raw copy preserves the header (opt-in debug view); the sanitized
+        # copy prefixes it with `llm_provider-` so bare `x-litellm-*` markers
+        # can't be spoofed (values still survive under the prefixed key for
+        # observability).
+        config = TinyfishSearchConfig()
+        mock_response = _make_mock_response(
+            MOCK_TINYFISH_RESPONSE,
+            headers={"x-litellm-attempted-fallbacks": "spoofed", "X-Request-ID": "r1"},
+        )
+        result = config.transform_search_response(
+            raw_response=mock_response, logging_obj=None
+        )
+        # Raw view still has the spoof.
+        assert result._hidden_params["headers"]["x-litellm-attempted-fallbacks"] == "spoofed"
+        # Sanitized view: the spoof survives only under the llm_provider- prefix
+        # (never under the bare x-litellm-* key that LiteLLM downstream trusts).
+        additional = result._hidden_params["additional_headers"]
+        assert "x-litellm-attempted-fallbacks" not in additional
+        assert additional.get("llm_provider-x-litellm-attempted-fallbacks") == "spoofed"
+
     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.
+        # Mirrors TinyFish Search's per-result `fetch` nested object.
+        # Confirms `fetch=...` requests surface their content to LiteLLM
+        # callers without provider changes.
         config = TinyfishSearchConfig()
         fetched = {
             "results": [
@@ -568,7 +687,7 @@ class TestTransformSearchResponse:
 
 class TestErrorHandling:
     def test_4xx_response_raises_with_attribution_and_unwrapped_message(self):
-        # Reproduces ux-labs' error envelope shape for an INVALID_INPUT response.
+        # Reproduces TinyFish Search's error envelope shape for an INVALID_INPUT response.
         config = TinyfishSearchConfig()
         body = {
             "error": {
@@ -590,7 +709,7 @@ class TestErrorHandling:
 
     def test_429_preserves_status_code_and_headers(self):
         config = TinyfishSearchConfig()
-        body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "60 rpm"}}
+        body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "rate limit exceeded"}}
         mock_response = _make_mock_response(
             body, status_code=429, headers={"Retry-After": "60"}
         )
@@ -600,10 +719,12 @@ class TestErrorHandling:
             )
         assert getattr(exc_info.value, "status_code", None) == 429
         headers = getattr(exc_info.value, "headers", {}) or {}
-        assert headers.get("Retry-After") == "60"
+        # httpx lowercases; the exception carries the same dict shape.
+        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.
+    def test_5xx_with_non_tinyfish_envelope_shape_falls_back_to_raw_text(self):
+        # A JSON body that doesn't match TinyFish Search's error envelope shape:
+        # unwrap fails, fall back to the raw body text.
         config = TinyfishSearchConfig()
         body = {"errors": [{"code": "10000", "message": "Internal"}]}
         mock_response = _make_mock_response(body, status_code=502)