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

@@ -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}",

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
+        )

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

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)

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

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:

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)

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