scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge ff474a0b9f into 559247fa84

Browse source
This commit is contained in:
Deepanshu Lulla 2026-09-13 04:04:06 +08:00 • committed by GitHub
commit 0bbc62f728
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 862 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/llms/xai/search/__init__.py Normal file
View file

@ -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

382
litellm/llms/xai/search/transformation.py Normal file
View file

@ -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,
)

1
litellm/types/utils.py
View file

@ -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

2
litellm/utils.py
View file

@ -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:

470
tests/test_litellm/llms/xai/search/test_xai_search_transformation.py Normal file
View file

@ -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",
[
"<html>502 Bad Gateway</html>",
'{"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)