Add Cerebras model support with official SDK

- Implement CerebrasModel subclassing OpenAIChatModel
- Disable web search (Cerebras-specific feature)
- Use cerebras-cloud-sdk for official API support
- All tests passing (6/6 Cerebras-specific tests)

Author: sebastiand-cerebras

docs/models/cerebras.md

@@ -1,65 +1,77 @@
 # Cerebras
 
-Cerebras provides ultra-fast inference using their Wafer-Scale Engine (WSE), delivering predictable performance for any workload.
+## Install
 
-## Installation
-
-To use Cerebras, you need to either install `pydantic-ai`, or install `pydantic-ai-slim` with the `cerebras` optional group:
+To use `CerebrasModel`, you need to either install `pydantic-ai`, or install `pydantic-ai-slim` with the `cerebras` optional group:
 
 ```bash
-# pip
-pip install "pydantic-ai-slim[cerebras]"
-
-# uv
-uv add "pydantic-ai-slim[cerebras]"
+pip/uv-add "pydantic-ai-slim[cerebras]"
 ```
 
 ## Configuration
 
-To use Cerebras, go to [cloud.cerebras.ai](https://cloud.cerebras.ai/?utm_source=3pi_pydantic-ai&utm_campaign=partner_doc) to get an API key.
+To use [Cerebras](https://cerebras.ai/) through their API, go to [cloud.cerebras.ai](https://cloud.cerebras.ai/?utm_source=3pi_pydantic-ai&utm_campaign=partner_doc) and follow your nose until you find the place to generate an API key.
+
+`CerebrasModelName` contains a list of available Cerebras models.
 
-### Environment Variable
+## Environment variable
 
-Set your API key as an environment variable:
+Once you have the API key, you can set it as an environment variable:
 
 ```bash
 export CEREBRAS_API_KEY='your-api-key'
 ```
 
-### Available Models
-
-Cerebras supports the following models:
-
-- `llama-3.3-70b` (recommended) - Latest Llama 3.3 model
-- `llama-3.1-8b` - Llama 3.1 8B (faster, smaller)
-- `qwen-3-235b-a22b-instruct-2507` - Qwen 3 235B
-- `qwen-3-32b` - Qwen 3 32B
-- `gpt-oss-120b` - GPT-OSS 120B
-- `zai-glm-4.6` - GLM 4.6 model
+You can then use `CerebrasModel` by name:
 
+```python
+from pydantic_ai import Agent
 
-See the [Cerebras documentation](https://inference-docs.cerebras.ai/introduction?utm_source=3pi_pydantic-ai&utm_campaign=partner_doc) for the latest models.
+agent = Agent('cerebras:llama-3.3-70b')
+...
+```
 
-## Usage
+Or initialise the model directly with just the model name:
 
 ```python
 from pydantic_ai import Agent
+from pydantic_ai.models.cerebras import CerebrasModel
 
-agent = Agent('cerebras:llama-3.3-70b')
-result = agent.run_sync('What is the capital of France?')
-print(result.output)
-#> The capital of France is Paris.
+model = CerebrasModel('llama-3.3-70b')
+agent = Agent(model)
+...
 ```
 
-## Why Cerebras?
+## `provider` argument
+
+You can provide a custom `Provider` via the `provider` argument:
 
-- **Ultra-fast inference** - Powered by the world's largest AI chip (WSE)
-- **Predictable performance** - Consistent latency for any workload
-- **OpenAI-compatible** - Drop-in replacement for OpenAI API
-- **Cost-effective** - Competitive pricing with superior performance
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.cerebras import CerebrasModel
+from pydantic_ai.providers.cerebras import CerebrasProvider
+
+model = CerebrasModel(
+    'llama-3.3-70b', provider=CerebrasProvider(api_key='your-api-key')
+)
+agent = Agent(model)
+...
+```
 
-## Resources
+You can also customize the `CerebrasProvider` with a custom `httpx.AsyncHTTPClient`:
 
-- [Cerebras Inference Documentation](https://inference-docs.cerebras.ai?utm_source=3pi_pydantic-ai&utm_campaign=partner_doc)
-- [Get API Key](https://cloud.cerebras.ai/?utm_source=3pi_pydantic-ai&utm_campaign=partner_doc)
-- [Model Pricing](https://cerebras.ai/pricing?utm_source=3pi_pydantic-ai&utm_campaign=partner_doc)
+```python
+from httpx import AsyncClient
+
+from pydantic_ai import Agent
+from pydantic_ai.models.cerebras import CerebrasModel
+from pydantic_ai.providers.cerebras import CerebrasProvider
+
+custom_http_client = AsyncClient(timeout=30)
+model = CerebrasModel(
+    'llama-3.3-70b',
+    provider=CerebrasProvider(api_key='your-api-key', http_client=custom_http_client),
+)
+agent = Agent(model)
+...
+```

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -800,7 +800,6 @@ def infer_model(  # noqa: C901
         'openai',
         'azure',
         'deepseek',
-        'cerebras',
         'fireworks',
         'github',
         'grok',
@@ -818,7 +817,11 @@ def infer_model(  # noqa: C901
     elif model_kind in ('google-gla', 'google-vertex'):
         model_kind = 'google'
 
-    if model_kind == 'openai-chat':
+    if model_kind == 'cerebras':
+        from .cerebras import CerebrasModel
+
+        return CerebrasModel(model_name, provider=provider)
+    elif model_kind == 'openai-chat':
         from .openai import OpenAIChatModel
 
         return OpenAIChatModel(model_name, provider=provider)

pydantic_ai_slim/pydantic_ai/models/cerebras.py

@@ -0,0 +1,95 @@
+"""Cerebras model implementation using OpenAI-compatible API."""
+
+from __future__ import annotations as _annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+try:
+    from cerebras.cloud.sdk import AsyncCerebras  # noqa: F401
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `cerebras-cloud-sdk` package to use the Cerebras model, '
+        'you can use the `cerebras` optional group — `pip install "pydantic-ai-slim[cerebras]"`'
+    ) from _import_error
+
+from ..profiles import ModelProfile, ModelProfileSpec
+from ..profiles.harmony import harmony_model_profile
+from ..profiles.meta import meta_model_profile
+from ..profiles.qwen import qwen_model_profile
+from ..providers import Provider
+from ..settings import ModelSettings
+from .openai import OpenAIChatModel, OpenAIModelProfile  # type: ignore[attr-defined]
+
+__all__ = ('CerebrasModel', 'CerebrasModelName')
+
+CerebrasModelName = Literal[
+    'llama-3.3-70b',
+    'llama-4-scout-17b-16e-instruct',
+    'qwen-3-235b-a22b-instruct-2507',
+    'qwen-3-32b',
+    'gpt-oss-120b',
+    'zai-glm-4.6',
+]
+
+
+@dataclass(init=False)
+class CerebrasModel(OpenAIChatModel):
+    """A model that uses Cerebras's OpenAI-compatible API.
+
+    Cerebras provides ultra-fast inference powered by the Wafer-Scale Engine (WSE).
+
+    Apart from `__init__`, all methods are private or match those of the base class.
+    """
+
+    def __init__(
+        self,
+        model_name: CerebrasModelName,
+        *,
+        provider: Literal['cerebras'] | Provider[Any] = 'cerebras',
+        profile: ModelProfileSpec | None = None,
+        settings: ModelSettings | None = None,
+    ):
+        """Initialize a Cerebras model.
+
+        Args:
+            model_name: The name of the Cerebras model to use.
+            provider: The provider to use. Can be 'cerebras' or a Provider instance.
+            profile: The model profile to use. Defaults to a profile based on the model name.
+            settings: Model-specific settings that will be used as defaults for this model.
+        """
+        if provider == 'cerebras':
+            from ..providers.cerebras import CerebrasProvider
+
+            # Extract api_key from settings if provided
+            api_key = settings.get('api_key') if settings else None
+            provider = CerebrasProvider(api_key=api_key) if api_key else CerebrasProvider()  # type: ignore[call-overload]
+
+        # Use our custom model_profile method if no profile is provided
+        if profile is None:
+            profile = self._cerebras_model_profile
+
+        super().__init__(model_name, provider=provider, profile=profile, settings=settings)  # type: ignore[arg-type]
+
+    def _cerebras_model_profile(self, model_name: str) -> ModelProfile:
+        """Get the model profile for this Cerebras model.
+
+        Returns a profile with web search disabled since Cerebras doesn't support it.
+        """
+        model_name_lower = model_name.lower()
+
+        # Get base profile based on model family
+        if model_name_lower.startswith('llama'):
+            base_profile = meta_model_profile(model_name)
+        elif model_name_lower.startswith('qwen'):
+            base_profile = qwen_model_profile(model_name)
+        elif model_name_lower.startswith('gpt-oss'):
+            base_profile = harmony_model_profile(model_name)
+        else:
+            # Default profile for unknown models
+            base_profile = ModelProfile()
+
+        # Wrap in OpenAIModelProfile with web search disabled
+        return OpenAIModelProfile(
+            openai_chat_supports_web_search=False,
+        ).update(base_profile)

pydantic_ai_slim/pydantic_ai/providers/cerebras.py

@@ -14,15 +14,15 @@
 from pydantic_ai.providers import Provider
 
 try:
-    from openai import AsyncOpenAI
+    from cerebras.cloud.sdk import AsyncCerebras
 except ImportError as _import_error:  # pragma: no cover
     raise ImportError(
-        'Please install the `openai` package to use the Cerebras provider, '
-        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+        'Please install the `cerebras-cloud-sdk` package to use the Cerebras provider, '
+        'you can use the `cerebras` optional group — `pip install "pydantic-ai-slim[cerebras]"`'
     ) from _import_error
 
 
-class CerebrasProvider(Provider[AsyncOpenAI]):
+class CerebrasProvider(Provider[AsyncCerebras]):
     """Provider for Cerebras API."""
 
     @property
@@ -34,7 +34,7 @@ def base_url(self) -> str:
         return 'https://api.cerebras.ai/v1'
 
     @property
-    def client(self) -> AsyncOpenAI:
+    def client(self) -> AsyncCerebras:
         return self._client
 
     def model_profile(self, model_name: str) -> ModelProfile | None:
@@ -52,35 +52,46 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
         return None
 
     @overload
-    def __init__(self) -> None: ...
+    def __init__(self, *, cerebras_client: AsyncCerebras | None = None) -> None: ...
 
     @overload
-    def __init__(self, *, api_key: str) -> None: ...
-
-    @overload
-    def __init__(self, *, api_key: str, http_client: httpx.AsyncClient) -> None: ...
-
-    @overload
-    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+    def __init__(
+        self, *, api_key: str | None = None, base_url: str | None = None, http_client: httpx.AsyncClient | None = None
+    ) -> None: ...
 
     def __init__(
         self,
         *,
         api_key: str | None = None,
-        openai_client: AsyncOpenAI | None = None,
+        base_url: str | None = None,
+        cerebras_client: AsyncCerebras | None = None,
         http_client: httpx.AsyncClient | None = None,
     ) -> None:
-        api_key = api_key or os.getenv('CEREBRAS_API_KEY')
-        if not api_key and openai_client is None:
-            raise UserError(
-                'Set the `CEREBRAS_API_KEY` environment variable or pass it via `CerebrasProvider(api_key=...)` '
-                'to use the Cerebras provider.'
-            )
-
-        if openai_client is not None:
-            self._client = openai_client
-        elif http_client is not None:
-            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        """Create a new Cerebras provider.
+
+        Args:
+            api_key: The API key to use for authentication, if not provided, the `CEREBRAS_API_KEY` environment variable
+                will be used if available.
+            base_url: The base url for the Cerebras requests. If not provided, defaults to Cerebras's base url.
+            cerebras_client: An existing `AsyncCerebras` client to use. If provided, `api_key` and `http_client` must be `None`.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+        """
+        if cerebras_client is not None:
+            assert http_client is None, 'Cannot provide both `cerebras_client` and `http_client`'
+            assert api_key is None, 'Cannot provide both `cerebras_client` and `api_key`'
+            assert base_url is None, 'Cannot provide both `cerebras_client` and `base_url`'
+            self._client = cerebras_client
         else:
-            http_client = cached_async_http_client(provider='cerebras')
-            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+            api_key = api_key or os.getenv('CEREBRAS_API_KEY')
+            base_url = base_url or 'https://api.cerebras.ai/v1'
+
+            if not api_key:
+                raise UserError(
+                    'Set the `CEREBRAS_API_KEY` environment variable or pass it via `CerebrasProvider(api_key=...)` '
+                    'to use the Cerebras provider.'
+                )
+            elif http_client is not None:
+                self._client = AsyncCerebras(base_url=base_url, api_key=api_key, http_client=http_client)
+            else:
+                http_client = cached_async_http_client(provider='cerebras')
+                self._client = AsyncCerebras(base_url=base_url, api_key=api_key, http_client=http_client)

pydantic_ai_slim/pyproject.toml

@@ -72,7 +72,7 @@ cohere = ["cohere>=5.18.0; platform_system != 'Emscripten'"]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
 google = ["google-genai>=1.51.0"]
 anthropic = ["anthropic>=0.70.0"]
-cerebras = ["openai>=1.107.2"]
+cerebras = ["cerebras-cloud-sdk>=1.0.0"]
 groq = ["groq>=0.25.0"]
 mistral = ["mistralai>=1.9.10"]
 bedrock = ["boto3>=1.40.14"]