Renmae GrokModel to XaiModel, and update dependencies and tests. Added XaiProvider and move api_key and client to this

Author: brightsparc

docs/models/grok.md

@@ -0,0 +1,73 @@
+# xAI
+
+## Install
+
+To use [`XaiModel`][pydantic_ai.models.xai.XaiModel], you need to either install `pydantic-ai`, or install `pydantic-ai-slim` with the `xai` optional group:
+
+```bash
+pip/uv-add "pydantic-ai-slim[xai]"
+```
+
+## Configuration
+
+To use xAI models from [xAI](https://x.ai/api) through their API, go to [console.x.ai](https://console.x.ai/team/default/api-keys) to create an API key.
+
+[`GrokModelName`][pydantic_ai.providers.grok.GrokModelName] contains a list of available xAI models.
+
+## Environment variable
+
+Once you have the API key, you can set it as an environment variable:
+
+```bash
+export XAI_API_KEY='your-api-key'
+```
+
+You can then use [`XaiModel`][pydantic_ai.models.xai.XaiModel] by name:
+
+```python
+from pydantic_ai import Agent
+
+agent = Agent('xai:grok-4-1-fast-non-reasoning')
+...
+```
+
+Or initialise the model directly:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.xai import XaiModel
+
+# Uses XAI_API_KEY environment variable
+model = XaiModel('grok-4-1-fast-non-reasoning')
+agent = Agent(model)
+...
+```
+
+You can also customize the [`XaiModel`][pydantic_ai.models.xai.XaiModel] with a custom provider:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.xai import XaiModel
+from pydantic_ai.providers.xai import XaiProvider
+
+# Custom API key
+provider = XaiProvider(api_key='your-api-key')
+model = XaiModel('grok-4-1-fast-non-reasoning', provider=provider)
+agent = Agent(model)
+...
+```
+
+Or with a custom `xai_sdk.AsyncClient`:
+
+```python
+from xai_sdk import AsyncClient
+from pydantic_ai import Agent
+from pydantic_ai.models.xai import XaiModel
+from pydantic_ai.providers.xai import XaiProvider
+
+xai_client = AsyncClient(api_key='your-api-key')
+provider = XaiProvider(xai_client=xai_client)
+model = XaiModel('grok-4-1-fast-non-reasoning', provider=provider)
+agent = Agent(model)
+...
+```

docs/models/xai.md

@@ -5,8 +5,6 @@
 2. Provides buy analysis for the user
 """
 
-import os
-
 import logfire
 from pydantic import BaseModel, Field
 
@@ -15,19 +13,16 @@
     BuiltinToolCallPart,
     WebSearchTool,
 )
-from pydantic_ai.models.grok import GrokModel
+from pydantic_ai.models.xai import XaiModel
 
 logfire.configure()
 logfire.instrument_pydantic_ai()
 
-# Configure for xAI API
-xai_api_key = os.getenv('XAI_API_KEY')
-if not xai_api_key:
-    raise ValueError('XAI_API_KEY environment variable is required')
-
+# Configure for xAI API - XAI_API_KEY environment variable is required
+# The model will automatically use XaiProvider with the API key from the environment
 
-# Create the model using GrokModel with server-side tools
-model = GrokModel('grok-4-fast', api_key=xai_api_key)
+# Create the model using XaiModel with server-side tools
+model = XaiModel('grok-4-fast')
 
 
 class StockAnalysis(BaseModel):

examples/pydantic_ai_examples/stock_analysis_agent.py

@@ -180,6 +180,19 @@
         'grok:grok-4-fast-non-reasoning',
         'grok:grok-4-fast-reasoning',
         'grok:grok-code-fast-1',
+        'xai:grok-2-image-1212',
+        'xai:grok-2-vision-1212',
+        'xai:grok-3',
+        'xai:grok-3-fast',
+        'xai:grok-3-mini',
+        'xai:grok-3-mini-fast',
+        'xai:grok-4',
+        'xai:grok-4-0709',
+        'xai:grok-4-1-fast-non-reasoning',
+        'xai:grok-4-1-fast-reasoning',
+        'xai:grok-4-fast-non-reasoning',
+        'xai:grok-4-fast-reasoning',
+        'xai:grok-code-fast-1',
         'groq:deepseek-r1-distill-llama-70b',
         'groq:deepseek-r1-distill-qwen-32b',
         'groq:distil-whisper-large-v3-en',
@@ -809,6 +822,7 @@ def infer_model(  # noqa: C901
         'fireworks',
         'github',
         'grok',
+        'xai',
         'heroku',
         'moonshotai',
         'ollama',

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -1,17 +1,22 @@
-"""Grok model implementation using xAI SDK."""
+"""xAI model implementation using xAI SDK."""
 
-import os
 from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
-import xai_sdk.chat as chat_types
+try:
+    import xai_sdk.chat as chat_types
 
-# Import xai_sdk components
-from xai_sdk import AsyncClient
-from xai_sdk.chat import assistant, image, system, tool, tool_result, user
-from xai_sdk.tools import code_execution, get_tool_call_type, mcp, web_search  # x_search not yet supported
+    # Import xai_sdk components
+    from xai_sdk import AsyncClient
+    from xai_sdk.chat import assistant, image, system, tool, tool_result, user
+    from xai_sdk.tools import code_execution, get_tool_call_type, mcp, web_search  # x_search not yet supported
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `xai-sdk` to use the xAI model, '
+        'you can use the `xai` optional group — `pip install "pydantic-ai-slim[xai]"`'
+    ) from _import_error
 
 from .._run_context import RunContext
 from .._utils import now_utc
@@ -41,42 +46,46 @@
     ModelRequestParameters,
     StreamedResponse,
 )
+from ..profiles import ModelProfileSpec
+from ..providers import Provider, infer_provider
+from ..providers.grok import GrokModelName
 from ..settings import ModelSettings
 from ..usage import RequestUsage
 
+# Type alias for consistency
+XaiModelName = GrokModelName
 
-class GrokModel(Model):
-    """A model that uses the xAI SDK to interact with Grok."""
+
+class XaiModel(Model):
+    """A model that uses the xAI SDK to interact with xAI models."""
 
     _model_name: str
-    _api_key: str
-    _client: AsyncClient | None
+    _provider: Provider[AsyncClient]
 
     def __init__(
         self,
-        model_name: str,
+        model_name: XaiModelName,
         *,
-        api_key: str | None = None,
-        client: AsyncClient | None = None,
+        provider: Literal['xai'] | Provider[AsyncClient] = 'xai',
+        profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
     ):
-        """Initialize the Grok model.
+        """Initialize the xAI model.
 
         Args:
-            model_name: The name of the Grok model to use (e.g., "grok-4-1-fast-non-reasoning")
-            api_key: The xAI API key. If not provided, uses XAI_API_KEY environment variable.
-            client: Optional AsyncClient instance for testing. If provided, api_key is ignored.
+            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.
         """
-        super().__init__(settings=settings)
         self._model_name = model_name
-        self._client = client
-        if client is None:
-            self._api_key = api_key or os.getenv('XAI_API_KEY') or ''
-            if not self._api_key:
-                raise ValueError('XAI API key is required')
-        else:
-            self._api_key = api_key or ''
+
+        if isinstance(provider, str):
+            provider = infer_provider(provider)
+        self._provider = provider
+        self.client = provider.client
+
+        super().__init__(settings=settings, profile=profile or provider.model_profile)
 
     @property
     def model_name(self) -> str:
@@ -188,7 +197,7 @@ def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -
                 )
             else:
                 raise UserError(
-                    f'`{builtin_tool.__class__.__name__}` is not supported by `GrokModel`. '
+                    f'`{builtin_tool.__class__.__name__}` is not supported by `XaiModel`. '
                     f'Supported built-in tools: WebSearchTool, CodeExecutionTool, MCPServerTool. '
                     f'If XSearchTool should be supported, please file an issue.'
                 )
@@ -200,9 +209,8 @@ async def request(
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
-        """Make a request to the Grok model."""
-        # Use injected client or create one in the current async context
-        client = self._client or AsyncClient(api_key=self._api_key)
+        """Make a request to the xAI model."""
+        client = self._provider.client
 
         # Convert messages to xAI format
         xai_messages = self._map_messages(messages)
@@ -253,9 +261,8 @@ async def request_stream(
         model_request_parameters: ModelRequestParameters,
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
-        """Make a streaming request to the Grok model."""
-        # Use injected client or create one in the current async context
-        client = self._client or AsyncClient(api_key=self._api_key)
+        """Make a streaming request to the xAI model."""
+        client = self._provider.client
 
         # Convert messages to xAI format
         xai_messages = self._map_messages(messages)
@@ -294,7 +301,7 @@ async def request_stream(
 
         # Stream the response
         response_stream = chat.stream()
-        streamed_response = GrokStreamedResponse(
+        streamed_response = XaiStreamedResponse(
             model_request_parameters=model_request_parameters,
             _model_name=self._model_name,
             _response=response_stream,
@@ -409,7 +416,7 @@ def _process_response(self, response: chat_types.Response) -> ModelResponse:
 
     def _map_usage(self, response: chat_types.Response) -> RequestUsage:
         """Extract usage information from xAI SDK response, including reasoning tokens and cache tokens."""
-        return GrokModel.extract_usage(response)
+        return XaiModel.extract_usage(response)
 
     @staticmethod
     def extract_usage(response: chat_types.Response) -> RequestUsage:
@@ -472,7 +479,7 @@ def extract_usage(response: chat_types.Response) -> RequestUsage:
 
 
 @dataclass
-class GrokStreamedResponse(StreamedResponse):
+class XaiStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for xAI SDK."""
 
     _model_name: str
@@ -486,7 +493,7 @@ def _update_response_state(self, response: Any) -> None:
 
         # Update usage
         if hasattr(response, 'usage'):
-            self._usage = GrokModel.extract_usage(response)
+            self._usage = XaiModel.extract_usage(response)
 
         # Set provider response ID
         if hasattr(response, 'id') and self.provider_response_id is None:

pydantic_ai_slim/pydantic_ai/models/grok.py renamed to pydantic_ai_slim/pydantic_ai/models/xai.py

@@ -7,18 +7,18 @@
 
 @dataclass(kw_only=True)
 class GrokModelProfile(ModelProfile):
-    """Profile for models used with GroqModel.
+    """Profile for Grok models (used with both GrokProvider and XaiProvider).
 
-    ALL FIELDS MUST BE `groq_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    ALL FIELDS MUST BE `grok_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
     """
 
-    builtin_tool: bool = False
+    grok_supports_builtin_tools: bool = False
     """Whether the model always has the web search built-in tool available."""
 
 
 def grok_model_profile(model_name: str) -> ModelProfile | None:
     """Get the model profile for a Grok model."""
     return GrokModelProfile(
         # Support tool calling for building tools
-        builtin_tool=model_name.startswith('grok-4'),
+        grok_supports_builtin_tools=model_name.startswith('grok-4'),
     )