diff --git a/litellm/llms/xai/search/__init__.py b/litellm/llms/xai/search/__init__.py new file mode 100644 index 00000000000..c061a045c5d --- /dev/null +++ b/litellm/llms/xai/search/__init__.py @@ -0,0 +1,7 @@ +""" +xAI Search API module. +""" + +from litellm.llms.xai.search.transformation import XAISearchConfig + +__all__ = ["XAISearchConfig"] # mutable-ok: module __all__ convention requires a list diff --git a/litellm/llms/xai/search/transformation.py b/litellm/llms/xai/search/transformation.py new file mode 100644 index 00000000000..5552c09848f --- /dev/null +++ b/litellm/llms/xai/search/transformation.py @@ -0,0 +1,382 @@ +""" +Calls xAI's Responses API with the `x_search` tool (xAI's Live Search over X/Twitter). + +xAI docs: https://docs.x.ai/docs/guides/live-search + +Setup: + Set XAI_API_KEY (or litellm.xai_key), the same credential xAI chat/responses calls use. + Optional: pass model=... in optional_params to override the default (grok-4-fast). + +Usage: + response = litellm.search( + query="what happened at the last SpaceX launch", + search_provider="xai", + ) +""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +import httpx +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.constants import XAI_API_BASE +from litellm.llms.base_llm.chat.transformation import BaseLLMException +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.utils import get_model_info + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelInfo + +_RESPONSES_PATH: Final = "/responses" +_DEFAULT_MODEL: Final = "grok-4-fast" +_UPSTREAM_ERROR_STATUS: Final = 502 +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" +_API_KEY_ENV_VAR: Final = "XAI_API_KEY" +_API_BASE_ENV_VAR: Final = "XAI_API_BASE" +# https://docs.x.ai/developers/pricing#tools-pricing — proxy for "sources used" only when +# xAI's own cost_in_usd_ticks isn't present on the response (see _resolved_cost). +_PER_CITATION_SURCHARGE_USD: Final = 0.025 + + +class _Annotation(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + url: str | None = None + title: str | None = None + + +class _ContentPart(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + text: str = "" + annotations: tuple[_Annotation, ...] = () + + +class _OutputItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + content: tuple[_ContentPart, ...] = () + + +class _ErrorBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + message: str | None = None + + +class _IncompleteDetails(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + reason: str | None = None + + +class _OutputTokensDetails(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + reasoning_tokens: int | None = None + + +class _Usage(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + input_tokens: int | None = None + output_tokens: int | None = None + output_tokens_details: _OutputTokensDetails | None = None + cost_in_usd_ticks: int | None = None + + +class _ResponsesEnvelope(BaseModel): + """An xAI 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 + usage: _Usage | None = None + + +class _XSearchTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["x_search"] = "x_search" + allowed_x_handles: tuple[str, ...] | None = None + excluded_x_handles: tuple[str, ...] | None = None + from_date: str | None = None + to_date: str | None = None + enable_image_understanding: bool | None = None + enable_video_understanding: bool | None = None + + +class _ResponsesRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + model: str + input: str + tools: tuple[_XSearchTool, ...] + + +def _typed_str(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _typed_bool(value: object) -> bool | None: + return value if isinstance(value, bool) else None + + +def _typed_str_tuple(value: object) -> tuple[str, ...] | None: + if isinstance(value, bool) or not isinstance(value, (list, tuple)): + return None + if not all(isinstance(item, str) for item in value): + return None + return tuple(value) + + +def _model(optional_params: Mapping[str, object]) -> str: + model: Final = optional_params.get("model") + return model if isinstance(model, str) and model else _DEFAULT_MODEL + + +def _requested_model(response_kwargs: Mapping[str, object]) -> str: + optional_params: Final = response_kwargs.get("optional_params") + if not isinstance(optional_params, Mapping): + return _DEFAULT_MODEL + 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")), + excluded_x_handles=_typed_str_tuple(optional_params.get("excluded_x_handles")), + from_date=_typed_str(optional_params.get("from_date")), + to_date=_typed_str(optional_params.get("to_date")), + enable_image_understanding=_typed_bool(optional_params.get("enable_image_understanding")), + enable_video_understanding=_typed_bool(optional_params.get("enable_video_understanding")), + ) + + +def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]: + """One result per distinct cited URL, in first-appearance order. + + xAI's url_citation annotations carry start_index/end_index, but they are reportedly + always 0 in practice, so there is no reliable span to slice a per-citation excerpt + from; the full synthesized answer is shared as `snippet` across every citation + instead of attempting a per-citation slice. + """ + full_text: Final = "\n\n".join( + part.text + for item in envelope.output + if item.type == "message" + for part in item.content + if part.type == "output_text" + ) + citations: Final = tuple( + (annotation.url, annotation.title or "") + 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 + ) + first_title_by_url: Final = MappingProxyType({url: title for url, title in reversed(citations)}) + return tuple( + SearchResult(title=first_title_by_url[url], url=url, snippet=full_text, date=None, last_updated=None) + for url in dict.fromkeys(url for url, _ in citations) + ) + + +def _model_info(model: str) -> ModelInfo | None: + try: + return get_model_info(model=f"xai/{model}", custom_llm_provider="xai") + except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model + return None + + +def _token_and_surcharge_cost(usage: _Usage, model: str, distinct_citation_count: int) -> float | None: + """ + Fallback cost when xAI doesn't report cost_in_usd_ticks: real per-token cost from + the model cost map plus a per-source surcharge proxied by distinct citation count. + + Tries the bare model name first, then the -reasoning/-non-reasoning suffix picked + by whether any reasoning tokens were billed, since only the suffixed variants carry + a price entry for grok-4-fast-family models. + """ + reasoning_tokens: Final = usage.output_tokens_details.reasoning_tokens if usage.output_tokens_details else None + suffix: Final = "-reasoning" if (reasoning_tokens or 0) > 0 else "-non-reasoning" + info: Final = _model_info(model) or _model_info(f"{model}{suffix}") + if info is None: + return None + input_cost: Final = (usage.input_tokens or 0) * float(info.get("input_cost_per_token") or 0.0) + output_cost: Final = (usage.output_tokens or 0) * float(info.get("output_cost_per_token") or 0.0) + return input_cost + output_cost + _PER_CITATION_SURCHARGE_USD * distinct_citation_count + + +def _resolved_cost(usage: _Usage | None, model: str, distinct_citation_count: int) -> float | None: + if usage is None: + return None + reported: Final = xai_reported_cost_in_usd(usage.cost_in_usd_ticks) + if reported is not None: + return reported + return _token_and_surcharge_cost(usage, model, distinct_citation_count) + + +class XAISearchConfig(BaseSearchConfig): + """ + x_search only exists as a tool inside an xAI Responses API turn: the model plans + its own queries, calls x_search itself, and synthesizes a prose answer with inline + url_citation annotations. There is no discrete results list the way every other + Search API provider returns one, so this config reverse-engineers a SearchResponse + out of that Responses API turn instead of calling a dedicated search endpoint. + """ + + @staticmethod + def ui_friendly_name() -> str: + return "xAI Live Search (x_search)" + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.validate_environment signature + api_key: str | None = None, + api_base: str | None = None, + **kwargs: object, # kwargs-ok: BaseSearchConfig.validate_environment signature + ) -> dict[str, str]: # mutable-ok: httpx requires a plain dict of headers + resolved_key: Final = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=(_API_KEY_ENV_VAR,), + base_env_var=_API_BASE_ENV_VAR, + default_api_base=XAI_API_BASE, + ) + if not resolved_key: + raise ValueError(f"{_API_KEY_ENV_VAR} is required. Set it, or pass api_key explicitly.") + return { # mutable-ok: httpx requires a plain dict of headers + **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 signature + data: dict[str, object] | list[dict[str, object]] | None = None, # mutable-ok: base signature + **kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url signature + ) -> str: + resolved_base: Final = (api_base or get_secret_str(_API_BASE_ENV_VAR) or XAI_API_BASE).rstrip("/") + if resolved_base.endswith(_RESPONSES_PATH): + return resolved_base + return f"{resolved_base}{_RESPONSES_PATH}" + + def transform_search_request( + self, + query: str | list[str], # mutable-ok: BaseSearchConfig.transform_search_request signature + optional_params: dict[str, object], # mutable-ok: base signature + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_request signature + ) -> dict[str, object]: # mutable-ok: the http handler passes this straight to httpx as the JSON body + request: Final = _ResponsesRequest( + model=_model(optional_params), + input=" ".join(query) if isinstance(query, list) else query, + tools=(_x_search_tool(optional_params),), + ) + return request.model_dump(mode="json", exclude_none=True) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_response signature + ) -> SearchResponse: + try: + parsed: Final = _ResponsesEnvelope.model_validate_json(raw_response.content) + except ValidationError as e: + raise self.get_error_class( + error_message=f"response does not match the xAI Responses API schema: {e}", + status_code=raw_response.status_code, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + if parsed.status == "failed": + detail: Final = parsed.error.message if parsed.error and parsed.error.message else "the search failed" + raise self._upstream_error(detail, raw_response) + results: Final = _citation_results(parsed) + if not results and parsed.status == "incomplete": + reason: Final = ( + parsed.incomplete_details.reason + if parsed.incomplete_details and parsed.incomplete_details.reason + else "unknown reason" + ) + raise self._upstream_error(f"the search was incomplete: {reason}", raw_response) + 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, distinct_citation_count) + if cost is not None: + response._hidden_params[ # pyright: ignore[reportPrivateUsage] # response_cost_calculator's own contract + "additional_headers" + ] = { # mutable-ok: response_cost_calculator writes into _hidden_params + _RESPONSE_COST_HEADER: cost + } + return response + + def _upstream_error(self, detail: str, raw_response: httpx.Response) -> Exception: + return self.get_error_class( + error_message=detail, + status_code=_UPSTREAM_ERROR_STATUS, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.get_error_class signature + ) -> Exception: + return BaseLLMException( + status_code=status_code, + message=f"xAI x_search: {error_message}", + headers=headers, + ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e3ea37dc0c8..f045a3179a9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4061,6 +4061,7 @@ class SearchProviders(str, Enum): AGENTCORE = "agentcore" NIMBLE = "nimble" BING_GROUNDING = "bing_grounding" + XAI = "xai" # Create a set of all search provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index 1a77655a5a4..e470c438bc8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9457,6 +9457,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 = { @@ -9480,6 +9481,7 @@ class ProviderConfigManager: SearchProviders.AGENTCORE: AgentCoreSearchConfig, SearchProviders.NIMBLE: NimbleSearchConfig, SearchProviders.BING_GROUNDING: BingGroundingSearchConfig, + SearchProviders.XAI: XAISearchConfig, } config_class: Final = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: 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 new file mode 100644 index 00000000000..f9b467f1529 --- /dev/null +++ b/tests/test_litellm/llms/xai/search/test_xai_search_transformation.py @@ -0,0 +1,470 @@ +""" +Tests for XAI Search API transformation (the x_search Live Search tool). + +Tests the XAISearchConfig class that reverse-engineers a SearchResponse out of an +xAI Responses API turn, since x_search only exists as a tool inside that turn. + +Source: litellm/llms/xai/search/transformation.py +""" + +import json +from unittest.mock import Mock + +import pytest + +from litellm.llms.xai.search import transformation +from litellm.llms.xai.search.transformation import XAISearchConfig +from litellm.types.utils import SearchProviders +from litellm.utils import ProviderConfigManager + +XAI_RESPONSES_URL = "https://api.x.ai/v1/responses" + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch: pytest.MonkeyPatch): + for var in ("XAI_API_KEY", "XAI_API_BASE"): + monkeypatch.delenv(var, raising=False) + + +def _config() -> XAISearchConfig: + return XAISearchConfig() + + +def _resp(payload, status_code: int = 200) -> Mock: + r = Mock() + r.status_code = status_code + r.headers = {} + r.content = (payload if isinstance(payload, str) else json.dumps(payload)).encode() + return r + + +def _citation(url: str, title: str) -> dict: + return {"type": "url_citation", "url": url, "title": title, "start_index": 0, "end_index": 0} + + +def _message_response(text: str, annotations: list, usage: dict | None = None) -> dict: + payload = { + "output": [ + {"type": "x_search_call", "status": "completed"}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": annotations}], + }, + ] + } + if usage is not None: + payload["usage"] = usage + return payload + + +REAL_SHAPED_FIXTURE = _message_response( + text=( + "SpaceX's latest Starship flight reached orbit and completed a controlled reentry, " + "according to the official SpaceX account and independent orbital trackers." + ), + annotations=[ + _citation("https://x.com/SpaceX/status/1234567890", "SpaceX on X"), + _citation("https://x.com/spacetrackorg/status/2234567890", "Space Track on X"), + _citation("https://x.com/SpaceX/status/1234567890", "SpaceX on X"), + ], + usage={ + "input_tokens": 512, + "output_tokens": 128, + "output_tokens_details": {"reasoning_tokens": 64}, + "cost_in_usd_ticks": 45000000, + }, +) + + +class TestXAISearchConfigRegistration: + def test_provider_registration(self): + config = ProviderConfigManager.get_provider_search_config(provider=SearchProviders.XAI) + assert isinstance(config, XAISearchConfig) + + def test_ui_friendly_name(self): + assert _config().ui_friendly_name() == "xAI Live Search (x_search)" + + +class TestXAISearchConfigValidateEnvironment: + def test_uses_caller_api_key(self): + headers = _config().validate_environment({}, api_key="caller-key") + assert headers["Authorization"] == "Bearer caller-key" + assert headers["Content-Type"] == "application/json" + + def test_reads_env_key(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("XAI_API_KEY", "env-key") + headers = _config().validate_environment({}) + assert headers["Authorization"] == "Bearer env-key" + + def test_missing_key_raises(self): + with pytest.raises(ValueError, match="XAI_API_KEY"): + _config().validate_environment({}) + + def test_does_not_mutate_caller_headers(self): + caller_headers = {"X-Custom": "keep-me"} + result = _config().validate_environment(caller_headers, api_key="k") + assert caller_headers == {"X-Custom": "keep-me"} + assert result["X-Custom"] == "keep-me" + + def test_refuses_env_key_for_untrusted_api_base(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("XAI_API_KEY", "env-key") + with pytest.raises(ValueError, match="Refusing to send the server-configured"): + _config().validate_environment({}, api_base="https://attacker.example.com") + + def test_allows_env_key_for_default_trusted_api_base(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("XAI_API_KEY", "env-key") + headers = _config().validate_environment({}, api_base="https://api.x.ai/v1") + assert headers["Authorization"] == "Bearer env-key" + + def test_allows_env_key_for_api_base_matching_xai_api_base_env_var(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("XAI_API_KEY", "env-key") + monkeypatch.setenv("XAI_API_BASE", "https://custom.x.ai/v1") + headers = _config().validate_environment({}, api_base="https://custom.x.ai/v1") + assert headers["Authorization"] == "Bearer env-key" + + def test_caller_api_key_bypasses_trust_check(self): + headers = _config().validate_environment( + {}, api_key="caller-key", api_base="https://attacker.example.com" + ) + assert headers["Authorization"] == "Bearer caller-key" + + +class TestXAISearchConfigGetCompleteUrl: + def test_default(self): + assert _config().get_complete_url(None, {}) == XAI_RESPONSES_URL + + def test_reads_env_var(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("XAI_API_BASE", "https://custom.x.ai/v1") + assert _config().get_complete_url(None, {}) == "https://custom.x.ai/v1/responses" + + def test_explicit_api_base_overrides_env(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("XAI_API_BASE", "https://ignored.x.ai/v1") + assert _config().get_complete_url("https://explicit.x.ai/v1", {}) == "https://explicit.x.ai/v1/responses" + + @pytest.mark.parametrize( + "api_base", + [ + "https://api.x.ai/v1", + "https://api.x.ai/v1/", + "https://api.x.ai/v1/responses", + "https://api.x.ai/v1/responses/", + ], + ) + def test_appends_responses_path_exactly_once(self, api_base: str): + assert _config().get_complete_url(api_base, {}) == XAI_RESPONSES_URL + + +class TestXAISearchConfigTransformRequest: + def test_default_model(self): + body = _config().transform_search_request("latest AI developments", {}) + assert body["model"] == "grok-4-fast" + assert body["input"] == "latest AI developments" + assert body["tools"] == [{"type": "x_search"}] + + def test_custom_model(self): + body = _config().transform_search_request("q", {"model": "grok-4"}) + assert body["model"] == "grok-4" + + def test_joins_list_query(self): + assert _config().transform_search_request(["foo", "bar"], {})["input"] == "foo bar" + + def test_x_search_filters_included(self): + body = _config().transform_search_request( + "q", + { + "allowed_x_handles": ["spacex", "nasa"], + "excluded_x_handles": ["spam"], + "from_date": "2026-01-01", + "to_date": "2026-02-01", + "enable_image_understanding": True, + "enable_video_understanding": False, + }, + ) + assert body["tools"] == [ + { + "type": "x_search", + "allowed_x_handles": ["spacex", "nasa"], + "excluded_x_handles": ["spam"], + "from_date": "2026-01-01", + "to_date": "2026-02-01", + "enable_image_understanding": True, + "enable_video_understanding": False, + } + ] + + def test_omits_absent_filters(self): + body = _config().transform_search_request("q", {}) + assert body["tools"] == [{"type": "x_search"}] + + def test_ignores_unrelated_optional_params(self): + body = _config().transform_search_request("q", {"max_results": 5, "search_domain_filter": ["x.com"]}) + assert body["tools"] == [{"type": "x_search"}] + + @pytest.mark.parametrize( + "key,bad_value", + [ + ("allowed_x_handles", "not-a-list"), + ("excluded_x_handles", 123), + ("from_date", 20260101), + ("to_date", True), + ("enable_image_understanding", "yes"), + ("enable_video_understanding", 1), + ], + ) + def test_ignores_wrong_typed_filters(self, key: str, bad_value: object): + body = _config().transform_search_request("q", {key: bad_value}) + assert body["tools"] == [{"type": "x_search"}] + + @pytest.mark.parametrize( + "key,bad_list", + [ + ("allowed_x_handles", ["spacex", None]), + ("allowed_x_handles", ["spacex", 123]), + ("excluded_x_handles", [True, "spam"]), + ], + ) + def test_ignores_list_filters_with_non_string_elements(self, key: str, bad_list: list): + body = _config().transform_search_request("q", {key: bad_list}) + assert body["tools"] == [{"type": "x_search"}] + + +class TestXAISearchConfigTransformResponse: + def test_real_shaped_fixture_dedupes_and_preserves_order(self): + resp = _config().transform_search_response(_resp(REAL_SHAPED_FIXTURE), logging_obj=Mock()) + assert resp.object == "search" + assert [r.url for r in resp.results] == [ + "https://x.com/SpaceX/status/1234567890", + "https://x.com/spacetrackorg/status/2234567890", + ] + assert resp.results[0].title == "SpaceX on X" + assert resp.results[1].title == "Space Track on X" + + def test_shares_full_text_as_snippet_across_citations(self): + resp = _config().transform_search_response(_resp(REAL_SHAPED_FIXTURE), logging_obj=Mock()) + expected_snippet = REAL_SHAPED_FIXTURE["output"][1]["content"][0]["text"] + assert all(r.snippet == expected_snippet for r in resp.results) + + def test_multiple_message_items_join_text_with_a_separator(self): + payload = { + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "first claim", + "annotations": [_citation("https://example.com/a", "A")], + } + ], + }, + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "second claim", + "annotations": [_citation("https://example.com/b", "B")], + } + ], + }, + ] + } + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert all(r.snippet == "first claim\n\nsecond claim" for r in resp.results) + + def test_ignores_non_citation_annotations(self): + payload = _message_response("text", [{"type": "file_citation", "url": "https://example.com"}]) + assert _config().transform_search_response(_resp(payload), logging_obj=Mock()).results == [] + + def test_ignores_citation_without_url(self): + payload = _message_response("text", [{"type": "url_citation", "title": "no url"}]) + assert _config().transform_search_response(_resp(payload), logging_obj=Mock()).results == [] + + def test_no_message_output_returns_empty_results(self): + payload = {"output": [{"type": "x_search_call", "status": "completed"}]} + assert _config().transform_search_response(_resp(payload), logging_obj=Mock()).results == [] + + @pytest.mark.parametrize( + "body", + [ + "502 Bad Gateway", + '{"output": "garbage"}', + '{"output": null}', + "{}", + ], + ) + def test_malformed_body_raises_instead_of_reporting_empty(self, body: str): + with pytest.raises(Exception, match="xAI x_search"): + _config().transform_search_response(_resp(body, status_code=502), logging_obj=Mock()) + + def test_failed_status_raises_with_error_message(self): + payload = {"output": [], "status": "failed", "error": {"message": "content was filtered"}} + with pytest.raises(Exception, match="content was filtered") as excinfo: + _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert excinfo.value.status_code == 502 + + def test_incomplete_with_no_results_raises_with_reason(self): + payload = {"output": [], "status": "incomplete", "incomplete_details": {"reason": "max_output_tokens"}} + with pytest.raises(Exception, match="incomplete: max_output_tokens"): + _config().transform_search_response(_resp(payload), logging_obj=Mock()) + + def test_incomplete_with_partial_results_returns_them(self): + payload = _message_response("claim", [_citation("https://example.com", "Example")]) + payload["status"] = "incomplete" + 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): + resp = _config().transform_search_response(_resp(REAL_SHAPED_FIXTURE), logging_obj=Mock()) + assert resp._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0045 + + def test_no_usage_leaves_no_cost_attached(self): + payload = _message_response("claim", [_citation("https://example.com", "Example")]) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert "additional_headers" not in resp._hidden_params + + def test_falls_back_to_token_cost_when_ticks_absent(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": 1e-6, "output_cost_per_token": 2e-6} + raise Exception("not mapped") + + monkeypatch.setattr(transformation, "get_model_info", fake_get_model_info) + payload = _message_response( + "claim", + [_citation("https://example.com", "Example")], + usage={"input_tokens": 100, "output_tokens": 50, "output_tokens_details": {"reasoning_tokens": 0}}, + ) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + expected = 100 * 1e-6 + 50 * 2e-6 + transformation._PER_CITATION_SURCHARGE_USD * 1 + assert resp._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == pytest.approx( + expected + ) + + def test_uses_reasoning_suffix_when_reasoning_tokens_present(self, monkeypatch: pytest.MonkeyPatch): + seen_models = [] + + def fake_get_model_info(model: str, custom_llm_provider: str): + seen_models.append(model) + if model == "xai/grok-4-fast-reasoning": + return {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6} + raise Exception("not mapped") + + monkeypatch.setattr(transformation, "get_model_info", fake_get_model_info) + payload = _message_response( + "claim", + [_citation("https://example.com", "Example")], + usage={"input_tokens": 100, "output_tokens": 50, "output_tokens_details": {"reasoning_tokens": 10}}, + ) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert "additional_headers" in resp._hidden_params + assert seen_models == ["xai/grok-4-fast", "xai/grok-4-fast-reasoning"] + + def test_tries_bare_model_before_suffix(self, monkeypatch: pytest.MonkeyPatch): + def fake_get_model_info(model: str, custom_llm_provider: str): + if model == "xai/grok-4-fast": + return {"input_cost_per_token": 3e-6, "output_cost_per_token": 4e-6} + raise Exception("not mapped") + + monkeypatch.setattr(transformation, "get_model_info", fake_get_model_info) + payload = _message_response( + "claim", + [_citation("https://example.com", "Example")], + usage={"input_tokens": 10, "output_tokens": 5, "output_tokens_details": {"reasoning_tokens": 0}}, + ) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + expected = 10 * 3e-6 + 5 * 4e-6 + transformation._PER_CITATION_SURCHARGE_USD * 1 + assert resp._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == pytest.approx( + expected + ) + + def test_unmapped_model_leaves_no_cost_attached(self, monkeypatch: pytest.MonkeyPatch): + def fake_get_model_info(model: str, custom_llm_provider: str): + raise Exception("not mapped") + + monkeypatch.setattr(transformation, "get_model_info", fake_get_model_info) + payload = _message_response( + "claim", + [_citation("https://example.com", "Example")], + usage={"input_tokens": 10, "output_tokens": 5, "output_tokens_details": {"reasoning_tokens": 0}}, + ) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert "additional_headers" not in resp._hidden_params + + def test_surcharge_uses_distinct_citation_count_not_raw_annotation_count(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) + payload = _message_response( + "claim", + [ + _citation("https://example.com/a", "A"), + _citation("https://example.com/a", "A"), + _citation("https://example.com/b", "B"), + ], + usage={"input_tokens": 0, "output_tokens": 0, "output_tokens_details": {"reasoning_tokens": 0}}, + ) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert resp._hidden_params["additional_headers"][ + "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): + error = _config().get_error_class(error_message="quota exceeded", status_code=429, headers={}) + assert error.status_code == 429 + assert "xAI x_search: quota exceeded" in str(error)