Merge pull request andrewyng#147 from andrewyng/add-cohere-provider

Add cohere provider

Author: ksolo

.github/workflows/run_pytest.yml

@@ -18,7 +18,7 @@ jobs:
  run: |
  python -m pip install --upgrade pip
  pip install poetry
- poetry install --with test
+ poetry install --all-extras --with test
  - name: Test with pytest
  run: poetry run pytest
 

.gitignore

@@ -9,4 +9,7 @@ env/
 .coverage
 
 # pyenv
-.python-version
+.python-version
+
+.DS_Store
+**/.DS_Store

aisuite/providers/cohere_provider.py

@@ -0,0 +1,37 @@
+import os
+import cohere
+
+from aisuite.framework import ChatCompletionResponse
+from aisuite.provider import Provider
+
+
+class CohereProvider(Provider):
+ def __init__(self, **config):
+ """
+ Initialize the Cohere provider with the given configuration.
+ Pass the entire configuration dictionary to the Cohere client constructor.
+ """
+ # Ensure API key is provided either in config or via environment variable
+ config.setdefault("api_key", os.getenv("CO_API_KEY"))
+ if not config["api_key"]:
+ raise ValueError(
+ " API key is missing. Please provide it in the config or set the CO_API_KEY environment variable."
+ )
+ self.client = cohere.ClientV2(**config)
+
+ def chat_completions_create(self, model, messages, **kwargs):
+ response = self.client.chat(
+ model=model,
+ messages=messages,
+ **kwargs # Pass any additional arguments to the Cohere API
+ )
+
+ return self.normalize_response(response)
+
+ def normalize_response(self, response):
+ """Normalize the reponse from Cohere API to match OpenAI's response format."""
+ normalized_response = ChatCompletionResponse()
+ normalized_response.choices[0].message.content = response.message.content[
+ 0
+ ].text
+ return normalized_response

guides/README.md

@@ -2,10 +2,11 @@
 
 These guides give directions for obtaining API keys from different providers. 
 
-Here're the instructions for:
+Here are the instructions for:
 - [Anthropic](anthropic.md) 
 - [AWS](aws.md)
 - [Azure](azure.md) 
+- [Cohere](cohere.md)
 - [Google](google.md)
 - [Hugging Face](huggingface.md)
 - [OpenAI](openai.md)

guides/cohere.md

@@ -0,0 +1,44 @@
+# Cohere
+
+To use Cohere with `aisuite`, you’ll need an [Cohere account](https://cohere.com/). After logging in, go to the [API Keys](https://dashboard.cohere.com/api-keys) section in your account settings, agree to the terms of service, connect your card, and generate a new key. Once you have your key, add it to your environment as follows:
+
+```shell
+export CO_API_KEY="your-cohere-api-key"
+```
+
+## Create a Chat Completion
+
+Install the `cohere` Python client:
+
+Example with pip:
+```shell
+pip install cohere
+```
+
+Example with poetry:
+```shell
+poetry add cohere
+```
+
+In your code:
+```python
+import aisuite as ai
+
+client = ai.Client()
+
+provider = "cohere"
+model_id = "command-r-plus-08-2024"
+
+messages = [
+ {"role": "user", "content": "Hi, how are you?"}
+]
+
+response = client.chat.completions.create(
+ model=f"{provider}:{model_id}",
+ messages=messages,
+)
+
+print(response.choices[0].message.content)
+```
+
+Happy coding! If you’d like to contribute, please read our [Contributing Guide](CONTRIBUTING.md).

poetry.lock

@@ -9,6 +9,7 @@ readme = "README.md"
 python = "^3.10"
 anthropic = { version = "^0.30.1", optional = true }
 boto3 = { version = "^1.34.144", optional = true }
+cohere = { version = "^5.12.0", optional = true }
 vertexai = { version = "^1.63.0", optional = true }
 groq = { version = "^0.9.0", optional = true }
 mistralai = { version = "^1.0.3", optional = true }
@@ -21,14 +22,15 @@ httpx = "~0.27.0"
 anthropic = ["anthropic"]
 aws = ["boto3"]
 azure = []
+cohere = ["cohere"]
 google = ["vertexai"]
 groq = ["groq"]
 huggingface = []
 mistral = ["mistralai"]
 ollama = []
 openai = ["openai"]
 watsonx = ["ibm-watsonx-ai"]
-all = ["anthropic", "aws", "google", "groq", "mistral", "openai", "watsonx"] # To install all providers
+all = ["anthropic", "aws", "google", "groq", "mistral", "openai", "cohere", "watsonx"] # To install all providers
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.7.1"

pyproject.toml

@@ -0,0 +1,46 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from aisuite.providers.cohere_provider import CohereProvider
+
+
+@pytest.fixture(autouse=True)
+def set_api_key_env_var(monkeypatch):
+ """Fixture to set environment variables for tests."""
+ monkeypatch.setenv("CO_API_KEY", "test-api-key")
+
+
+def test_cohere_provider():
+ """High-level test that the provider is initialized and chat completions are requested successfully."""
+
+ user_greeting = "Hello!"
+ message_history = [{"role": "user", "content": user_greeting}]
+ selected_model = "our-favorite-model"
+ chosen_temperature = 0.75
+ response_text_content = "mocked-text-response-from-model"
+
+ provider = CohereProvider()
+ mock_response = MagicMock()
+ mock_response.message = MagicMock()
+ mock_response.message.content = [MagicMock()]
+ mock_response.message.content[0].text = response_text_content
+
+ with patch.object(
+ provider.client,
+ "chat",
+ return_value=mock_response,
+ ) as mock_create:
+ response = provider.chat_completions_create(
+ messages=message_history,
+ model=selected_model,
+ temperature=chosen_temperature,
+ )
+
+ mock_create.assert_called_with(
+ messages=message_history,
+ model=selected_model,
+ temperature=chosen_temperature,
+ )
+
+ assert response.choices[0].message.content == response_text_content