pydantic_ai.models.xai

Setup

For details on how to set up authentication with this model, see model configuration for xAI.

xAI model implementation using xAI SDK.

XAI_EFFORT_MAP module-attribute

XAI_EFFORT_MAP: dict[
    ThinkingLevel, Literal["low", "high"]
] = {
    True: "high",
    "minimal": "low",
    "low": "low",
    "medium": "high",
    "high": "high",
    "xhigh": "high",
}

Maps unified thinking values to xAI reasoning_effort. xAI only supports 'low' and 'high'.

XaiModelName module-attribute

XaiModelName = str | ChatModel

Possible xAI model names.

XaiModelSettings

Bases: ModelSettings

Settings specific to xAI models.

See xAI SDK documentation for more details on these parameters.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

class XaiModelSettings(ModelSettings, total=False):
    """Settings specific to xAI models.

    See [xAI SDK documentation](https://docs.x.ai/docs) for more details on these parameters.
    """

    xai_logprobs: bool
    """Whether to return log probabilities of the output tokens or not."""

    xai_top_logprobs: int
    """An integer between 0 and 20 specifying the number of most likely tokens to return at each position."""

    xai_user: str
    """A unique identifier representing your end-user, which can help xAI to monitor and detect abuse."""

    xai_store_messages: bool
    """Whether to store messages on xAI's servers for conversation continuity."""

    xai_previous_response_id: str
    """The ID of the previous response to continue the conversation."""

    xai_include_encrypted_content: bool
    """Whether to include the encrypted content in the response.

    Corresponds to the `use_encrypted_content` value of the model settings in the Responses API.
    """

    xai_include_code_execution_output: bool
    """Whether to include the code execution results in the response.

    Corresponds to the `code_interpreter_call.outputs` value of the `include` parameter in the Responses API.
    """

    xai_include_web_search_output: bool
    """Whether to include the web search results in the response.

    Corresponds to the `web_search_call.action.sources` value of the `include` parameter in the Responses API.
    """

    xai_include_inline_citations: bool
    """Whether to include inline citations in the response.

    Corresponds to the `inline_citations` option in the xAI `include` parameter.
    """

    xai_include_mcp_output: bool
    """Whether to include the MCP results in the response.

    Corresponds to the `mcp_call.outputs` value of the `include` parameter in the Responses API.
    """

    xai_include_x_search_output: bool
    """Whether to include the X search results in the response.

    Corresponds to the `x_search_call.outputs` value of the `include` parameter in the Responses API.
    """

    xai_include_collections_search_output: bool
    """Whether to include the collections search results in the response.

    Corresponds to the `collections_search_call.outputs` value of the `include` parameter in the Responses API.
    """

    xai_reasoning_effort: Literal['low', 'high']
    """Reasoning effort level for Grok reasoning models.

    See https://docs.x.ai for details.
    """

xai_logprobs instance-attribute

xai_logprobs: bool

Whether to return log probabilities of the output tokens or not.

xai_top_logprobs instance-attribute

xai_top_logprobs: int

An integer between 0 and 20 specifying the number of most likely tokens to return at each position.

xai_user instance-attribute

xai_user: str

A unique identifier representing your end-user, which can help xAI to monitor and detect abuse.

xai_store_messages instance-attribute

xai_store_messages: bool

Whether to store messages on xAI's servers for conversation continuity.

xai_previous_response_id instance-attribute

xai_previous_response_id: str

The ID of the previous response to continue the conversation.

xai_include_encrypted_content instance-attribute

xai_include_encrypted_content: bool

Whether to include the encrypted content in the response.

Corresponds to the use_encrypted_content value of the model settings in the Responses API.

xai_include_code_execution_output instance-attribute

xai_include_code_execution_output: bool

Whether to include the code execution results in the response.

Corresponds to the code_interpreter_call.outputs value of the include parameter in the Responses API.

xai_include_web_search_output instance-attribute

xai_include_web_search_output: bool

Whether to include the web search results in the response.

Corresponds to the web_search_call.action.sources value of the include parameter in the Responses API.

xai_include_inline_citations instance-attribute

xai_include_inline_citations: bool

Whether to include inline citations in the response.

Corresponds to the inline_citations option in the xAI include parameter.

xai_include_mcp_output instance-attribute

xai_include_mcp_output: bool

Whether to include the MCP results in the response.

Corresponds to the mcp_call.outputs value of the include parameter in the Responses API.

xai_include_x_search_output instance-attribute

xai_include_x_search_output: bool

Whether to include the X search results in the response.

Corresponds to the x_search_call.outputs value of the include parameter in the Responses API.

xai_include_collections_search_output instance-attribute

xai_include_collections_search_output: bool

Whether to include the collections search results in the response.

Corresponds to the collections_search_call.outputs value of the include parameter in the Responses API.

xai_reasoning_effort instance-attribute

xai_reasoning_effort: Literal['low', 'high']

Reasoning effort level for Grok reasoning models.

See https://docs.x.ai for details.

XSearch dataclass

Bases: NativeOrLocalTool[AgentDepsT]

X (Twitter) search capability for xAI models.

Uses the xAI model's native x_search builtin tool. Only works with xAI models.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

@dataclass(init=False)
class XSearch(NativeOrLocalTool[AgentDepsT]):
    """X (Twitter) search capability for xAI models.

    Uses the xAI model's native x_search builtin tool. Only works with xAI models.
    """

    allowed_x_handles: list[str] | None
    """If provided, only posts from these X handles will be included (max 10). Requires builtin support."""

    excluded_x_handles: list[str] | None
    """If provided, posts from these X handles will be excluded (max 10). Requires builtin support."""

    from_date: datetime | None
    """If provided, only posts created on or after this datetime will be included."""

    to_date: datetime | None
    """If provided, only posts created on or before this datetime will be included."""

    enable_image_understanding: bool
    """Enable image analysis from X posts. Defaults to `False`."""

    enable_video_understanding: bool
    """Enable video analysis from X content. Defaults to `False`."""

    include_output: bool
    """Include raw X search results in the response as
    [`NativeToolReturnPart`][pydantic_ai.messages.NativeToolReturnPart]. Defaults to `False`.
    """

    def __init__(
        self,
        *,
        native: XSearchTool
        | Callable[[RunContext[AgentDepsT]], Awaitable[XSearchTool | None] | XSearchTool | None]
        | bool = True,
        local: Tool[AgentDepsT] | Callable[..., Any] | Literal[False] | None = None,
        allowed_x_handles: list[str] | None = None,
        excluded_x_handles: list[str] | None = None,
        from_date: datetime | None = None,
        to_date: datetime | None = None,
        enable_image_understanding: bool = False,
        enable_video_understanding: bool = False,
        include_output: bool = False,
    ) -> None:
        self.native = native
        self.local = local
        self.allowed_x_handles = allowed_x_handles
        self.excluded_x_handles = excluded_x_handles
        self.from_date = from_date
        self.to_date = to_date
        self.enable_image_understanding = enable_image_understanding
        self.enable_video_understanding = enable_video_understanding
        self.include_output = include_output
        self.__post_init__()

    def _default_native(self) -> XSearchTool:
        return XSearchTool(
            allowed_x_handles=self.allowed_x_handles,
            excluded_x_handles=self.excluded_x_handles,
            from_date=self.from_date,
            to_date=self.to_date,
            enable_image_understanding=self.enable_image_understanding,
            enable_video_understanding=self.enable_video_understanding,
            include_output=self.include_output,
        )

    def _native_unique_id(self) -> str:
        return XSearchTool.kind

    def _default_local(self) -> Tool[AgentDepsT] | AbstractToolset[AgentDepsT] | None:
        return None

    def _requires_native(self) -> bool:
        return self.allowed_x_handles is not None or self.excluded_x_handles is not None

allowed_x_handles instance-attribute

allowed_x_handles: list[str] | None = allowed_x_handles

If provided, only posts from these X handles will be included (max 10). Requires builtin support.

excluded_x_handles instance-attribute

excluded_x_handles: list[str] | None = excluded_x_handles

If provided, posts from these X handles will be excluded (max 10). Requires builtin support.

from_date instance-attribute

from_date: datetime | None = from_date

If provided, only posts created on or after this datetime will be included.

to_date instance-attribute

to_date: datetime | None = to_date

If provided, only posts created on or before this datetime will be included.

enable_image_understanding instance-attribute

enable_image_understanding: bool = (
    enable_image_understanding
)

Enable image analysis from X posts. Defaults to False.

enable_video_understanding instance-attribute

enable_video_understanding: bool = (
    enable_video_understanding
)

Enable video analysis from X content. Defaults to False.

include_output instance-attribute

include_output: bool = include_output

Include raw X search results in the response as NativeToolReturnPart. Defaults to False.

XaiModel

Bases: Model[AsyncClient]

A model that uses the xAI SDK to interact with xAI models.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

class XaiModel(Model[AsyncClient]):
    """A model that uses the xAI SDK to interact with xAI models."""

    _model_name: str
    _provider: Provider[AsyncClient]

    def __init__(
        self,
        model_name: XaiModelName,
        *,
        provider: Literal['xai'] | Provider[AsyncClient] = 'xai',
        profile: ModelProfileSpec | None = None,
        settings: ModelSettings | None = None,
    ):
        """Initialize the xAI model.

        Args:
            model_name: The name of the xAI model to use (e.g., "grok-4-1-fast-non-reasoning")
            provider: The provider to use for API calls. Defaults to `'xai'`.
            profile: Optional model profile specification. Defaults to a profile picked by the provider based on the model name.
            settings: Optional model settings.
        """
        self._model_name = model_name

        if isinstance(provider, str):
            provider = infer_provider(provider)
        self._provider = provider

        super().__init__(settings=settings, profile=profile)

    @property
    def client(self) -> 'AsyncClient':
        return self._provider.client

    @property
    def model_name(self) -> str:
        """The model name."""
        return self._model_name

    @property
    def system(self) -> str:
        """The model provider."""
        return 'xai'

    @cached_property
    def profile(self) -> GrokModelProfile:
        return cast(GrokModelProfile, super().profile)

    @classmethod
    def supported_native_tools(cls) -> frozenset[type]:
        """Return the set of builtin tool types this model can handle."""
        return frozenset({WebSearchTool, CodeExecutionTool, MCPServerTool, XSearchTool, FileSearchTool})

    async def _map_messages(
        self,
        messages: list[ModelMessage],
        model_request_parameters: ModelRequestParameters,
    ) -> list[chat_types.chat_pb2.Message]:
        """Convert pydantic_ai messages to xAI SDK messages."""
        xai_messages: list[chat_types.chat_pb2.Message] = []
        # xAI expects tool results in the same order as tool calls.
        #
        # Pydantic AI doesn't guarantee tool-result part ordering, so we track
        # tool call order as we walk message history and reorder tool results.
        pending_tool_call_ids: list[str] = []

        for message in messages:
            if isinstance(message, ModelRequest):
                mapped_request_parts = await self._map_request_parts(
                    message.parts,
                    pending_tool_call_ids,
                )
                xai_messages.extend(mapped_request_parts)
            elif isinstance(message, ModelResponse):
                xai_messages.extend(self._map_response_parts(message.parts))
                pending_tool_call_ids.extend(
                    part.tool_call_id for part in message.parts if isinstance(part, ToolCallPart) and part.tool_call_id
                )
            else:
                assert_never(message)

        # Insert instructions as system messages after existing system messages if present
        if instruction_parts := self._get_instruction_parts(messages, model_request_parameters):
            system_prompt_count = next(
                (i for i, m in enumerate(xai_messages) if m.role != chat_types.chat_pb2.MessageRole.ROLE_SYSTEM),
                len(xai_messages),
            )
            xai_messages[system_prompt_count:system_prompt_count] = [system(part.content) for part in instruction_parts]

        return xai_messages

    async def _map_request_parts(
        self,
        parts: Sequence[ModelRequestPart],
        pending_tool_call_ids: list[str],
    ) -> list[chat_types.chat_pb2.Message]:
        """Map ModelRequest parts to xAI messages."""
        xai_messages: list[chat_types.chat_pb2.Message] = []
        tool_results: list[ToolReturnPart | RetryPromptPart] = []

        for part in parts:
            if isinstance(part, SystemPromptPart):
                xai_messages.append(system(part.content))
            elif isinstance(part, UserPromptPart):
                if user_msg := await self._map_user_prompt(part):
                    xai_messages.append(user_msg)
            elif isinstance(part, ToolReturnPart):
                tool_results.append(part)
            elif isinstance(part, RetryPromptPart):
                if part.tool_name is None:
                    xai_messages.append(user(part.model_response()))
                else:
                    tool_results.append(part)
            else:
                assert_never(part)

        # Sort tool results by requested order, then emit
        if tool_results:
            order = {id: i for i, id in enumerate(pending_tool_call_ids)}
            tool_results.sort(key=lambda p: order.get(p.tool_call_id, float('inf')))
            file_content: list[UserContent] = []
            for part in tool_results:
                if isinstance(part, ToolReturnPart):
                    text, files = part.model_response_str_and_user_content()
                    xai_messages.append(tool_result(text))
                    file_content.extend(files)
                else:
                    xai_messages.append(tool_result(part.model_response()))
            if file_content and (
                user_msg := await self._map_user_prompt(UserPromptPart(content=file_content))
            ):  # pragma: no branch
                xai_messages.append(user_msg)

        return xai_messages

    def _map_response_parts(self, parts: Sequence[ModelResponsePart]) -> list[chat_types.chat_pb2.Message]:
        """Map ModelResponse parts to xAI assistant messages (one message per part)."""
        messages: list[chat_types.chat_pb2.Message] = []

        # Track builtin tool calls by tool_call_id to update their status with return parts
        builtin_calls: dict[str, chat_types.chat_pb2.ToolCall] = {}

        for item in parts:
            if isinstance(item, TextPart):
                messages.append(assistant(item.content))
            elif isinstance(item, ThinkingPart):
                if (thinking_msg := self._map_thinking_part(item)) is not None:
                    messages.append(thinking_msg)
            elif isinstance(item, ToolCallPart):
                client_side_tool_call = self._map_tool_call(item)
                self._append_tool_call(messages, client_side_tool_call)
            elif isinstance(item, NativeToolCallPart):
                builtin_call = self._map_builtin_tool_call_part(item)
                if item.provider_name in _XAI_PROVIDER_NAMES and builtin_call:
                    self._append_tool_call(messages, builtin_call)
                    # Track specific tool calls for status updates
                    # Note: tool_call_id is always truthy here since _map_builtin_tool_call_part
                    # returns None when tool_call_id is empty
                    if item.tool_call_id:  # pragma: no branch
                        builtin_calls[item.tool_call_id] = builtin_call
            elif isinstance(item, NativeToolReturnPart):
                if (
                    item.provider_name in _XAI_PROVIDER_NAMES
                    and item.tool_call_id
                    and (details := item.provider_details) is not None
                    and details.get('status') == 'failed'
                    and (call := builtin_calls.get(item.tool_call_id))
                ):
                    call.status = chat_types.chat_pb2.TOOL_CALL_STATUS_FAILED
                    if error_msg := details.get('error'):
                        call.error_message = str(error_msg)
            elif isinstance(item, FilePart):
                # Files generated by models (e.g., from CodeExecutionTool) are not sent back
                pass
            elif isinstance(item, CompactionPart):  # pragma: no cover
                # Compaction parts are not sent back to models that don't support compaction.
                pass
            else:
                assert_never(item)

        return messages

    @staticmethod
    def _append_tool_call(messages: list[chat_types.chat_pb2.Message], tool_call: chat_types.chat_pb2.ToolCall) -> None:
        """Append a tool call to the most recent tool-call assistant message, or create a new one.

        We keep tool calls grouped to avoid generating one assistant message per tool call.
        """
        if messages and messages[-1].tool_calls:
            messages[-1].tool_calls.append(tool_call)
        else:
            msg = assistant('')
            msg.tool_calls.append(tool_call)
            messages.append(msg)

    def _map_thinking_part(self, item: ThinkingPart) -> chat_types.chat_pb2.Message | None:
        """Map a `ThinkingPart` into a single xAI assistant message.

        - Native xAI thinking (with optional signature) is sent via `reasoning_content`/`encrypted_content`
        - Non-xAI (or non-native) thinking is preserved by wrapping in the model profile's thinking tags
        """
        if item.provider_name in _XAI_PROVIDER_NAMES and (item.content or item.signature):
            msg = assistant('')
            if item.content:
                msg.reasoning_content = item.content
            if item.signature:
                msg.encrypted_content = item.signature
            return msg
        elif item.content:
            start_tag, end_tag = self.profile.get('thinking_tags', DEFAULT_THINKING_TAGS)
            return assistant('\n'.join([start_tag, item.content, end_tag]))
        else:
            return None

    def _map_tool_call(self, tool_call_part: ToolCallPart) -> chat_types.chat_pb2.ToolCall:
        """Map a ToolCallPart to an xAI SDK ToolCall."""
        return chat_types.chat_pb2.ToolCall(
            id=tool_call_part.tool_call_id,
            type=chat_types.chat_pb2.TOOL_CALL_TYPE_CLIENT_SIDE_TOOL,
            status=chat_types.chat_pb2.TOOL_CALL_STATUS_COMPLETED,
            function=chat_types.chat_pb2.FunctionCall(
                name=tool_call_part.tool_name,
                arguments=tool_call_part.args_as_json_str(),
            ),
        )

    def _map_builtin_tool_call_part(self, item: NativeToolCallPart) -> chat_types.chat_pb2.ToolCall | None:
        """Map a NativeToolCallPart to an xAI SDK ToolCall with appropriate type and status."""
        if not item.tool_call_id:
            return None

        if item.tool_name == CodeExecutionTool.kind:
            return chat_types.chat_pb2.ToolCall(
                id=item.tool_call_id,
                type=chat_types.chat_pb2.TOOL_CALL_TYPE_CODE_EXECUTION_TOOL,
                status=chat_types.chat_pb2.TOOL_CALL_STATUS_COMPLETED,
                function=chat_types.chat_pb2.FunctionCall(
                    name=CodeExecutionTool.kind,
                    arguments=item.args_as_json_str(),
                ),
            )
        elif item.tool_name == WebSearchTool.kind:
            return chat_types.chat_pb2.ToolCall(
                id=item.tool_call_id,
                type=chat_types.chat_pb2.TOOL_CALL_TYPE_WEB_SEARCH_TOOL,
                status=chat_types.chat_pb2.TOOL_CALL_STATUS_COMPLETED,
                function=chat_types.chat_pb2.FunctionCall(
                    name=WebSearchTool.kind,
                    arguments=item.args_as_json_str(),
                ),
            )
        elif item.tool_name == XSearchTool.kind:
            function_name = (item.provider_details or {}).get('function_name', XSearchTool.kind)
            return chat_types.chat_pb2.ToolCall(
                id=item.tool_call_id,
                type=chat_types.chat_pb2.TOOL_CALL_TYPE_X_SEARCH_TOOL,
                status=chat_types.chat_pb2.TOOL_CALL_STATUS_COMPLETED,
                function=chat_types.chat_pb2.FunctionCall(
                    name=function_name,
                    arguments=item.args_as_json_str(),
                ),
            )
        elif item.tool_name.startswith(MCPServerTool.kind):
            # Extract server label from tool_name (format: 'mcp_server:server_label')
            server_label = item.tool_name.split(':', 1)[1] if ':' in item.tool_name else item.tool_name
            args_dict = item.args_as_dict() or {}
            # Extract tool_name and tool_args from the structured args (matches OpenAI/Anthropic pattern)
            actual_tool_name = args_dict.get('tool_name', '')
            tool_args = args_dict.get('tool_args', {})
            # Construct the full function name in xAI's format: 'server_label.tool_name'
            function_name = f'{server_label}.{actual_tool_name}' if actual_tool_name else server_label
            return chat_types.chat_pb2.ToolCall(
                id=item.tool_call_id,
                type=chat_types.chat_pb2.TOOL_CALL_TYPE_MCP_TOOL,
                status=chat_types.chat_pb2.TOOL_CALL_STATUS_COMPLETED,
                function=chat_types.chat_pb2.FunctionCall(
                    name=function_name,
                    arguments=json.dumps(tool_args),
                ),
            )
        elif item.tool_name == FileSearchTool.kind:
            function_name = (item.provider_details or {}).get('function_name', FileSearchTool.kind)
            return chat_types.chat_pb2.ToolCall(
                id=item.tool_call_id,
                type=chat_types.chat_pb2.TOOL_CALL_TYPE_COLLECTIONS_SEARCH_TOOL,
                status=chat_types.chat_pb2.TOOL_CALL_STATUS_COMPLETED,
                function=chat_types.chat_pb2.FunctionCall(
                    name=function_name,
                    arguments=item.args_as_json_str(),
                ),
            )
        return None

    async def _upload_file_to_xai(self, data: bytes, filename: str) -> str:
        """Upload a file to xAI files API and return the file ID.

        Args:
            data: The file content as bytes
            filename: The filename to use for the upload

        Returns:
            The file ID from xAI
        """
        uploaded_file = await self._provider.client.files.upload(data, filename=filename)
        return uploaded_file.id

    async def _map_user_prompt(self, part: UserPromptPart) -> chat_types.chat_pb2.Message | None:  # noqa: C901
        """Map a UserPromptPart to an xAI user message."""
        if isinstance(part.content, str):
            return user(part.content)

        # Handle complex content (images, text, etc.)
        content_items: list[chat_types.Content] = []

        for item in part.content:
            if isinstance(item, str | TextContent):
                text = item if isinstance(item, str) else item.content
                content_items.append(text)
            elif isinstance(item, ImageUrl):
                # Get detail from vendor_metadata if available
                detail: chat_types.ImageDetail = 'auto'
                if item.vendor_metadata and 'detail' in item.vendor_metadata:
                    detail = item.vendor_metadata['detail']
                image_url = item.url
                if item.force_download:
                    downloaded = await download_item(item, data_format='base64_uri', type_format='extension')
                    image_url = downloaded['data']
                content_items.append(image(image_url, detail=detail))
            elif isinstance(item, BinaryContent):
                if item.is_image:
                    # Convert binary content to data URI and use image()
                    image_detail: chat_types.ImageDetail = 'auto'
                    if item.vendor_metadata and 'detail' in item.vendor_metadata:
                        image_detail = item.vendor_metadata['detail']
                    content_items.append(image(item.data_uri, detail=image_detail))
                elif item.is_audio:
                    raise NotImplementedError('BinaryContent with audio is not supported in xAI user prompts')
                elif item.is_document:
                    # Upload document to xAI files API and reference it
                    filename = item.identifier or f'document.{item.format}'
                    file_id = await self._upload_file_to_xai(item.data, filename)
                    content_items.append(file(file_id))
                elif item.is_video:
                    raise NotImplementedError('BinaryContent with video is not supported in xAI user prompts')
                else:  # pragma: no cover
                    raise RuntimeError(f'Unsupported binary content type: {item.media_type}')
            elif isinstance(item, AudioUrl):
                raise NotImplementedError('AudioUrl is not supported in xAI user prompts')
            elif isinstance(item, DocumentUrl):
                # Download and upload to xAI files API
                downloaded = await download_item(item, data_format='bytes')
                filename = item.identifier or 'document'
                # Add extension if data_type is available from download
                if 'data_type' in downloaded and downloaded['data_type']:
                    filename = f'{filename}.{downloaded["data_type"]}'

                file_id = await self._upload_file_to_xai(downloaded['data'], filename)
                content_items.append(file(file_id))
            elif isinstance(item, VideoUrl):
                raise NotImplementedError('VideoUrl is not supported in xAI user prompts')
            elif isinstance(item, UploadedFile):
                if item.provider_name != self.system:
                    raise UserError(
                        f'UploadedFile with `provider_name={item.provider_name!r}` cannot be used with XaiModel. '
                        f'Expected `provider_name` to be `{self.system!r}`.'
                    )
                content_items.append(file(item.file_id))
            elif isinstance(item, CachePoint):
                # xAI doesn't support prompt caching via CachePoint, so we filter it out
                pass
            else:
                assert_never(item)

        if content_items:
            return user(*content_items)

        return None

    def _get_tool_choice(
        self,
        model_settings: XaiModelSettings,
        model_request_parameters: ModelRequestParameters,
    ) -> tuple[dict[str, ToolDefinition], Literal['none', 'required', 'auto'] | chat_pb2.ToolChoice | None]:
        """Determine which tools to send and the API tool_choice value.

        Returns:
            A tuple of (filtered_tool_defs, tool_choice).
        """
        resolved_tool_choice = resolve_tool_choice(model_settings, model_request_parameters)
        tool_defs = model_request_parameters.tool_defs

        profile = self.profile

        tool_choice: Literal['none', 'required', 'auto'] | chat_pb2.ToolChoice
        if resolved_tool_choice in ('auto', 'none'):
            tool_choice = resolved_tool_choice
        elif resolved_tool_choice == 'required':
            tool_choice = 'required' if profile.get('grok_supports_tool_choice_required', True) else 'auto'
        elif isinstance(resolved_tool_choice, tuple):
            tool_choice_mode, tool_names = resolved_tool_choice
            if tool_choice_mode == 'required' and len(tool_names) == 1:
                if profile.get('grok_supports_tool_choice_required', True):
                    tool_choice = required_tool(next(iter(tool_names)))
                else:
                    # Forcing not supported: filter so the model can only see the requested tool.
                    # Breaks caching, but xAI doesn't support limiting tools via API arg.
                    tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names}
                    tool_choice = 'auto'
            else:
                tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names}
                if tool_choice_mode == 'required' and profile.get('grok_supports_tool_choice_required', True):
                    tool_choice = 'required'
                else:
                    tool_choice = 'auto'
        else:
            assert_never(resolved_tool_choice)

        if not tool_defs:
            return tool_defs, None

        return tool_defs, tool_choice

    async def _create_chat(
        self,
        messages: list[ModelMessage],
        model_settings: XaiModelSettings,
        model_request_parameters: ModelRequestParameters,
    ) -> Any:
        """Create an xAI chat instance with common setup for both request and stream.

        Returns:
            The xAI SDK chat object, ready to call .sample() or .stream() on.
        """
        # Convert messages to xAI format
        xai_messages = await self._map_messages(messages, model_request_parameters)

        # Resolve tool_choice and filter function tools
        filtered_tool_defs, resolved_tool_choice = self._get_tool_choice(model_settings, model_request_parameters)

        # Convert tools: combine built-in (server-side) tools and custom (client-side) tools
        tools: list[chat_types.chat_pb2.Tool] = []
        if model_request_parameters.native_tools:
            tools.extend(_get_native_tools(model_request_parameters))
        if filtered_tool_defs:
            tools.extend(_map_tools(filtered_tool_defs))
        tools_param = tools if tools else None

        # Determine final tool_choice
        tool_choice: Literal['none', 'required', 'auto'] | chat_pb2.ToolChoice | None
        if resolved_tool_choice is not None:
            tool_choice = resolved_tool_choice
        elif tools_param:
            tool_choice = 'auto'
        else:
            tool_choice = None

        # Set response_format based on the output_mode
        profile = self.profile
        response_format: chat_pb2.ResponseFormat | None = None
        if model_request_parameters.output_mode == 'native':
            output_object = model_request_parameters.output_object
            assert output_object is not None
            response_format = _map_json_schema(output_object)
        elif (
            model_request_parameters.output_mode == 'prompted'
            and not tools_param
            and profile.get('supports_json_object_output', False)
        ):  # pragma: no branch
            response_format = _map_json_object()

        # Map model settings to xAI SDK parameters
        xai_settings = _map_model_settings(model_settings)

        # Fall back to unified thinking when xai_reasoning_effort is not set
        if 'reasoning_effort' not in xai_settings and model_request_parameters.thinking is not None:
            thinking = model_request_parameters.thinking
            if thinking is not False:
                xai_settings['reasoning_effort'] = XAI_EFFORT_MAP[thinking]

        # Populate use_encrypted_content and include based on model settings
        include: list[chat_pb2.IncludeOption] = []
        use_encrypted_content = model_settings.get('xai_include_encrypted_content') or False
        if model_settings.get('xai_include_code_execution_output'):
            include.append(chat_pb2.IncludeOption.INCLUDE_OPTION_CODE_EXECUTION_CALL_OUTPUT)
        if model_settings.get('xai_include_web_search_output'):
            include.append(chat_pb2.IncludeOption.INCLUDE_OPTION_WEB_SEARCH_CALL_OUTPUT)
        if model_settings.get('xai_include_inline_citations'):
            include.append(chat_pb2.IncludeOption.INCLUDE_OPTION_INLINE_CITATIONS)
        if model_settings.get('xai_include_x_search_output') or any(
            isinstance(bt, XSearchTool) and bt.include_output for bt in model_request_parameters.native_tools
        ):
            include.append(chat_pb2.IncludeOption.INCLUDE_OPTION_X_SEARCH_CALL_OUTPUT)
        if model_settings.get('xai_include_collections_search_output'):
            include.append(chat_pb2.IncludeOption.INCLUDE_OPTION_COLLECTIONS_SEARCH_CALL_OUTPUT)
        if model_settings.get('xai_include_mcp_output'):
            include.append(chat_pb2.IncludeOption.INCLUDE_OPTION_MCP_CALL_OUTPUT)

        # Create and return chat instance
        return self._provider.client.chat.create(
            model=self._model_name,
            messages=xai_messages,
            tools=tools_param,
            tool_choice=tool_choice,
            response_format=response_format,
            use_encrypted_content=use_encrypted_content,
            include=include,
            **xai_settings,
        )

    async def request(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
    ) -> ModelResponse:
        """Make a request to the xAI model."""
        check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(
            model_settings,
            model_request_parameters,
        )

        chat = await self._create_chat(messages, cast(XaiModelSettings, model_settings or {}), model_request_parameters)
        with _map_api_errors(self.model_name):
            response = await chat.sample()
        return self._process_response(response)

    @asynccontextmanager
    async def request_stream(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
        run_context: RunContext[Any] | None = None,
    ) -> AsyncIterator[StreamedResponse]:
        """Make a streaming request to the xAI model."""
        check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(
            model_settings,
            model_request_parameters,
        )

        chat = await self._create_chat(messages, cast(XaiModelSettings, model_settings or {}), model_request_parameters)
        response_stream = chat.stream()
        try:
            yield await self._process_streamed_response(response_stream, model_request_parameters)
        finally:
            aclose = getattr(response_stream, 'aclose', None)
            if aclose is not None:  # pragma: no branch
                await aclose()

    def _process_response(self, response: chat_types.Response) -> ModelResponse:
        """Convert xAI SDK response to pydantic_ai ModelResponse.

        Processes response.proto.outputs to extract (in order):
        - ThinkingPart: For reasoning/thinking content
        - TextPart: For text content
        - ToolCallPart: For client-side tool calls
        - NativeToolCallPart + NativeToolReturnPart: For server-side (builtin) tool calls
        """
        parts: list[ModelResponsePart] = []
        outputs = response.proto.outputs

        for output in outputs:
            message = output.message

            # Add reasoning/thinking content if present
            if message.reasoning_content or message.encrypted_content:
                signature = message.encrypted_content or None
                parts.append(
                    ThinkingPart(
                        content=message.reasoning_content or '',
                        signature=signature,
                        provider_name=self.system if signature else None,
                    )
                )

            # Add text content from assistant messages
            if message.content and message.role == chat_types.chat_pb2.MessageRole.ROLE_ASSISTANT:
                part_provider_details: dict[str, Any] | None = None
                if output.logprobs and output.logprobs.content:
                    part_provider_details = {'logprobs': _map_logprobs(output.logprobs)}
                parts.append(TextPart(content=message.content, provider_details=part_provider_details))

            # Process tool calls in this output
            for tool_call in message.tool_calls:
                tool_result_content = _get_tool_result_content(message.content)
                _, part = _create_tool_call_part(
                    tool_call,
                    tool_result_content,
                    self.system,
                    message_role=message.role,
                )
                parts.append(part)

        # xAI returns x_search results as top-level `response.citations` (URL strings) rather than
        # on the ROLE_TOOL message's `content`. Surface them on the corresponding return part so
        # users who set `include_output=True` can actually see the citations.
        _attach_x_search_citations(
            (p for p in parts if isinstance(p, NativeToolReturnPart) and p.tool_name == XSearchTool.kind),
            response.citations,
        )

        # Convert usage with detailed token information
        usage = _extract_usage(response, self._model_name, self._provider.name, self._provider.base_url)

        # Map finish reason.
        #
        # The xAI SDK exposes `response.finish_reason` as a *string* for the overall response, but in
        # multi-output responses (e.g. server-side tools) it can reflect an intermediate TOOL_CALLS
        # output rather than the final STOP output. We derive the finish reason from the final output
        # when available.
        if outputs:
            last_reason = outputs[-1].finish_reason
            finish_reason = _FINISH_REASON_PROTO_MAP.get(last_reason, 'stop')
        else:  # pragma: no cover
            finish_reason = _FINISH_REASON_MAP.get(response.finish_reason, 'stop')

        return ModelResponse(
            parts=parts,
            usage=usage,
            model_name=self._model_name,
            timestamp=response.created,
            provider_name=self.system,
            provider_url=self._provider.base_url,
            provider_response_id=response.id,
            finish_reason=finish_reason,
        )

    async def _process_streamed_response(
        self,
        response: AsyncIterator[tuple[chat_types.Response, Any]],
        model_request_parameters: ModelRequestParameters,
    ) -> 'XaiStreamedResponse':
        """Process a streamed response, and prepare a streaming response to return."""
        peekable_response: _utils.PeekableAsyncStream[
            tuple[chat_types.Response, chat_types.Chunk],
            AsyncIterator[tuple[chat_types.Response, Any]],
        ] = _utils.PeekableAsyncStream(response)
        with _map_api_errors(self.model_name):
            first_item = await peekable_response.peek()
        if isinstance(first_item, _utils.Unset):
            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')

        first_response, _ = first_item

        return XaiStreamedResponse(
            model_request_parameters=model_request_parameters,
            _model_name=self._model_name,
            _response=peekable_response,
            _timestamp=first_response.created,
            _provider=self._provider,
        )

__init__

__init__(
    model_name: XaiModelName,
    *,
    provider: (
        Literal["xai"] | Provider[AsyncClient]
    ) = "xai",
    profile: ModelProfileSpec | None = None,
    settings: ModelSettings | None = None
)

Initialize the xAI model.

Parameters:

Name	Type	Description	Default
model_name	XaiModelName	The name of the xAI model to use (e.g., "grok-4-1-fast-non-reasoning")	required
provider	Literal['xai'] | Provider[AsyncClient]	The provider to use for API calls. Defaults to 'xai'.	'xai'
profile	ModelProfileSpec | None	Optional model profile specification. Defaults to a profile picked by the provider based on the model name.	None
settings	ModelSettings | None	Optional model settings.	None

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

def __init__(
    self,
    model_name: XaiModelName,
    *,
    provider: Literal['xai'] | Provider[AsyncClient] = 'xai',
    profile: ModelProfileSpec | None = None,
    settings: ModelSettings | None = None,
):
    """Initialize the xAI model.

    Args:
        model_name: The name of the xAI model to use (e.g., "grok-4-1-fast-non-reasoning")
        provider: The provider to use for API calls. Defaults to `'xai'`.
        profile: Optional model profile specification. Defaults to a profile picked by the provider based on the model name.
        settings: Optional model settings.
    """
    self._model_name = model_name

    if isinstance(provider, str):
        provider = infer_provider(provider)
    self._provider = provider

    super().__init__(settings=settings, profile=profile)

model_name property

model_name: str

The model name.

system property

system: str

The model provider.

supported_native_tools classmethod

supported_native_tools() -> frozenset[type]

Return the set of builtin tool types this model can handle.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

@classmethod
def supported_native_tools(cls) -> frozenset[type]:
    """Return the set of builtin tool types this model can handle."""
    return frozenset({WebSearchTool, CodeExecutionTool, MCPServerTool, XSearchTool, FileSearchTool})

request async

request(
    messages: list[ModelMessage],
    model_settings: ModelSettings | None,
    model_request_parameters: ModelRequestParameters,
) -> ModelResponse

Make a request to the xAI model.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

async def request(
    self,
    messages: list[ModelMessage],
    model_settings: ModelSettings | None,
    model_request_parameters: ModelRequestParameters,
) -> ModelResponse:
    """Make a request to the xAI model."""
    check_allow_model_requests()
    model_settings, model_request_parameters = self.prepare_request(
        model_settings,
        model_request_parameters,
    )

    chat = await self._create_chat(messages, cast(XaiModelSettings, model_settings or {}), model_request_parameters)
    with _map_api_errors(self.model_name):
        response = await chat.sample()
    return self._process_response(response)

request_stream async

request_stream(
    messages: list[ModelMessage],
    model_settings: ModelSettings | None,
    model_request_parameters: ModelRequestParameters,
    run_context: RunContext[Any] | None = None,
) -> AsyncIterator[StreamedResponse]

Make a streaming request to the xAI model.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

@asynccontextmanager
async def request_stream(
    self,
    messages: list[ModelMessage],
    model_settings: ModelSettings | None,
    model_request_parameters: ModelRequestParameters,
    run_context: RunContext[Any] | None = None,
) -> AsyncIterator[StreamedResponse]:
    """Make a streaming request to the xAI model."""
    check_allow_model_requests()
    model_settings, model_request_parameters = self.prepare_request(
        model_settings,
        model_request_parameters,
    )

    chat = await self._create_chat(messages, cast(XaiModelSettings, model_settings or {}), model_request_parameters)
    response_stream = chat.stream()
    try:
        yield await self._process_streamed_response(response_stream, model_request_parameters)
    finally:
        aclose = getattr(response_stream, 'aclose', None)
        if aclose is not None:  # pragma: no branch
            await aclose()

XaiStreamedResponse dataclass

Bases: StreamedResponse

Implementation of StreamedResponse for xAI SDK.

Source code in pydantic_ai_slim/pydantic_ai/models/xai.py

@dataclass
class XaiStreamedResponse(StreamedResponse):
    """Implementation of `StreamedResponse` for xAI SDK."""

    _model_name: str
    _response: _utils.PeekableAsyncStream[
        tuple[chat_types.Response, chat_types.Chunk],
        AsyncIterator[tuple[chat_types.Response, Any]],
    ]
    _timestamp: datetime
    _provider: Provider[AsyncClient]

    def get_stream_cancel_errors(self) -> tuple[type[BaseException], ...]:
        return (grpc.RpcError,)

    async def close_stream(self) -> None:
        # In xai-sdk 1.5.0, `chat.stream()` returns a Python async generator that
        # wraps the underlying gRPC `GetCompletionChunk(...)` call.
        #
        # Calling `aclose()` shuts down that local async-generator wrapper and
        # stops consumption on our side, but the SDK does not expose the inner
        # `grpc.aio.UnaryStreamCall`, so this is not a documented transport-level
        # RPC cancellation hook.
        try:
            await self._response.source.aclose()  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
        except RuntimeError as exc:
            if not _utils.is_async_generator_already_running(exc):
                raise

    @property
    def system(self) -> str:
        """The model provider system name."""
        return self._provider.name

    @property
    def provider_url(self) -> str:
        """Get the provider base URL."""
        return self._provider.base_url

    def _update_response_state(self, response: chat_types.Response) -> None:
        """Update response state including usage, response ID, and finish reason."""
        # Update usage (SDK Response always provides a usage object)
        self._usage = _extract_usage(response, self._model_name, self._provider.name, self._provider.base_url)

        # Set provider response ID (only set once)
        if response.id and self.provider_response_id is None:
            self.provider_response_id = response.id

        # Handle finish reason (SDK Response always provides a finish_reason)
        self.finish_reason = _FINISH_REASON_MAP.get(response.finish_reason, 'stop')

    def _collect_reasoning_events(
        self,
        *,
        response: chat_types.Response,
        prev_reasoning_content: str,
        prev_encrypted_content: str,
    ) -> tuple[str, str, list[ModelResponseStreamEvent]]:
        """Collect thinking/reasoning events and return updated previous values.

        Note: xAI exposes reasoning via the accumulated Response object (not the per-chunk delta), so we compute
        deltas ourselves to avoid re-emitting the entire accumulated content on every chunk.
        """
        events: list[ModelResponseStreamEvent] = []

        if response.reasoning_content and response.reasoning_content != prev_reasoning_content:
            if response.reasoning_content.startswith(prev_reasoning_content):
                reasoning_delta = response.reasoning_content[len(prev_reasoning_content) :]
            else:
                reasoning_delta = response.reasoning_content
            prev_reasoning_content = response.reasoning_content
            if reasoning_delta:  # pragma: no branch
                events.extend(
                    self._parts_manager.handle_thinking_delta(
                        vendor_part_id='reasoning',
                        content=reasoning_delta,
                        # Only set provider_name when we have an encrypted signature to send back.
                        provider_name=self.system if response.encrypted_content else None,
                    )
                )

        if response.encrypted_content and response.encrypted_content != prev_encrypted_content:
            prev_encrypted_content = response.encrypted_content
            events.extend(
                self._parts_manager.handle_thinking_delta(
                    vendor_part_id='reasoning',
                    signature=response.encrypted_content,
                    provider_name=self.system,
                )
            )

        return prev_reasoning_content, prev_encrypted_content, events

    def _handle_server_side_tool_call(
        self,
        *,
        tool_call: chat_pb2.ToolCall,
        delta: chat_pb2.Delta,
        seen_tool_call_ids: set[str],
        seen_tool_return_ids: set[str],
        last_tool_return_content: dict[str, dict[str, Any] | str | None],
        x_search_return_parts: dict[str, NativeToolReturnPart],
    ) -> Iterator[ModelResponseStreamEvent]:
        """Handle a single server-side tool call delta, yielding stream events."""
        builtin_tool_name = _get_builtin_tool_name(tool_call)

        if delta.role == chat_pb2.MessageRole.ROLE_ASSISTANT:
            # Emit the call part once per tool_call_id.
            if tool_call.id in seen_tool_call_ids:
                return
            seen_tool_call_ids.add(tool_call.id)

            if builtin_tool_name.startswith(MCPServerTool.kind):
                parsed_args = _build_mcp_tool_call_args(tool_call)
            else:
                parsed_args = _parse_tool_args(tool_call.function.arguments)
            call_part = NativeToolCallPart(
                tool_name=builtin_tool_name,
                args=parsed_args,
                tool_call_id=tool_call.id,
                provider_name=self.system,
                provider_details={'function_name': tool_call.function.name},
            )
            if builtin_tool_name == CodeExecutionTool.kind:
                call_part.otel_metadata = {'code_arg_name': 'code', 'code_arg_language': 'python'}
            yield self._parts_manager.handle_part(vendor_part_id=tool_call.id, part=call_part)
            return

        if delta.role == chat_pb2.MessageRole.ROLE_TOOL:
            # Emit the return part once per tool_call_id.
            return_vendor_id = f'{tool_call.id}_return'
            tool_result_content = _get_tool_result_content(delta.content)
            if return_vendor_id in seen_tool_return_ids and tool_result_content == last_tool_return_content.get(
                return_vendor_id
            ):
                return
            seen_tool_return_ids.add(return_vendor_id)
            last_tool_return_content[return_vendor_id] = tool_result_content
            return_part = NativeToolReturnPart(
                tool_name=builtin_tool_name,
                content=tool_result_content,
                tool_call_id=tool_call.id,
                provider_name=self.system,
            )
            if builtin_tool_name == XSearchTool.kind:
                x_search_return_parts[return_vendor_id] = return_part
            yield self._parts_manager.handle_part(vendor_part_id=return_vendor_id, part=return_part)

    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        with _map_api_errors(self._model_name):
            # Local state to avoid re-emmiting duplicate events.
            prev_reasoning_content = ''
            prev_encrypted_content = ''
            seen_tool_call_ids: set[str] = set()
            seen_tool_return_ids: set[str] = set()
            last_tool_return_content: dict[str, dict[str, Any] | str | None] = {}
            # Track previous tool call args to compute deltas (like we do for reasoning content).
            prev_tool_call_args: dict[str, str] = {}
            # xAI exposes x_search results as top-level `response.citations` that only arrive with the
            # final chunk. Track the emitted x_search return parts so we can backfill their content
            # once the stream completes.
            x_search_return_parts: dict[str, NativeToolReturnPart] = {}
            last_citations: Sequence[str] = ()

            async for response, chunk in self._response:
                self._update_response_state(response)
                last_citations = response.citations

                prev_reasoning_content, prev_encrypted_content, reasoning_events = self._collect_reasoning_events(
                    response=response,
                    prev_reasoning_content=prev_reasoning_content,
                    prev_encrypted_content=prev_encrypted_content,
                )
                for event in reasoning_events:
                    yield event

                # Handle text content (property filters for ROLE_ASSISTANT)
                if chunk.content:
                    for event in self._parts_manager.handle_text_delta(
                        vendor_part_id='content',
                        content=chunk.content,
                    ):
                        yield event

                # Handle tool calls/tool results from *this chunk*.
                #
                # Important: xAI SDK `Response` is an accumulated view; `response.tool_calls` includes tool calls from
                # previous chunks. Iterating over it would re-emit tool calls repeatedly. Instead, we read tool calls
                # from the chunk's deltas which represent what changed in this frame.
                for output_chunk in chunk.proto.outputs:
                    delta = output_chunk.delta
                    if not delta.tool_calls:
                        continue
                    for tool_call in delta.tool_calls:
                        if not tool_call.function.name:
                            continue

                        if tool_call.type != chat_pb2.ToolCallType.TOOL_CALL_TYPE_CLIENT_SIDE_TOOL:
                            for event in self._handle_server_side_tool_call(
                                tool_call=tool_call,
                                delta=delta,
                                seen_tool_call_ids=seen_tool_call_ids,
                                seen_tool_return_ids=seen_tool_return_ids,
                                last_tool_return_content=last_tool_return_content,
                                x_search_return_parts=x_search_return_parts,
                            ):
                                yield event
                        else:
                            # Client-side tools: emit args as deltas so UI adapters receive PartDeltaEvents
                            # (not repeated PartStartEvents). Use accumulated args from response.tool_calls
                            # and compute the delta like we do for reasoning content.
                            accumulated = next((tc for tc in response.tool_calls if tc.id == tool_call.id), None)
                            accumulated_args = (
                                accumulated.function.arguments
                                if accumulated is not None and accumulated.function.arguments
                                else tool_call.function.arguments
                            )
                            prev_args = prev_tool_call_args.get(tool_call.id, '')
                            is_new_tool_call = tool_call.id not in prev_tool_call_args
                            args_changed = accumulated_args != prev_args

                            if is_new_tool_call or args_changed:
                                # Compute delta: if accumulated starts with prev, extract the new portion.
                                if accumulated_args.startswith(prev_args):
                                    args_delta = accumulated_args[len(prev_args) :] or None
                                else:
                                    args_delta = accumulated_args or None
                                prev_tool_call_args[tool_call.id] = accumulated_args
                                maybe_event = self._parts_manager.handle_tool_call_delta(
                                    vendor_part_id=tool_call.id,
                                    # Only pass tool_name on the first call; it would be appended otherwise.
                                    tool_name=tool_call.function.name if is_new_tool_call else None,
                                    args=args_delta,
                                    tool_call_id=tool_call.id,
                                )
                                if maybe_event is not None:  # pragma: no branch
                                    yield maybe_event

            # Backfill x_search return parts with the top-level response citations, which xAI only
            # populates alongside the final chunk. We mutate the existing part in place rather than
            # re-emitting via `handle_part`: the parts manager holds the same object reference we
            # tracked in `x_search_return_parts`, so the mutation is reflected in the final
            # `ModelResponse` without emitting a duplicate `PartStartEvent` at the same index.
            _attach_x_search_citations(x_search_return_parts.values(), last_citations)

    @property
    def model_name(self) -> str:
        """Get the model name of the response."""
        return self._model_name

    @property
    def provider_name(self) -> str:
        """The model provider."""
        return self.system

    @property
    def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        return self._timestamp

system property

system: str

The model provider system name.

provider_url property

provider_url: str

Get the provider base URL.

model_name property

model_name: str

Get the model name of the response.

provider_name property

provider_name: str

The model provider.

timestamp property

timestamp: datetime

Get the timestamp of the response.