Address PR feedback for Cerebras integration

sebastiand-cerebras · sebastiand-cerebras · commit 613483ee5d51 · 2025-11-25T15:30:59.000-08:00
- Update docs/models/cerebras.md: use pip/uv-add format, link to Cerebras docs
- Fix httpx.AsyncClient typo in cerebras.md, groq.md, mistral.md
- Add docs/api/models/cerebras.md and update mkdocs.yml
- Remove Cerebras section from openai.md, move to main list in overview.md
- Add str | to CerebrasModelName for arbitrary model names
- Add CerebrasModelSettings with cerebras_disable_reasoning field
- Add zai_model_profile, restore unsupported_model_settings and json_schema_transformer
- Pass lowercase model name to profile functions
- Add tests/providers/test_cerebras.py with full coverage
- Remove type ignore in models/__init__.py
diff --git a/docs/api/models/cerebras.md b/docs/api/models/cerebras.md
@@ -0,0 +1,7 @@
+# `pydantic_ai.models.cerebras`
+
+## Setup
+
+For details on how to set up authentication with this model, see [model configuration for Cerebras](../../models/cerebras.md).
+
+::: pydantic_ai.models.cerebras
diff --git a/docs/models/cerebras.md b/docs/models/cerebras.md
@@ -5,20 +5,14 @@
 To use `CerebrasModel`, you need to either install `pydantic-ai`, or install `pydantic-ai-slim` with the `cerebras` optional group:
 
 ```bash
-pip install "pydantic-ai-slim[cerebras]"
-```
-
-or
-
-```bash
-uv add "pydantic-ai-slim[cerebras]"
+pip/uv-add "pydantic-ai-slim[cerebras]"
 ```
 
 ## Configuration
 
-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.
+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 generate an API key.
 
-`CerebrasModelName` contains a list of available Cerebras models.
+For a list of available models, see the [Cerebras models documentation](https://inference-docs.cerebras.ai/models).
 
 ## Environment variable
 
@@ -64,7 +58,7 @@ agent = Agent(model)
 ...
 ```
 
-You can also customize the `CerebrasProvider` with a custom `httpx.AsyncHTTPClient`:
+You can also customize the `CerebrasProvider` with a custom `httpx.AsyncClient`:
 
 ```python
 from httpx import AsyncClient
diff --git a/docs/models/groq.md b/docs/models/groq.md
@@ -58,7 +58,7 @@ agent = Agent(model)
 ...
 ```
 
-You can also customize the `GroqProvider` with a custom `httpx.AsyncHTTPClient`:
+You can also customize the `GroqProvider` with a custom `httpx.AsyncClient`:
 
 ```python
 from httpx import AsyncClient
diff --git a/docs/models/mistral.md b/docs/models/mistral.md
@@ -58,7 +58,7 @@ agent = Agent(model)
 ...
 ```
 
-You can also customize the provider with a custom `httpx.AsyncHTTPClient`:
+You can also customize the provider with a custom `httpx.AsyncClient`:
 
 ```python
 from httpx import AsyncClient
diff --git a/docs/models/openai.md b/docs/models/openai.md
@@ -632,39 +632,6 @@ agent = Agent(model)
 ...
 ```
 
-### Cerebras
-
-To use [Cerebras](https://cerebras.ai/), you need to create an API key in the [Cerebras Console](https://cloud.cerebras.ai/).
-
-You can set the `CEREBRAS_API_KEY` environment variable and use [`CerebrasProvider`][pydantic_ai.providers.cerebras.CerebrasProvider] by name:
-
-```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.
-```
-
-Or initialise the model and provider directly:
-
-```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.
-```
-
 ### LiteLLM
 
 To use [LiteLLM](https://www.litellm.ai/), set the configs as outlined in the [doc](https://docs.litellm.ai/docs/set_keys). In `LiteLLMProvider`, you can pass `api_base` and `api_key`. The value of these configs will depend on your setup. For example, if you are using OpenAI models, then you need to pass `https://api.openai.com/v1` as the `api_base` and your OpenAI API key as the `api_key`. If you are using a LiteLLM proxy server running on your local machine, then you need to pass `http://localhost:<port>` as the `api_base` and your LiteLLM API key (or a placeholder) as the `api_key`.
diff --git a/docs/models/overview.md b/docs/models/overview.md
@@ -9,6 +9,7 @@ Pydantic AI is model-agnostic and has built-in support for multiple model provid
 * [Mistral](mistral.md)
 * [Cohere](cohere.md)
 * [Bedrock](bedrock.md)
+* [Cerebras](cerebras.md)
 * [Hugging Face](huggingface.md)
 * [Outlines](outlines.md)
 
@@ -27,7 +28,6 @@ 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)
 - [LiteLLM](openai.md#litellm)
 - [Nebius AI Studio](openai.md#nebius-ai-studio)
 - [OVHcloud AI Endpoints](openai.md#ovhcloud-ai-endpoints)
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -147,6 +147,7 @@ nav:
           - api/models/openai.md
           - api/models/anthropic.md
           - api/models/bedrock.md
+          - api/models/cerebras.md
           - api/models/cohere.md
           - api/models/google.md
           - api/models/groq.md
diff --git a/pydantic_ai_slim/pydantic_ai/models/__init__.py b/pydantic_ai_slim/pydantic_ai/models/__init__.py
@@ -819,7 +819,7 @@ def infer_model(  # noqa: C901
     if model_kind == 'cerebras':
         from .cerebras import CerebrasModel
 
-        return CerebrasModel(model_name, provider=provider)  # type: ignore[arg-type]
+        return CerebrasModel(model_name, provider=provider)
     elif model_kind == 'openai-chat':
         from .openai import OpenAIChatModel
 
diff --git a/pydantic_ai_slim/pydantic_ai/models/cerebras.py b/pydantic_ai_slim/pydantic_ai/models/cerebras.py
@@ -18,9 +18,9 @@
         'you can use the `cerebras` optional group — `pip install "pydantic-ai-slim[cerebras]"'
     ) from _import_error
 
-__all__ = ('CerebrasModel', 'CerebrasModelName')
+__all__ = ('CerebrasModel', 'CerebrasModelName', 'CerebrasModelSettings')
 
-CerebrasModelName = Literal[
+_KnownCerebrasModelName = Literal[
     'gpt-oss-120b',
     'llama-3.3-70b',
     'llama3.1-8b',
@@ -29,6 +29,28 @@
     'zai-glm-4.6',
 ]
 
+CerebrasModelName = str | _KnownCerebrasModelName
+"""Possible Cerebras model names.
+
+Since Cerebras supports a variety of models and the list changes frequently, we explicitly list known models
+but allow any name in the type hints.
+
+See <https://inference-docs.cerebras.ai/models/overview> for an up to date list of models.
+"""
+
+
+class CerebrasModelSettings(ModelSettings, total=False):
+    """Settings used for a Cerebras model request.
+
+    ALL FIELDS MUST BE `cerebras_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    """
+
+    cerebras_disable_reasoning: bool
+    """Disable reasoning for the model.
+
+    See [the Cerebras docs](https://inference-docs.cerebras.ai/resources/openai#passing-non-standard-parameters) for more details.
+    """
+
 
 @dataclass(init=False)
 class CerebrasModel(OpenAIChatModel):
@@ -45,7 +67,7 @@ def __init__(
         *,
         provider: Literal['cerebras'] | Provider[AsyncOpenAI] = 'cerebras',
         profile: ModelProfileSpec | None = None,
-        settings: ModelSettings | None = None,
+        settings: CerebrasModelSettings | None = None,
     ):
         """Initialize a Cerebras model.
 
diff --git a/pydantic_ai_slim/pydantic_ai/providers/cerebras.py b/pydantic_ai_slim/pydantic_ai/providers/cerebras.py
@@ -10,7 +10,7 @@
 from pydantic_ai.models import cached_async_http_client
 from pydantic_ai.profiles.harmony import harmony_model_profile
 from pydantic_ai.profiles.meta import meta_model_profile
-from pydantic_ai.profiles.openai import OpenAIModelProfile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
 from pydantic_ai.profiles.qwen import qwen_model_profile
 from pydantic_ai.providers import Provider
 
@@ -23,6 +23,15 @@
     ) from _import_error
 
 
+def zai_model_profile(model_name: str) -> ModelProfile | None:
+    """The model profile for ZAI models on Cerebras."""
+    return ModelProfile(
+        supports_json_object_output=True,
+        supports_json_schema_output=True,
+        json_schema_transformer=OpenAIJsonSchemaTransformer,
+    )
+
+
 class CerebrasProvider(Provider[AsyncOpenAI]):
     """Provider for Cerebras API."""
 
@@ -43,18 +52,30 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
             'llama': meta_model_profile,
             'qwen': qwen_model_profile,
             'gpt-oss': harmony_model_profile,
+            'zai': zai_model_profile,
         }
 
         profile = None
         model_name_lower = model_name.lower()
         for prefix, profile_func in prefix_to_profile.items():
             if model_name_lower.startswith(prefix):
-                profile = profile_func(model_name)
+                profile = profile_func(model_name_lower)
                 break
 
-        # Wrap in OpenAIModelProfile with web search disabled
-        # Cerebras doesn't support web search
-        return OpenAIModelProfile(openai_chat_supports_web_search=False).update(profile)
+        # According to https://inference-docs.cerebras.ai/resources/openai#currently-unsupported-openai-features,
+        # Cerebras doesn't support some model settings.
+        # openai_chat_supports_web_search=False is default, so not required here
+        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: ...
diff --git a/tests/providers/test_cerebras.py b/tests/providers/test_cerebras.py
@@ -0,0 +1,122 @@
+from __future__ import annotations as _annotations
+
+import re
+
+import httpx
+import pytest
+from pytest_mock import MockerFixture
+
+from pydantic_ai._json_schema import InlineDefsJsonSchemaTransformer
+from pydantic_ai.exceptions import UserError
+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 ..conftest import TestEnv, try_import
+
+with try_import() as imports_successful:
+    from openai import AsyncOpenAI
+
+    from pydantic_ai.models import infer_model
+    from pydantic_ai.models.cerebras import CerebrasModel
+    from pydantic_ai.providers.cerebras import CerebrasProvider, zai_model_profile
+
+
+pytestmark = pytest.mark.skipif(not imports_successful(), reason='openai not installed')
+
+
+def test_cerebras_provider():
+    provider = CerebrasProvider(api_key='api-key')
+    assert provider.name == 'cerebras'
+    assert provider.base_url == 'https://api.cerebras.ai/v1'
+    assert isinstance(provider.client, AsyncOpenAI)
+    assert provider.client.api_key == 'api-key'
+
+
+def test_cerebras_provider_need_api_key(env: TestEnv) -> None:
+    env.remove('CEREBRAS_API_KEY')
+    with pytest.raises(
+        UserError,
+        match=re.escape(
+            'Set the `CEREBRAS_API_KEY` environment variable or pass it via `CerebrasProvider(api_key=...)` '
+            'to use the Cerebras provider.'
+        ),
+    ):
+        CerebrasProvider()
+
+
+def test_cerebras_provider_pass_http_client() -> None:
+    http_client = httpx.AsyncClient()
+    provider = CerebrasProvider(http_client=http_client, api_key='api-key')
+    assert provider.client._client == http_client  # type: ignore[reportPrivateUsage]
+
+
+def test_cerebras_provider_pass_openai_client() -> None:
+    openai_client = AsyncOpenAI(api_key='api-key')
+    provider = CerebrasProvider(openai_client=openai_client)
+    assert provider.client == openai_client
+
+
+def test_cerebras_provider_model_profile(mocker: MockerFixture):
+    openai_client = AsyncOpenAI(api_key='api-key')
+    provider = CerebrasProvider(openai_client=openai_client)
+
+    ns = 'pydantic_ai.providers.cerebras'
+    meta_model_profile_mock = mocker.patch(f'{ns}.meta_model_profile', wraps=meta_model_profile)
+    qwen_model_profile_mock = mocker.patch(f'{ns}.qwen_model_profile', wraps=qwen_model_profile)
+    harmony_model_profile_mock = mocker.patch(f'{ns}.harmony_model_profile', wraps=harmony_model_profile)
+    zai_model_profile_mock = mocker.patch(f'{ns}.zai_model_profile', wraps=zai_model_profile)
+
+    # Test llama model - uses meta profile which has InlineDefsJsonSchemaTransformer
+    meta_profile = provider.model_profile('llama-3.3-70b')
+    meta_model_profile_mock.assert_called_with('llama-3.3-70b')
+    assert meta_profile is not None
+    assert isinstance(meta_profile, OpenAIModelProfile)
+    assert meta_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
+
+    # Test qwen model - uses qwen profile which has InlineDefsJsonSchemaTransformer
+    qwen_profile = provider.model_profile('qwen-3-32b')
+    qwen_model_profile_mock.assert_called_with('qwen-3-32b')
+    assert qwen_profile is not None
+    assert isinstance(qwen_profile, OpenAIModelProfile)
+    assert qwen_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
+
+    # Test gpt-oss model (harmony) - uses OpenAIJsonSchemaTransformer
+    harmony_profile = provider.model_profile('gpt-oss-120b')
+    harmony_model_profile_mock.assert_called_with('gpt-oss-120b')
+    assert harmony_profile is not None
+    assert isinstance(harmony_profile, OpenAIModelProfile)
+    assert harmony_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test zai model
+    zai_profile = provider.model_profile('zai-glm-4.6')
+    zai_model_profile_mock.assert_called_with('zai-glm-4.6')
+    assert zai_profile is not None
+    assert isinstance(zai_profile, OpenAIModelProfile)
+    assert zai_profile.supports_json_object_output is True
+    assert zai_profile.supports_json_schema_output is True
+    assert zai_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Test unknown model - should still return a profile with OpenAIJsonSchemaTransformer
+    unknown_profile = provider.model_profile('unknown-model')
+    assert unknown_profile is not None
+    assert isinstance(unknown_profile, OpenAIModelProfile)
+    assert unknown_profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Verify unsupported model settings are set for all profiles
+    for profile in [meta_profile, qwen_profile, harmony_profile, zai_profile, unknown_profile]:
+        assert isinstance(profile, OpenAIModelProfile)
+        assert 'frequency_penalty' in profile.openai_unsupported_model_settings
+        assert 'logit_bias' in profile.openai_unsupported_model_settings
+        assert 'presence_penalty' in profile.openai_unsupported_model_settings
+        assert 'parallel_tool_calls' in profile.openai_unsupported_model_settings
+        assert 'service_tier' in profile.openai_unsupported_model_settings
+
+
+def test_infer_cerebras_model(env: TestEnv):
+    """Test that infer_model correctly creates a CerebrasModel from a model name string."""
+    env.set('CEREBRAS_API_KEY', 'test-api-key')
+    model = infer_model('cerebras:llama-3.3-70b')
+    assert isinstance(model, CerebrasModel)
+    assert model.model_name == 'llama-3.3-70b'
