From ff474a0b9fd6ad96ce56ac5e013f23e910b6a677 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Wed, 9 Sep 2026 09:20:49 -0400 Subject: [PATCH] fix(xai): cap search results to the caller's max_results transform_search_response returned every distinct citation regardless of the caller's max_results, unlike every other search provider. Caps the returned list while still pricing the per-citation surcharge off the true distinct citation count, since xAI consulted all of them internally regardless of how many are returned. --- litellm/llms/xai/search/transformation.py | 20 ++++++++- .../search/test_xai_search_transformation.py | 45 +++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/litellm/llms/xai/search/transformation.py b/litellm/llms/xai/search/transformation.py index 3bab84f31b3..5552c09848f 100644 --- a/litellm/llms/xai/search/transformation.py +++ b/litellm/llms/xai/search/transformation.py @@ -163,6 +163,19 @@ def _requested_model(response_kwargs: Mapping[str, object]) -> str: return _model(optional_params) +def _valid_max_results(max_results: object) -> int | None: + if isinstance(max_results, bool) or not isinstance(max_results, int): + return None + return max_results if max_results > 0 else None + + +def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None: + optional_params: Final = response_kwargs.get("optional_params") + if not isinstance(optional_params, Mapping): + return None + return _valid_max_results(optional_params.get("max_results")) + + def _x_search_tool(optional_params: Mapping[str, object]) -> _XSearchTool: return _XSearchTool( allowed_x_handles=_typed_str_tuple(optional_params.get("allowed_x_handles")), @@ -325,19 +338,22 @@ class XAISearchConfig(BaseSearchConfig): else "unknown reason" ) raise self._upstream_error(f"the search was incomplete: {reason}", raw_response) - return self._priced(results, parsed.usage, _requested_model(kwargs)) + max_results: Final = _requested_max_results(kwargs) + capped_results: Final = results[:max_results] if max_results is not None else results + return self._priced(capped_results, parsed.usage, _requested_model(kwargs), len(results)) def _priced( self, results: tuple[SearchResult, ...], usage: _Usage | None, model: str, + distinct_citation_count: int, ) -> SearchResponse: response: Final = SearchResponse( results=list(results), # mutable-ok: SearchResponse.results is list[SearchResult] object="search", ) - cost: Final = _resolved_cost(usage, model, len(results)) + cost: Final = _resolved_cost(usage, model, distinct_citation_count) if cost is not None: response._hidden_params[ # pyright: ignore[reportPrivateUsage] # response_cost_calculator's own contract "additional_headers" diff --git a/tests/test_litellm/llms/xai/search/test_xai_search_transformation.py b/tests/test_litellm/llms/xai/search/test_xai_search_transformation.py index df005f76a69..f9b467f1529 100644 --- a/tests/test_litellm/llms/xai/search/test_xai_search_transformation.py +++ b/tests/test_litellm/llms/xai/search/test_xai_search_transformation.py @@ -317,6 +317,30 @@ class TestXAISearchConfigTransformResponse: resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) assert [r.url for r in resp.results] == ["https://example.com"] + def test_caps_results_to_max_results(self): + annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(5)] + resp = _config().transform_search_response( + _resp(_message_response("claim", annotations)), logging_obj=Mock(), optional_params={"max_results": 2} + ) + assert [r.url for r in resp.results] == ["https://example.com/0", "https://example.com/1"] + + def test_without_max_results_returns_all_citations(self): + annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(4)] + resp = _config().transform_search_response( + _resp(_message_response("claim", annotations)), logging_obj=Mock() + ) + assert len(resp.results) == 4 + + @pytest.mark.parametrize("max_results", [True, False, 0, -1, "5"]) + def test_ignores_invalid_max_results_and_returns_all_citations(self, max_results: object): + annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(3)] + resp = _config().transform_search_response( + _resp(_message_response("claim", annotations)), + logging_obj=Mock(), + optional_params={"max_results": max_results}, + ) + assert len(resp.results) == 3 + class TestXAISearchConfigCost: def test_uses_xai_reported_ticks_when_present(self): @@ -417,6 +441,27 @@ class TestXAISearchConfigCost: "llm_provider-x-litellm-response-cost" ] == pytest.approx(transformation._PER_CITATION_SURCHARGE_USD * 2) + def test_surcharge_uses_full_citation_count_even_when_results_are_capped(self, monkeypatch: pytest.MonkeyPatch): + def fake_get_model_info(model: str, custom_llm_provider: str): + if model == "xai/grok-4-fast-non-reasoning": + return {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0} + raise Exception("not mapped") + + monkeypatch.setattr(transformation, "get_model_info", fake_get_model_info) + annotations = [_citation(f"https://example.com/{i}", f"T{i}") for i in range(5)] + payload = _message_response( + "claim", + annotations, + usage={"input_tokens": 0, "output_tokens": 0, "output_tokens_details": {"reasoning_tokens": 0}}, + ) + resp = _config().transform_search_response( + _resp(payload), logging_obj=Mock(), optional_params={"max_results": 2} + ) + assert len(resp.results) == 2 + assert resp._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] == pytest.approx(transformation._PER_CITATION_SURCHARGE_USD * 5) + class TestXAISearchConfigGetErrorClass: def test_attributes_the_provider(self):