mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
Follow-up to #31411 (superseded and merged as #31997). Two related fixes so LiteLLM callers see what TinyFish actually returns, plus small correctness cleanups. ## Response headers surfaced on _hidden_params TinyFish sets useful response headers (x-request-id on every response, retry-after and x-ratelimit-limit on 429s). Previously these were only accessible via BaseLLMException.headers on error paths; on the success path they were dropped entirely. Fix: stash headers on both LiteLLM-conventional channels, matching the pattern used by Gemini / Volcengine / Manus / ChatGPT / OpenAI-responses providers. - `_hidden_params["headers"]` -- raw dict from httpx, all keys lowercased. - `_hidden_params["additional_headers"]` -- passed through process_response_headers, which prefixes any x-litellm-* provider header with `llm_provider-` so downstream LiteLLM code that trusts bare x-litellm-* markers can't be spoofed (values still survive under the prefixed key for observability). ## Top-level response extras (query, total_results, page, future fields) transform_search_response was building a fresh SearchResponse from just `results`, silently dropping every top-level field TinyFish's response carries beyond `results` / `object`. Fix: mutate parsed.results to its truncated slice and return the same SearchResponse instance rather than reconstructing. Every field pydantic populated during model_validate -- declared attributes AND extras (query, total_results, page, parameter_warnings, and any future TinyFish additions) -- survives regardless of which storage bucket holds it. Robust against upstream schema evolution: if LiteLLM later promotes a field from extras to declared, this code needs no change. ## Code cleanup - List-valued custom params JSON-encoded on the wire (matching the existing dict handling), so callers can pass a natural Python list for JSON-array wire params. - URL-encodable-params adapter accepts float in addition to str / int / bool; server-side rejection of a wrong-typed float now surfaces cleanly with `TinyFish Search:` attribution + docs link. - Assorted comment / docstring / test-fixture hygiene (no logic changes). ## Tests 70 unit + integration tests pass locally. Live-tested against production TinyFish with 6 diverse queries (basic / max_results / country=US / language=ja / domain filter / fetch={"format":"html"}) -- all 6 pass every expected-behavior check.
823 lines
34 KiB
Python
823 lines
34 KiB
Python
"""
|
|
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
|
|
# 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)
|
|
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 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"}},
|
|
)
|
|
assert result["_tinyfish_params"]["fetch"] == '{"format":"html"}'
|
|
|
|
def test_bool_param_serialized_as_lowercase(self):
|
|
# 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={"some_bool_param": True}
|
|
)
|
|
false_result = config.transform_search_request(
|
|
query="test", optional_params={"some_bool_param": 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.
|
|
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_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 TinyFish Search's per-result `fetch` nested object.
|
|
# 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 TinyFish Search's 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": "rate limit exceeded"}}
|
|
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 {}
|
|
# httpx lowercases; the exception carries the same dict shape.
|
|
assert headers.get("retry-after") == "60"
|
|
|
|
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)
|
|
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="<html>Bad Gateway</html>"
|
|
)
|
|
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": ""}
|