Add Cohere docs and additional (live) tests (#810)

Co-authored-by: David Montague <[email protected]>

Author: sydney-runkle

.github/workflows/ci.yml

@@ -98,6 +98,7 @@ jobs:
           --extra groq
           --extra anthropic
           --extra mistral
+          --extra cohere
           pytest tests/test_live.py -v
           --durations=100
         env:
@@ -108,6 +109,7 @@ jobs:
           GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
           MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}
+          CO_API_KEY: ${{ secrets.COHERE_API_KEY }}
 
   test:
     name: test on ${{ matrix.python-version }}

docs/api/models/cohere.md

@@ -0,0 +1,7 @@
+# `pydantic_ai.models.cohere`
+
+## Setup
+
+For details on how to set up authentication with this model, see [model configuration for Cohere](../../models.md#cohere).
+
+::: pydantic_ai.models.cohere

docs/install.md

@@ -53,6 +53,7 @@ pip/uv-add 'pydantic-ai-slim[openai]'
 * `anthropic` — installs `anthropic` [PyPI ↗](https://pypi.org/project/anthropic){:target="_blank"}
 * `groq` — installs `groq` [PyPI ↗](https://pypi.org/project/groq){:target="_blank"}
 * `mistral` — installs `mistralai` [PyPI ↗](https://pypi.org/project/mistralai){:target="_blank"}
+* `cohere` - installs `cohere` [PyPI ↗](https://pypi.org/project/cohere){:target="_blank"}
 
 See the [models](models.md) documentation for information on which optional dependencies are required for each model.
 

docs/models.md

@@ -7,6 +7,7 @@ PydanticAI is Model-agnostic and has built in support for the following model pr
 * [Deepseek](#deepseek)
 * [Groq](#groq)
 * [Mistral](#mistral)
+* [Cohere](#cohere)
 
 See [OpenAI-compatible models](#openai-compatible-models) for more examples on how to use models such as [OpenRouter](#openrouter), and [Grok (xAI)](#grok-xai) that support the OpenAI SDK.
 
@@ -419,6 +420,63 @@ agent = Agent(model)
 ...
 ```
 
+## Cohere
+
+### Install
+
+To use [`CohereModel`][pydantic_ai.models.cohere.CohereModel], you need to either install [`pydantic-ai`](install.md), or install [`pydantic-ai-slim`](install.md#slim-install) with the `cohere` optional group:
+
+```bash
+pip/uv-add 'pydantic-ai-slim[cohere]'
+```
+
+### Configuration
+
+To use [Cohere](https://cohere.com/) through their API, go to [dashboard.cohere.com/api-keys](https://dashboard.cohere.com/api-keys) and follow your nose until you find the place to generate an API key.
+
+[`NamedCohereModels`][pydantic_ai.models.cohere.NamedCohereModels] contains a list of the most popular Cohere models.
+
+### Environment variable
+
+Once you have the API key, you can set it as an environment variable:
+
+```bash
+export CO_API_KEY='your-api-key'
+```
+
+You can then use [`CohereModel`][pydantic_ai.models.cohere.CohereModel] by name:
+
+```python {title="cohere_model_by_name.py"}
+from pydantic_ai import Agent
+
+agent = Agent('cohere:command')
+...
+```
+
+Or initialise the model directly with just the model name:
+
+```python {title="cohere_model_init.py"}
+from pydantic_ai import Agent
+from pydantic_ai.models.cohere import CohereModel
+
+model = CohereModel('command', api_key='your-api-key')
+agent = Agent(model)
+...
+```
+
+### `api_key` argument
+
+If you don't want to or can't set the environment variable, you can pass it at runtime via the [`api_key` argument][pydantic_ai.models.cohere.CohereModel.__init__]:
+
+```python {title="cohere_model_api_key.py"}
+from pydantic_ai import Agent
+from pydantic_ai.models.cohere import CohereModel
+
+model = CohereModel('command', api_key='your-api-key')
+agent = Agent(model)
+...
+```
+
 ## OpenAI-compatible Models
 
 Many of the models are compatible with OpenAI API, and thus can be used with [`OpenAIModel`][pydantic_ai.models.openai.OpenAIModel] in PydanticAI.

mkdocs.yml

@@ -50,6 +50,7 @@ nav:
     - api/models/base.md
     - api/models/openai.md
     - api/models/anthropic.md
+    - api/models/cohere.md
     - api/models/gemini.md
     - api/models/vertexai.md
     - api/models/groq.md

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -3,9 +3,10 @@
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 from itertools import chain
-from typing import Literal, TypeAlias, Union, cast
+from typing import Literal, Union, cast
 
 from cohere import TextAssistantMessageContentItem
+from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
 from .. import result
@@ -51,24 +52,24 @@
         "you can use the `cohere` optional group — `pip install 'pydantic-ai-slim[cohere]'`"
     ) from _import_error
 
-CohereModelName: TypeAlias = Union[
-    str,
-    Literal[
-        'c4ai-aya-expanse-32b',
-        'c4ai-aya-expanse-8b',
-        'command',
-        'command-light',
-        'command-light-nightly',
-        'command-nightly',
-        'command-r',
-        'command-r-03-2024',
-        'command-r-08-2024',
-        'command-r-plus',
-        'command-r-plus-04-2024',
-        'command-r-plus-08-2024',
-        'command-r7b-12-2024',
-    ],
+NamedCohereModels = Literal[
+    'c4ai-aya-expanse-32b',
+    'c4ai-aya-expanse-8b',
+    'command',
+    'command-light',
+    'command-light-nightly',
+    'command-nightly',
+    'command-r',
+    'command-r-03-2024',
+    'command-r-08-2024',
+    'command-r-plus',
+    'command-r-plus-04-2024',
+    'command-r-plus-08-2024',
+    'command-r7b-12-2024',
 ]
+"""Latest / most popular named Cohere models."""
+
+CohereModelName = Union[NamedCohereModels, str]
 
 
 class CohereModelSettings(ModelSettings):
@@ -96,23 +97,26 @@ def __init__(
         *,
         api_key: str | None = None,
         cohere_client: AsyncClientV2 | None = None,
+        http_client: AsyncHTTPClient | None = None,
     ):
         """Initialize an Cohere model.
 
         Args:
             model_name: The name of the Cohere model to use. List of model names
                 available [here](https://docs.cohere.com/docs/models#command).
             api_key: The API key to use for authentication, if not provided, the
-                `COHERE_API_KEY` environment variable will be used if available.
+                `CO_API_KEY` environment variable will be used if available.
             cohere_client: An existing Cohere async client to use. If provided,
-                `api_key` must be `None`.
+                `api_key` and `http_client` must be `None`.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
         self.model_name: CohereModelName = model_name
         if cohere_client is not None:
+            assert http_client is None, 'Cannot provide both `cohere_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
             self.client = cohere_client
         else:
-            self.client = AsyncClientV2(api_key=api_key)  # type: ignore
+            self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)  # type: ignore
 
     async def agent_model(
         self,

tests/models/test_model.py

@@ -56,7 +56,7 @@
         'MistralModel',
     ),
     (
-        'COHERE_API_KEY',
+        'CO_API_KEY',
         'cohere:command',
         'cohere:command',
         'cohere',

tests/test_live.py

@@ -66,6 +66,12 @@ def mistral(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
     return MistralModel('mistral-small-latest', http_client=http_client)
 
 
+def cohere(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
+    from pydantic_ai.models.cohere import CohereModel
+
+    return CohereModel('command-r7b-12-2024', http_client=http_client)
+
+
 params = [
     pytest.param(openai, id='openai'),
     pytest.param(gemini, marks=pytest.mark.skip(reason='API seems very flaky'), id='gemini'),
@@ -74,6 +80,7 @@ def mistral(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
     pytest.param(anthropic, id='anthropic'),
     pytest.param(ollama, id='ollama'),
     pytest.param(mistral, id='mistral'),
+    pytest.param(cohere, id='cohere'),
 ]
 GetModel = Callable[[httpx.AsyncClient, Path], Model]
 
@@ -95,7 +102,7 @@ async def test_text(http_client: httpx.AsyncClient, tmp_path: Path, get_model: G
     assert usage.total_tokens is not None and usage.total_tokens > 0
 
 
-stream_params = [p for p in params if p.id != 'anthropic']
+stream_params = [p for p in params if p.id != 'cohere']
 
 
 @pytest.mark.parametrize('get_model', stream_params)