Add cohere provider class (#1225)

Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: hrahmadi71

docs/models.md

@@ -611,20 +611,39 @@ Or initialise the model directly with just the model name:
 from pydantic_ai import Agent
 from pydantic_ai.models.cohere import CohereModel
 
-model = CohereModel('command', api_key='your-api-key')
+model = CohereModel('command')
 agent = Agent(model)
 ...
 ```
 
-### `api_key` argument
+### `provider` 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__]:
+You can provide a custom [`Provider`][pydantic_ai.providers.Provider] via the [`provider` argument][pydantic_ai.models.cohere.CohereModel.__init__]:
 
-```python {title="cohere_model_api_key.py"}
+```python {title="cohere_model_provider.py"}
 from pydantic_ai import Agent
 from pydantic_ai.models.cohere import CohereModel
+from pydantic_ai.providers.cohere import CohereProvider
 
-model = CohereModel('command', api_key='your-api-key')
+model = CohereModel('command', provider=CohereProvider(api_key='your-api-key'))
+agent = Agent(model)
+...
+```
+
+You can also customize the `CohereProvider` with a custom `http_client`:
+
+```python {title="cohere_model_custom_provider.py"}
+from httpx import AsyncClient
+
+from pydantic_ai import Agent
+from pydantic_ai.models.cohere import CohereModel
+from pydantic_ai.providers.cohere import CohereProvider
+
+custom_http_client = AsyncClient(timeout=30)
+model = CohereModel(
+    'command',
+    provider=CohereProvider(api_key='your-api-key', http_client=custom_http_client),
+)
 agent = Agent(model)
 ...
 ```

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -402,8 +402,7 @@ def infer_model(model: Model | KnownModelName) -> Model:
     if provider == 'cohere':
         from .cohere import CohereModel
 
-        # TODO(Marcelo): Missing provider API.
-        return CohereModel(model_name)
+        return CohereModel(model_name, provider=provider)
     elif provider in ('deepseek', 'openai'):
         from .openai import OpenAIModel
 

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -3,11 +3,11 @@
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 from itertools import chain
-from typing import Literal, Union, cast
+from typing import Literal, Union, cast, overload
 
 from cohere import TextAssistantMessageContentItem
 from httpx import AsyncClient as AsyncHTTPClient
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
 
 from .. import ModelHTTPError, result
 from .._utils import guard_tool_call_id as _guard_tool_call_id
@@ -23,11 +23,13 @@
     ToolReturnPart,
     UserPromptPart,
 )
+from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     Model,
     ModelRequestParameters,
+    cached_async_http_client,
     check_allow_model_requests,
 )
 
@@ -100,10 +102,34 @@ class CohereModel(Model):
     _model_name: CohereModelName = field(repr=False)
     _system: str = field(default='cohere', repr=False)
 
+    @overload
     def __init__(
         self,
         model_name: CohereModelName,
         *,
+        provider: Literal['cohere'] | Provider[AsyncClientV2] = 'cohere',
+        api_key: None = None,
+        cohere_client: None = None,
+        http_client: None = None,
+    ) -> None: ...
+
+    @deprecated('Use the `provider` parameter instead of `api_key`, `cohere_client`, and `http_client`.')
+    @overload
+    def __init__(
+        self,
+        model_name: CohereModelName,
+        *,
+        provider: None = None,
+        api_key: str | None = None,
+        cohere_client: AsyncClientV2 | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        model_name: CohereModelName,
+        *,
+        provider: Literal['cohere'] | Provider[AsyncClientV2] | None = None,
         api_key: str | None = None,
         cohere_client: AsyncClientV2 | None = None,
         http_client: AsyncHTTPClient | None = None,
@@ -113,19 +139,27 @@ def __init__(
         Args:
             model_name: The name of the Cohere model to use. List of model names
                 available [here](https://docs.cohere.com/docs/models#command).
+            provider: The provider to use for authentication and API access. Can be either the string
+                'cohere' or an instance of `Provider[AsyncClientV2]`. If not provided, a new provider will be
+                created using the other parameters.
             api_key: The API key to use for authentication, if not provided, the
                 `CO_API_KEY` environment variable will be used if available.
             cohere_client: An existing Cohere async 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.
         """
         self._model_name: CohereModelName = model_name
-        if cohere_client is not None:
+
+        if provider is not None:
+            if isinstance(provider, str):
+                provider = infer_provider(provider)
+            self.client = provider.client
+        elif 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, httpx_client=http_client)
+            self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client or cached_async_http_client())
 
     @property
     def base_url(self) -> str:

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -77,5 +77,9 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .mistral import MistralProvider
 
         return MistralProvider()
+    elif provider == 'cohere':
+        from .cohere import CohereProvider
+
+        return CohereProvider()
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')

pydantic_ai_slim/pydantic_ai/providers/cohere.py

@@ -0,0 +1,72 @@
+from __future__ import annotations as _annotations
+
+import os
+
+from httpx import AsyncClient as AsyncHTTPClient
+
+from pydantic_ai.models import cached_async_http_client
+
+try:
+    from cohere import AsyncClientV2
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `cohere` package to use the Cohere provider, '
+        'you can use the `cohere` optional group — `pip install "pydantic-ai-slim[cohere]"`'
+    ) from _import_error
+
+
+from . import Provider
+
+
+class CohereProvider(Provider[AsyncClientV2]):
+    """Provider for Cohere API."""
+
+    @property
+    def name(self) -> str:
+        return 'cohere'
+
+    @property
+    def base_url(self) -> str:
+        client_wrapper = self.client._client_wrapper  # type: ignore
+        return str(client_wrapper.get_base_url())
+
+    @property
+    def client(self) -> AsyncClientV2:
+        return self._client
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        cohere_client: AsyncClientV2 | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        """Create a new Cohere provider.
+
+        Args:
+            api_key: The API key to use for authentication, if not provided, the `CO_API_KEY` environment variable
+                will be used if available.
+            cohere_client: An existing
+                [AsyncClientV2](https://github.com/cohere-ai/cohere-python)
+                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 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:
+            api_key = api_key or os.environ.get('CO_API_KEY')
+            if api_key is None:
+                raise ValueError(
+                    'Set the `CO_API_KEY` environment variable or pass it via `CohereProvider(api_key=...)`'
+                    'to use the Cohere provider.'
+                )
+
+            base_url = os.environ.get('CO_BASE_URL')
+            if http_client is not None:
+                self._client = AsyncClientV2(api_key=api_key, httpx_client=http_client, base_url=base_url)
+            else:
+                self._client = AsyncClientV2(
+                    api_key=api_key, httpx_client=cached_async_http_client(), base_url=base_url
+                )

tests/conftest.py

@@ -246,6 +246,11 @@ def anthropic_api_key() -> str:
     return os.getenv('ANTHROPIC_API_KEY', 'mock-api-key')
 
 
+@pytest.fixture(scope='session')
+def co_api_key() -> str:
+    return os.getenv('CO_API_KEY', 'mock-api-key')
+
+
 @pytest.fixture
 def mock_snapshot_id(mocker: MockerFixture):
     i = 0

tests/models/cassettes/test_cohere/test_request_simple_success_with_vcr.yaml

@@ -0,0 +1,64 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '93'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.com
+    method: POST
+    parsed_body:
+      messages:
+      - content: hello
+        role: user
+      model: command-r7b-12-2024
+      stream: false
+    uri: https://api.cohere.com/v2/chat
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-length:
+      - '286'
+      content-type:
+      - application/json
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      num_chars:
+      - '2583'
+      num_tokens:
+      - '10'
+      pragma:
+      - no-cache
+      vary:
+      - Origin
+    parsed_body:
+      finish_reason: COMPLETE
+      id: f17a5f6c-1734-4098-bd0d-733ef000ac7b
+      message:
+        content:
+        - text: Hello! How can I assist you today?
+          type: text
+        role: assistant
+      usage:
+        billed_units:
+          input_tokens: 1
+          output_tokens: 9
+        tokens:
+          input_tokens: 496
+          output_tokens: 11
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/test_cohere.py

@@ -38,6 +38,7 @@
     from cohere.core.api_error import ApiError
 
     from pydantic_ai.models.cohere import CohereModel
+    from pydantic_ai.providers.cohere import CohereProvider
 
     # note: we use Union here for compatibility with Python 3.9
     MockChatResponse = Union[ChatResponse, Exception]
@@ -49,7 +50,7 @@
 
 
 def test_init():
-    m = CohereModel('command-r7b-12-2024', api_key='foobar')
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(api_key='foobar'))
     assert m.model_name == 'command-r7b-12-2024'
     assert m.system == 'cohere'
     assert m.base_url == 'https://api.cohere.com'
@@ -96,7 +97,7 @@ async def test_request_simple_success(allow_model_requests: None):
         )
     )
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
 
     result = await agent.run('hello')
@@ -135,7 +136,7 @@ async def test_request_simple_usage(allow_model_requests: None):
         ),
     )
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
 
     result = await agent.run('Hello')
@@ -169,7 +170,7 @@ async def test_request_structured_response(allow_model_requests: None):
         )
     )
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m, result_type=list[int])
 
     result = await agent.run('Hello')
@@ -243,7 +244,7 @@ async def test_request_tool_call(allow_model_requests: None):
         ),
     ]
     mock_client = MockAsyncClientV2.create_mock(responses)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m, system_prompt='this is the system prompt')
 
     @agent.tool_plain
@@ -326,7 +327,7 @@ async def get_location(loc_name: str) -> str:
 async def test_multimodal(allow_model_requests: None):
     c = completion_message(AssistantMessageResponse(content=[TextAssistantMessageResponseContentItem(text='world')]))
     mock_client = MockAsyncClientV2.create_mock(c)
-    m = CohereModel('command-r7b-12-2024', cohere_client=mock_client)
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
 
     with pytest.raises(RuntimeError, match='Cohere does not yet support multi-modal inputs.'):
@@ -347,8 +348,16 @@ def test_model_status_error(allow_model_requests: None) -> None:
             body={'error': 'test error'},
         )
     )
-    m = CohereModel('command-r', cohere_client=mock_client)
+    m = CohereModel('command-r', provider=CohereProvider(cohere_client=mock_client))
     agent = Agent(m)
     with pytest.raises(ModelHTTPError) as exc_info:
         agent.run_sync('hello')
     assert str(exc_info.value) == snapshot("status_code: 500, model_name: command-r, body: {'error': 'test error'}")
+
+
+@pytest.mark.vcr()
+async def test_request_simple_success_with_vcr(allow_model_requests: None, co_api_key: str):
+    m = CohereModel('command-r7b-12-2024', provider=CohereProvider(api_key=co_api_key))
+    agent = Agent(m)
+    result = await agent.run('hello')
+    assert result.data == snapshot('Hello! How can I assist you today?')