Add Cerebras provider (#2643)

Author: Kludex

docs/api/providers.md

@@ -16,6 +16,8 @@
 
 ::: pydantic_ai.providers.cohere
 
+::: pydantic_ai.providers.cerebras.CerebrasProvider
+
 ::: pydantic_ai.providers.mistral.MistralProvider
 
 ::: pydantic_ai.providers.fireworks.FireworksProvider

docs/models/index.md

@@ -26,6 +26,7 @@ In addition, many providers are compatible with the OpenAI API, and can be used
 - [Azure AI Foundry](openai.md#azure-ai-foundry)
 - [Heroku](openai.md#heroku-ai)
 - [GitHub Models](openai.md#github-models)
+- [Cerebras](openai.md#cerebras)
 
 Pydantic AI also comes with [`TestModel`](../api/models/test.md) and [`FunctionModel`](../api/models/function.md)
 for testing and development.

docs/models/openai.md

@@ -530,3 +530,36 @@ You can set the `HEROKU_INFERENCE_KEY` and `HEROKU_INFERENCE_URL` environment va
 export HEROKU_INFERENCE_KEY='your-heroku-inference-key'
 export HEROKU_INFERENCE_URL='https://us.inference.heroku.com'
 ```
+
+### Cerebras
+
+To use [Cerebras](https://cerebras.ai/), you need to create an API key in the [Cerebras Console](https://cloud.cerebras.ai/).
+
+Once you've set the `CEREBRAS_API_KEY` environment variable, you can run the following:
+
+```python
+from pydantic_ai import Agent
+
+agent = Agent('cerebras:llama3.3-70b')
+result = agent.run_sync('What is the capital of France?')
+print(result.output)
+#> The capital of France is Paris.
+```
+
+If you need to configure the provider, you can use the [`CerebrasProvider`][pydantic_ai.providers.cerebras.CerebrasProvider] class:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.cerebras import CerebrasProvider
+
+model = OpenAIChatModel(
+    'llama3.3-70b',
+    provider=CerebrasProvider(api_key='your-cerebras-api-key'),
+)
+agent = Agent(model)
+
+result = agent.run_sync('What is the capital of France?')
+print(result.output)
+#> The capital of France is Paris.
+```

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -111,6 +111,15 @@
         'bedrock:mistral.mixtral-8x7b-instruct-v0:1',
         'bedrock:mistral.mistral-large-2402-v1:0',
         'bedrock:mistral.mistral-large-2407-v1:0',
+        'cerebras:gpt-oss-120b',
+        'cerebras:llama3.1-8b',
+        'cerebras:llama-3.3-70b',
+        'cerebras:llama-4-scout-17b-16e-instruct',
+        'cerebras:llama-4-maverick-17b-128e-instruct',
+        'cerebras:qwen-3-235b-a22b-instruct-2507',
+        'cerebras:qwen-3-32b',
+        'cerebras:qwen-3-coder-480b',
+        'cerebras:qwen-3-235b-a22b-thinking-2507',
         'claude-3-5-haiku-20241022',
         'claude-3-5-haiku-latest',
         'claude-3-5-sonnet-20240620',
@@ -695,18 +704,19 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
 
         return CohereModel(model_name, provider=provider)
     elif provider in (
-        'openai',
-        'openai-chat',
-        'deepseek',
         'azure',
-        'openrouter',
-        'vercel',
+        'deepseek',
+        'cerebras',
+        'fireworks',
+        'github',
         'grok',
+        'heroku',
         'moonshotai',
-        'fireworks',
+        'openai',
+        'openai-chat',
+        'openrouter',
         'together',
-        'heroku',
-        'github',
+        'vercel',
     ):
         from .openai import OpenAIChatModel
 

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -211,19 +211,20 @@ def __init__(
         model_name: OpenAIModelName,
         *,
         provider: Literal[
-            'openai',
-            'openai-chat',
-            'deepseek',
             'azure',
-            'openrouter',
-            'moonshotai',
-            'vercel',
-            'grok',
+            'deepseek',
+            'cerebras',
             'fireworks',
-            'together',
-            'heroku',
             'github',
+            'grok',
+            'heroku',
+            'moonshotai',
             'ollama',
+            'openai',
+            'openai-chat',
+            'openrouter',
+            'together',
+            'vercel',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -237,19 +238,20 @@ def __init__(
         model_name: OpenAIModelName,
         *,
         provider: Literal[
-            'openai',
-            'openai-chat',
-            'deepseek',
             'azure',
-            'openrouter',
-            'moonshotai',
-            'vercel',
-            'grok',
+            'deepseek',
+            'cerebras',
             'fireworks',
-            'together',
-            'heroku',
             'github',
+            'grok',
+            'heroku',
+            'moonshotai',
             'ollama',
+            'openai',
+            'openai-chat',
+            'openrouter',
+            'together',
+            'vercel',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -262,19 +264,20 @@ def __init__(
         model_name: OpenAIModelName,
         *,
         provider: Literal[
-            'openai',
-            'openai-chat',
-            'deepseek',
             'azure',
-            'openrouter',
-            'moonshotai',
-            'vercel',
-            'grok',
+            'deepseek',
+            'cerebras',
             'fireworks',
-            'together',
-            'heroku',
             'github',
+            'grok',
+            'heroku',
+            'moonshotai',
             'ollama',
+            'openai',
+            'openai-chat',
+            'openrouter',
+            'together',
+            'vercel',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -402,6 +405,11 @@ async def _completions_create(
         ):  # pragma: no branch
             response_format = {'type': 'json_object'}
 
+        unsupported_model_settings = OpenAIModelProfile.from_profile(self.profile).openai_unsupported_model_settings
+        for setting in unsupported_model_settings:
+            model_settings.pop(setting, None)
+
+        # TODO(Marcelo): Deprecate this in favor of `openai_unsupported_model_settings`.
         sampling_settings = (
             model_settings
             if OpenAIModelProfile.from_profile(self.profile).openai_supports_sampling_settings
@@ -646,9 +654,7 @@ async def _map_user_message(self, message: ModelRequest) -> AsyncIterable[chat.C
                 )
             elif isinstance(part, RetryPromptPart):
                 if part.tool_name is None:
-                    yield chat.ChatCompletionUserMessageParam(  # pragma: no cover
-                        role='user', content=part.model_response()
-                    )
+                    yield chat.ChatCompletionUserMessageParam(role='user', content=part.model_response())
                 else:
                     yield chat.ChatCompletionToolMessageParam(
                         role='tool',

pydantic_ai_slim/pydantic_ai/profiles/harmony.py

@@ -0,0 +1,13 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+from .openai import OpenAIModelProfile, openai_model_profile
+
+
+def harmony_model_profile(model_name: str) -> ModelProfile | None:
+    """The model profile for the OpenAI Harmony Response format.
+
+    See <https://cookbook.openai.com/articles/openai-harmony> for more details.
+    """
+    profile = openai_model_profile(model_name)
+    return OpenAIModelProfile(openai_supports_tool_choice_required=False).update(profile)

pydantic_ai_slim/pydantic_ai/profiles/openai.py

@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 import re
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any, Literal
 
@@ -20,9 +21,13 @@ class OpenAIModelProfile(ModelProfile):
     openai_supports_strict_tool_definition: bool = True
     """This can be set by a provider or user if the OpenAI-"compatible" API doesn't support strict tool definitions."""
 
+    # TODO(Marcelo): Deprecate this in favor of `openai_unsupported_model_settings`.
     openai_supports_sampling_settings: bool = True
     """Turn off to don't send sampling settings like `temperature` and `top_p` to models that don't support them, like OpenAI's o-series reasoning models."""
 
+    openai_unsupported_model_settings: Sequence[str] = ()
+    """A list of model settings that are not supported by the model."""
+
     # Some OpenAI-compatible providers (e.g. MoonshotAI) currently do **not** accept
     # `tool_choice="required"`.  This flag lets the calling model know whether it's
     # safe to pass that value along.  Default is `True` to preserve existing

pydantic_ai_slim/pydantic_ai/profiles/qwen.py

@@ -1,10 +1,18 @@
 from __future__ import annotations as _annotations
 
+from ..profiles.openai import OpenAIModelProfile
 from . import InlineDefsJsonSchemaTransformer, ModelProfile
 
 
 def qwen_model_profile(model_name: str) -> ModelProfile | None:
     """Get the model profile for a Qwen model."""
+    if model_name.startswith('qwen-3-coder'):
+        return OpenAIModelProfile(
+            json_schema_transformer=InlineDefsJsonSchemaTransformer,
+            openai_supports_tool_choice_required=False,
+            openai_supports_strict_tool_definition=False,
+            ignore_streamed_leading_whitespace=True,
+        )
     return ModelProfile(
         json_schema_transformer=InlineDefsJsonSchemaTransformer,
         ignore_streamed_leading_whitespace=True,

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -95,6 +95,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .mistral import MistralProvider
 
         return MistralProvider
+    elif provider == 'cerebras':
+        from .cerebras import CerebrasProvider
+
+        return CerebrasProvider
     elif provider == 'cohere':
         from .cohere import CohereProvider
 

pydantic_ai_slim/pydantic_ai/providers/cerebras.py

@@ -0,0 +1,96 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+import httpx
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.harmony import harmony_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
+from pydantic_ai.profiles.qwen import qwen_model_profile
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+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]"`'
+    ) from _import_error
+
+
+class CerebrasProvider(Provider[AsyncOpenAI]):
+    """Provider for Cerebras API."""
+
+    @property
+    def name(self) -> str:
+        return 'cerebras'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://api.cerebras.ai/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        prefix_to_profile = {'llama': meta_model_profile, 'qwen': qwen_model_profile, 'gpt-oss': harmony_model_profile}
+
+        profile = None
+        for prefix, profile_func in prefix_to_profile.items():
+            model_name = model_name.lower()
+            if model_name.startswith(prefix):
+                profile = profile_func(model_name)
+
+        # According to https://inference-docs.cerebras.ai/resources/openai#currently-unsupported-openai-features,
+        # Cerebras doesn't support some model settings.
+        unsupported_model_settings = (
+            'frequency_penalty',
+            'logit_bias',
+            'presence_penalty',
+            'parallel_tool_calls',
+            'service_tier',
+        )
+        return OpenAIModelProfile(
+            json_schema_transformer=OpenAIJsonSchemaTransformer,
+            openai_unsupported_model_settings=unsupported_model_settings,
+        ).update(profile)
+
+    @overload
+    def __init__(self) -> 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,
+        openai_client: AsyncOpenAI | 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)
+        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)