from collections.abc import AsyncIterator, Iterator, Mapping from typing import Any, Final import httpx import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, ) from litellm.llms.xai.common_utils import XAIModelInfo from litellm.llms.xai.cost_calculator import ( apply_server_side_tool_usage_details_to_usage, ) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( Choices, ModelResponse, ModelResponseStream, Usage, ) from ...openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) class XAIChatConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: return "xai" def _get_openai_compatible_provider_info( self, api_base: str | None, api_key: str | None ) -> tuple[str | None, str | None]: api_base = api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) return api_base, dynamic_api_key def validate_environment( self, headers: dict, model: str, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: str | None = None, api_base: str | None = None, ) -> dict: from litellm.llms.xai.oauth import ( XAIOAuthAuthenticator, XAIOAuthError, should_use_xai_oauth, ) dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: try: headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}" except XAIOAuthError as exc: raise AuthenticationError( model=model, llm_provider=self.custom_llm_provider or "xai", message=str(exc), ) from exc if "content-type" not in headers and "Content-Type" not in headers: headers["Content-Type"] = "application/json" return headers return super().validate_environment( headers=headers, model=model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, api_key=dynamic_api_key, api_base=api_base, ) def get_complete_url( self, api_base: str | None, api_key: str | None, model: str, optional_params: dict, litellm_params: dict, stream: bool | None = None, ) -> str: from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth dynamic_api_key: Final = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: api_base = XAIOAuthAuthenticator().get_api_base() return super().get_complete_url( api_base=api_base, api_key=dynamic_api_key, model=model, optional_params=optional_params, litellm_params=litellm_params, stream=stream, ) def get_supported_openai_params(self, model: str) -> list: base_openai_params: Final = [ "logit_bias", "logprobs", "max_tokens", "n", "parallel_tool_calls", "presence_penalty", "response_format", "seed", "stream", "stream_options", "temperature", "tool_choice", "tools", "top_logprobs", "top_p", "user", "web_search_options", ] # for some reason, grok-3-mini does not support stop tokens ######################################################### # stop tokens check ######################################################### if self._supports_stop_reason(model): base_openai_params.append("stop") ######################################################### # frequency penalty check ######################################################### if self._supports_frequency_penalty(model): base_openai_params.append("frequency_penalty") ######################################################### # reasoning check ######################################################### try: if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_openai_params.append("reasoning_effort") except Exception as e: verbose_logger.debug("Error checking if model supports reasoning: %s", e) return base_openai_params def _supports_stop_reason(self, model: str) -> bool: if "grok-3-mini" in model or "grok-4" in model or "grok-code-fast" in model: return False return True def _supports_frequency_penalty(self, model: str) -> bool: """ From manual testing grok-4 does not support `frequency_penalty` When sent the model fails from xAI API """ if "grok-4" in model: return False if "grok-code-fast" in model: return False return True def map_openai_params( self, non_default_params: dict, optional_params: dict, model: str, drop_params: bool = False, ) -> dict: supported_openai_params: Final = self.get_supported_openai_params(model=model) for param, value in non_default_params.items(): if param == "max_completion_tokens": optional_params["max_tokens"] = value elif param == "tools" and value is not None: tools = [] for tool in value: tool = filter_value_from_dict(tool, "strict") if tool is not None: tools.append(tool) if len(tools) > 0: optional_params["tools"] = tools elif param in supported_openai_params: if value is not None: optional_params[param] = value return optional_params def get_model_response_iterator( self, streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, ) -> Any: return XAIChatCompletionStreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, ) def transform_request( self, model: str, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, ) -> dict: """ Handle https://github.com/BerriAI/litellm/issues/9720 Filter out 'name' from messages """ messages = strip_name_from_messages(messages) return super().transform_request(model, messages, optional_params, litellm_params, headers) @staticmethod def _fix_choice_finish_reason_for_tool_calls(choice: Choices) -> None: """ Helper to fix finish_reason for tool calls when XAI API returns empty string. XAI API returns empty string for finish_reason when using tools, so we need to set it to "tool_calls" when tool_calls are present. """ if choice.finish_reason == "" and choice.message.tool_calls and len(choice.message.tool_calls) > 0: choice.finish_reason = "tool_calls" def transform_response( self, model: str, raw_response: httpx.Response, model_response: ModelResponse, logging_obj, request_data: dict, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, encoding, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: """ Transform the response from the XAI API. XAI API returns empty string for finish_reason when using tools, so we need to fix this after the standard OpenAI transformation. Also handles X.AI web search usage tracking. """ # First, let the parent class handle the standard transformation response: Final = super().transform_response( model=model, raw_response=raw_response, model_response=model_response, logging_obj=logging_obj, request_data=request_data, messages=messages, optional_params=optional_params, litellm_params=litellm_params, encoding=encoding, api_key=api_key, json_mode=json_mode, ) # Fix finish_reason for tool calls across all choices if response.choices: for choice in response.choices: if isinstance(choice, Choices): self._fix_choice_finish_reason_for_tool_calls(choice) # Handle X.AI web search usage tracking try: raw_response_json: Final = raw_response.json() self._enhance_usage_with_xai_web_search_fields(response, raw_response_json) except Exception as e: verbose_logger.debug("Error extracting X.AI web search usage: %s", e) self._fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) return response @staticmethod def _fold_reasoning_tokens_into_completion( target: ModelResponse | Usage | dict[str, Any] | None, ) -> None: """Reconcile xAI Usage to the OpenAI invariant. xAI accounts ``reasoning_tokens`` separately from ``completion_tokens`` while still summing them into ``total_tokens``. OpenAI's contract (o1/o3) folds reasoning into ``completion_tokens``, so fold here to keep ``total = prompt + completion``. Idempotent. Accepts a ``ModelResponse`` (non-streaming), a ``Usage`` object, or a raw usage ``dict`` (streaming chunk) so streaming and non-streaming paths stay in sync. """ if target is None: return if isinstance(target, ModelResponse): usage: Usage | dict[str, Any] | None = getattr(target, "usage", None) else: usage = target if usage is None: return if isinstance(usage, dict): details = usage.get("completion_tokens_details") or {} if isinstance(details, dict): reasoning_tokens = int(details.get("reasoning_tokens") or 0) else: reasoning_tokens = int(getattr(details, "reasoning_tokens", 0) or 0) if reasoning_tokens <= 0: return prompt_tokens = int(usage.get("prompt_tokens") or 0) completion_tokens = int(usage.get("completion_tokens") or 0) total_tokens = int(usage.get("total_tokens") or 0) if total_tokens == prompt_tokens + completion_tokens: return # Guard against double-counting if xAI changes accounting. if total_tokens != prompt_tokens + completion_tokens + reasoning_tokens: return usage["completion_tokens"] = completion_tokens + reasoning_tokens return details = getattr(usage, "completion_tokens_details", None) reasoning_tokens = int(getattr(details, "reasoning_tokens", 0) or 0) if details else 0 if reasoning_tokens <= 0: return prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0) completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) total_tokens = int(getattr(usage, "total_tokens", 0) or 0) if total_tokens == prompt_tokens + completion_tokens: return # Guard against double-counting if xAI changes accounting. if total_tokens != prompt_tokens + completion_tokens + reasoning_tokens: return usage.completion_tokens = completion_tokens + reasoning_tokens def _enhance_usage_with_xai_web_search_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ Copy usage.server_side_tool_usage_details from the provider usage block onto model_response.usage for tool cost calculation. """ if not hasattr(model_response, "usage") or model_response.usage is None: return usage: Final[Usage] = model_response.usage response_usage: Final = raw_response_json.get("usage") if not isinstance(response_usage, dict): return details: Final = response_usage.get("server_side_tool_usage_details") if isinstance(details, Mapping): apply_server_side_tool_usage_details_to_usage(usage, details) verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details) @staticmethod def _normalize_openai_compatible_usage_totals( usage: Usage | dict[str, Any] | None, ) -> None: if usage is None: return if isinstance(usage, dict): prompt_tokens = int(usage.get("prompt_tokens") or 0) completion_tokens = int(usage.get("completion_tokens") or 0) expected_total = prompt_tokens + completion_tokens if int(usage.get("total_tokens") or 0) < expected_total: usage["total_tokens"] = expected_total return prompt_tokens = int(usage.prompt_tokens or 0) completion_tokens = int(usage.completion_tokens or 0) expected_total = prompt_tokens + completion_tokens if int(usage.total_tokens or 0) < expected_total: usage.total_tokens = expected_total class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): def chunk_parser(self, chunk: dict) -> ModelResponseStream: """ Handle xAI-specific streaming behavior. xAI Grok sends a final chunk with empty choices array but with usage data when stream_options={"include_usage": True} is set. Example from xAI API: {"id":"...","object":"chat.completion.chunk","created":...,"model":"grok-4-1-fast-non-reasoning", "choices":[],"usage":{"prompt_tokens":171,"completion_tokens":2,"total_tokens":173,...}} """ # Handle chunks with empty choices but with usage data choices: Final = chunk.get("choices", []) if len(choices) == 0 and "usage" in chunk: # xAI sends usage in a chunk with empty choices array # Add a dummy choice with empty delta to ensure proper processing chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] if "usage" in chunk and chunk["usage"] is not None: XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) return super().chunk_parser(chunk)