From 4ad5a003b1928005ef2c2ada3dfe8269f8b5985d Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Tue, 27 Oct 2026 17:28:51 -0700 Subject: [PATCH] feat(search): add native xAI X Search support (#38505) (#48677) * feat(search): add native xAI X Search support (#38505) * fix(search): support xAI search on Python 3.10 (#38505) --- litellm/llms/azure/search/transformation.py | 75 +----- litellm/llms/base_llm/search/responses.py | 74 ++++++ litellm/llms/xai/search/__init__.py | 0 litellm/llms/xai/search/transformation.py | 138 +++++++++++ litellm/types/llms/xai.py | 4 +- litellm/types/utils.py | 1 + litellm/utils.py | 2 + tests/unit/llms/xai/search/__init__.py | 0 tests/unit/llms/xai/search/conftest.py | 8 + .../llms/xai/search/test_transformation.py | 233 ++++++++++++++++++ 10 files changed, 471 insertions(+), 64 deletions(-) create mode 100644 litellm/llms/base_llm/search/responses.py create mode 100644 litellm/llms/xai/search/__init__.py create mode 100644 litellm/llms/xai/search/transformation.py create mode 100644 tests/unit/llms/xai/search/__init__.py create mode 100644 tests/unit/llms/xai/search/conftest.py create mode 100644 tests/unit/llms/xai/search/test_transformation.py diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index f7cdf88ade0..e255a667606 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -38,6 +38,12 @@ from pydantic import ConfigDict, ValidationError from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.search.responses import ( + ResponsesSearchAnnotation, + ResponsesSearchEnvelope, + ResponsesSearchError, + responses_search_citations, +) from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -64,62 +70,10 @@ _UPSTREAM_ERROR_STATUS: Final = 502 _RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" -class _Annotation(LiteLLMBaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - type: str = "" - url: str | None = None - title: str | None = None - start_index: int | None = None - end_index: int | None = None - - -class _ContentPart(LiteLLMBaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - type: str = "" - text: str = "" - annotations: tuple[_Annotation, ...] = () - - -class _OutputItem(LiteLLMBaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - type: str = "" - content: tuple[_ContentPart, ...] = () - - -class _ErrorBody(LiteLLMBaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - message: str | None = None - - -class _IncompleteDetails(LiteLLMBaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - reason: str | None = None - - -class _ResponsesEnvelope(LiteLLMBaseModel): - """A Foundry Responses API body. `output` is required: a body without it is not a - Responses API response and must not be reported as a successful empty search. - - A 200 body can still carry `status` `failed` or `incomplete`; those are surfaced as - errors rather than reported as a successful empty search.""" - - model_config = ConfigDict(extra="ignore", frozen=True) - - output: tuple[_OutputItem, ...] - status: str | None = None - error: _ErrorBody | None = None - incomplete_details: _IncompleteDetails | None = None - - class _ErrorEnvelope(LiteLLMBaseModel): model_config = ConfigDict(extra="ignore", frozen=True) - error: _ErrorBody | None = None + error: ResponsesSearchError | None = None def _unwrap_error_detail(error_message: str) -> str: @@ -138,13 +92,13 @@ def _unwrap_error_detail(error_message: str) -> str: if message is None: return error_message try: - nested: Final = _ErrorBody.model_validate_json(message) + nested: Final = ResponsesSearchError.model_validate_json(message) except ValidationError: return message return nested.message or message -def _snippet(text: str, annotation: _Annotation) -> str: +def _snippet(text: str, annotation: ResponsesSearchAnnotation) -> str: """ The text a citation supports, not the citation marker itself. @@ -160,7 +114,7 @@ def _snippet(text: str, annotation: _Annotation) -> str: return text[:_SNIPPET_FALLBACK_LENGTH] -def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]: +def _citation_results(envelope: ResponsesSearchEnvelope) -> tuple[SearchResult, ...]: """One result per cited URL: first occurrence wins, order preserved as answered.""" cited: Final = tuple( SearchResult( @@ -170,12 +124,7 @@ def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]: date=None, last_updated=None, ) - for item in envelope.output - if item.type == "message" - for part in item.content - if part.type == "output_text" - for annotation in part.annotations - if annotation.type == "url_citation" and annotation.url + for part, annotation in responses_search_citations(envelope.output) ) first_by_url: Final = MappingProxyType({result.url: result for result in reversed(cited)}) return tuple(first_by_url[url] for url in dict.fromkeys(result.url for result in cited)) @@ -380,7 +329,7 @@ class BingGroundingSearchConfig(BaseSearchConfig): **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_response signature ) -> SearchResponse: try: - parsed: Final = _ResponsesEnvelope.model_validate_json(raw_response.content) + parsed: Final = ResponsesSearchEnvelope.model_validate_json(raw_response.content) except ValidationError as e: raise self.get_error_class( error_message=f"response does not match the Foundry Responses API schema: {e}", diff --git a/litellm/llms/base_llm/search/responses.py b/litellm/llms/base_llm/search/responses.py new file mode 100644 index 00000000000..42e8672c8ac --- /dev/null +++ b/litellm/llms/base_llm/search/responses.py @@ -0,0 +1,74 @@ +from collections.abc import Iterator, Sequence + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMBaseModel + + +class ResponsesSearchAnnotation(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + url: str | None = None + title: str | None = None + start_index: int | None = None + end_index: int | None = None + + +class ResponsesSearchContent(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + text: str = "" + annotations: tuple[ResponsesSearchAnnotation, ...] = () + + +class ResponsesSearchOutput(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + content: tuple[ResponsesSearchContent, ...] = () + + +class ResponsesSearchError(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + message: str | None = None + + +class ResponsesSearchIncompleteDetails(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + reason: str | None = None + + +class ResponsesSearchEnvelope(LiteLLMBaseModel): + """A Responses API body. `output` is required: a body without it is not a + Responses API response and must not be reported as a successful empty search. + + A 200 body can still carry `status` `failed` or `incomplete`; those are surfaced as + errors rather than reported as a successful empty search.""" + + model_config = ConfigDict(extra="ignore", frozen=True) + + output: tuple[ResponsesSearchOutput, ...] + status: str | None = None + error: ResponsesSearchError | None = None + incomplete_details: ResponsesSearchIncompleteDetails | None = None + + +def responses_search_text_parts(output: Sequence[ResponsesSearchOutput]) -> Iterator[ResponsesSearchContent]: + for item in output: + if item.type == "message": + yield from (part for part in item.content if part.type == "output_text") + + +def responses_search_citations( + output: Sequence[ResponsesSearchOutput], +) -> Iterator[tuple[ResponsesSearchContent, ResponsesSearchAnnotation]]: + for part in responses_search_text_parts(output): + yield from ( + (part, annotation) + for annotation in part.annotations + if annotation.type == "url_citation" and annotation.url + ) diff --git a/litellm/llms/xai/search/__init__.py b/litellm/llms/xai/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/xai/search/transformation.py b/litellm/llms/xai/search/transformation.py new file mode 100644 index 00000000000..9577ebe544a --- /dev/null +++ b/litellm/llms/xai/search/transformation.py @@ -0,0 +1,138 @@ +from collections.abc import Mapping +from typing import Annotated, Final + +import httpx +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError + +from litellm.constants import XAI_API_BASE +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.search.responses import ( + ResponsesSearchEnvelope, + responses_search_citations, + responses_search_text_parts, +) +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) +from litellm.llms.xai.common_utils import xai_reported_cost_in_usd +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.xai import XAIXSearchTool + +_X_SEARCH_TOOL: Final = TypeAdapter(XAIXSearchTool) + + +class _SearchOptions(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + search_model: str | None = None + max_results: Annotated[int, Field(gt=0, strict=True)] | None = None + + +class _XAIResponsesSearchEnvelope(ResponsesSearchEnvelope): + usage: Mapping[str, object] | None = None + + +class XAISearchConfig(BaseSearchConfig): + """Search X using the Responses API. ``search_model`` selects the Grok model. + + Each citation becomes a result with the answer as its snippet. ``max_results`` + caps returned citations, not upstream tool usage or the cost of the request. + """ + + @staticmethod + def ui_friendly_name() -> str: + return "xAI" + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.validate_environment contract + api_key: str | None = None, + api_base: str | None = None, + **kwargs: object, # kwargs-ok: BaseSearchConfig.validate_environment contract + ) -> dict[str, str]: # mutable-ok: HTTP handler consumes the headers + resolved_key: Final = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("XAI_API_KEY",), + base_env_var="XAI_API_BASE", + default_api_base=XAI_API_BASE, + ) + if not resolved_key: + raise ValueError("Set XAI_API_KEY or pass api_key to use xAI search.") + return { + **headers, + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + } + + def get_complete_url( + self, + api_base: str | None, + optional_params: dict[str, object], # mutable-ok: BaseSearchConfig.get_complete_url contract + data: (dict[str, object] | list[dict[str, object]] | None) = None, # mutable-ok: base contract + **kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url contract + ) -> str: + return ensure_api_route(api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE, "responses") + + def transform_search_request( + self, + query: (str | list[str]), # mutable-ok: BaseSearchConfig.transform_search_request contract + optional_params: dict[str, object], # mutable-ok: base contract + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_request contract + ) -> dict[str, object]: # mutable-ok: HTTP handler sends the JSON body + options: Final = _SearchOptions.model_validate(optional_params) + tool: Final = _X_SEARCH_TOOL.validate_python({**optional_params, "type": "x_search"}) + return { + "model": options.search_model or "grok-4.7", + "input": "\n".join(query) if isinstance(query, list) else query, + "tools": [tool], + "tool_choice": "required", + "store": False, + } + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_response contract + ) -> SearchResponse: + try: + parsed: Final = _XAIResponsesSearchEnvelope.model_validate_json(raw_response.content) + except ValidationError: + raise BaseLLMException( + message="xAI returned an invalid Responses API search response.", + status_code=502, + headers=dict(raw_response.headers), + ) + if parsed.error or parsed.status in ("failed", "incomplete"): + detail: Final = parsed.error.message if parsed.error else f"Search response status was {parsed.status}" + raise BaseLLMException( + message=detail or "xAI search failed", + status_code=502, + headers=dict(raw_response.headers), + ) + cost: Final = xai_reported_cost_in_usd(parsed.usage.get("cost_in_usd_ticks") if parsed.usage else None) + if cost is None: + raise BaseLLMException( + message="xAI search response is missing a valid usage.cost_in_usd_ticks for billing.", + status_code=502, + headers=dict(raw_response.headers), + ) + options: Final = _SearchOptions.model_validate(kwargs.get("optional_params") or {}) + answer: Final = "\n".join(part.text for part in responses_search_text_parts(parsed.output)) + results: Final = tuple( + SearchResult( + title=annotation.title or annotation.url or "", + url=annotation.url or "", + snippet=answer, + ) + for _, annotation in responses_search_citations(parsed.output) + ) + response: Final = SearchResponse(results=list(results[: options.max_results])) + response.hidden_params = {"response_cost": cost} + return response diff --git a/litellm/types/llms/xai.py b/litellm/types/llms/xai.py index 000b0cf5a0b..e52939cdb22 100644 --- a/litellm/types/llms/xai.py +++ b/litellm/types/llms/xai.py @@ -1,5 +1,7 @@ from dataclasses import dataclass -from typing import Literal, TypedDict +from typing import Literal + +from typing_extensions import TypedDict @dataclass(frozen=True, slots=True) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 7e8645fd0dd..2869f9b2577 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4643,6 +4643,7 @@ class SearchProviders(str, Enum): NIMBLE = "nimble" BING_GROUNDING = "bing_grounding" DEEPSEEK = "deepseek" + XAI = "xai" # Create a set of all search provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index 6f0b1ca2e3e..1e095966a43 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10636,6 +10636,7 @@ class ProviderConfigManager: from litellm.llms.serper.search.transformation import SerperSearchConfig from litellm.llms.tavily.search.transformation import TavilySearchConfig from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig + from litellm.llms.xai.search.transformation import XAISearchConfig from litellm.llms.you_com.search.transformation import YouComSearchConfig PROVIDER_TO_CONFIG_MAP: Final = { @@ -10660,6 +10661,7 @@ class ProviderConfigManager: SearchProviders.NIMBLE: NimbleSearchConfig, SearchProviders.BING_GROUNDING: BingGroundingSearchConfig, SearchProviders.DEEPSEEK: DeepSeekSearchConfig, + SearchProviders.XAI: XAISearchConfig, } config_class: Final = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/tests/unit/llms/xai/search/__init__.py b/tests/unit/llms/xai/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/xai/search/conftest.py b/tests/unit/llms/xai/search/conftest.py new file mode 100644 index 00000000000..f9fd5e0716e --- /dev/null +++ b/tests/unit/llms/xai/search/conftest.py @@ -0,0 +1,8 @@ +import pytest + + +@pytest.fixture(autouse=True) +def _search_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "true") + for name in ("XAI_API_KEY", "XAI_API_BASE"): + monkeypatch.delenv(name, raising=False) diff --git a/tests/unit/llms/xai/search/test_transformation.py b/tests/unit/llms/xai/search/test_transformation.py new file mode 100644 index 00000000000..6fb11ee1948 --- /dev/null +++ b/tests/unit/llms/xai/search/test_transformation.py @@ -0,0 +1,233 @@ +import json +from typing import Final + +import openai +import pytest +import respx + +import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse +from litellm.llms.xai.search.transformation import XAISearchConfig + +RESPONSES_URL: Final = "https://api.x.ai/v1/responses" +ANSWER: Final = "Grok update from X." +URLS: Final = ("https://x.com/xai/status/123", "https://x.com/xai/status/456") + + +def _response(ticks: int = 123456789) -> dict[str, object]: + return { + "status": "completed", + "output": [ + {"type": "custom_tool_call", "name": "x_keyword_search"}, + { + "type": "message", + "content": [ + { + "type": "output_text", + "text": ANSWER, + "annotations": [ + { + "type": "url_citation", + "url": URLS[0], + "title": "Update", + "start_index": 0, + "end_index": 0, + }, + { + "type": "url_citation", + "url": URLS[1], + "start_index": 0, + "end_index": 0, + }, + {"type": "url_citation", "url": ""}, + {"type": "file_citation", "file_id": "ignored"}, + ], + } + ], + }, + ], + "usage": { + "input_tokens": 500, + "output_tokens": 100, + "cost_in_usd_ticks": ticks, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_search_returns_citations_and_preserves_the_full_bill_after_limiting_results( + respx_mock: respx.MockRouter, use_async: bool +) -> None: + route = respx_mock.post(RESPONSES_URL).respond(json=_response()) + params = { + "query": "Grok update", + "search_provider": "xai", + "api_key": "test-xai-key", + "max_results": 1, + } + result = await litellm.asearch(**params) if use_async else litellm.search(**params) + + assert isinstance(result, SearchResponse) + assert result.model_dump() == { + "object": "search", + "results": [ + { + "title": "Update", + "url": URLS[0], + "snippet": ANSWER, + "date": None, + "last_updated": None, + } + ], + } + assert result.hidden_params["response_cost"] == pytest.approx(123456789 / 10_000_000_000) + assert json.loads(route.calls.last.request.content) == { + "model": "grok-4.7", + "input": "Grok update", + "tools": [{"type": "x_search"}], + "tool_choice": "required", + "store": False, + } + + +@pytest.mark.asyncio +async def test_configured_search_model_and_filters_reach_xai_through_the_router( + respx_mock: respx.MockRouter, +) -> None: + route = respx_mock.post(RESPONSES_URL).respond(json=_response(0)) + router = litellm.Router( + search_tools=[ + { + "search_tool_name": "social-search", + "litellm_params": { + "search_provider": "xai", + "api_key": "test-xai-key", + "search_model": "grok-test", + "allowed_x_handles": ["xai"], + "from_date": "2026-01-01", + "to_date": "2026-02-01", + "enable_image_understanding": True, + "enable_video_understanding": False, + }, + } + ] + ) + result = await router.asearch(search_tool_name="social-search", query=["first", "second"]) + + assert [r.url for r in result.results] == list(URLS) + assert [r.snippet for r in result.results] == [ANSWER, ANSWER] + assert result.results[1].title == URLS[1] + assert result.hidden_params["response_cost"] == 0.0 + assert json.loads(route.calls.last.request.content) == { + "model": "grok-test", + "input": "first\nsecond", + "tools": [ + { + "type": "x_search", + "allowed_x_handles": ["xai"], + "from_date": "2026-01-01", + "to_date": "2026-02-01", + "enable_image_understanding": True, + "enable_video_understanding": False, + } + ], + "tool_choice": "required", + "store": False, + } + + +@pytest.mark.parametrize( + "payload", + [ + {"output": [], "status": "failed", "error": {"message": "search failed"}}, + {"output": [], "status": "incomplete"}, + {"output": [], "status": "completed", "usage": {}}, + {"output": [], "status": "completed", "usage": {"cost_in_usd_ticks": -1}}, + {"output": [], "status": "completed", "usage": {"cost_in_usd_ticks": True}}, + {"output": [], "status": "completed", "usage": {"cost_in_usd_ticks": "123"}}, + {"unexpected": "body"}, + ], +) +def test_search_does_not_report_failed_or_unbillable_upstream_responses_as_success( + respx_mock: respx.MockRouter, payload: dict[str, object] +) -> None: + respx_mock.post(RESPONSES_URL).respond(json=payload) + with pytest.raises(openai.APIStatusError) as error: + litellm.search(query="test", search_provider="xai", api_key="test-xai-key") + assert error.value.status_code == 502 + + +@pytest.mark.parametrize("status", [401, 429, 500]) +def test_upstream_http_error_keeps_its_status(respx_mock: respx.MockRouter, status: int) -> None: + respx_mock.post(RESPONSES_URL).respond(status, json={"error": "upstream error"}) + with pytest.raises(openai.APIStatusError) as error: + litellm.search(query="test", search_provider="xai", api_key="test-xai-key") + assert error.value.status_code == status + + +def test_search_with_no_citations_still_records_the_bill( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(RESPONSES_URL).respond(json={"output": [], "usage": {"cost_in_usd_ticks": 99}}) + result = litellm.search(query="test", search_provider="xai", api_key="test-xai-key") + assert isinstance(result, SearchResponse) + assert result.results == [] + assert result.hidden_params["response_cost"] == pytest.approx(99 / 10_000_000_000) + + +def test_environment_credentials_cannot_be_sent_to_a_caller_controlled_host( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("XAI_API_KEY", "server-test-key") + monkeypatch.delenv("XAI_API_BASE", raising=False) + with pytest.raises(ValueError, match="Refusing to send the server-configured"): + XAISearchConfig().validate_environment({}, api_base="https://untrusted.example/v1") + + +def test_operator_base_and_explicit_model_are_used_without_mutating_headers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("XAI_API_BASE", "https://gateway.example/v1/") + monkeypatch.setenv("XAI_API_KEY", "server-test-key") + config = XAISearchConfig() + headers = {"X-Custom": "preserved"} + resolved = config.validate_environment(headers) + assert resolved == { + **headers, + "Authorization": "Bearer server-test-key", + "Content-Type": "application/json", + } + assert headers == {"X-Custom": "preserved"} + assert config.get_complete_url(None, {}) == "https://gateway.example/v1/responses" + assert ( + config.get_complete_url("https://gateway.example/v1/responses/", {}) == "https://gateway.example/v1/responses" + ) + assert config.transform_search_request("test", {"search_model": "explicit-model"})["model"] == "explicit-model" + + +@pytest.mark.asyncio +async def test_websearch_interception_returns_xai_citations_and_cost( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger + from litellm.proxy import proxy_server + + respx_mock.post(RESPONSES_URL).respond(json=_response()) + router = litellm.Router( + search_tools=[ + { + "search_tool_name": "social-search", + "litellm_params": {"search_provider": "xai", "api_key": "test-xai-key"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + logger = WebSearchInterceptionLogger(search_tool_name="social-search") + text, result = await logger._execute_search("Grok update") + + assert result is not None + assert [r.url for r in result.results] == list(URLS) + assert result.hidden_params["response_cost"] == pytest.approx(123456789 / 10_000_000_000) + assert all(url in text for url in URLS) + assert ANSWER in text