Add new Cohere models

Added `command-r` and `command-r-plus` models
jupyterlab · Jun 14, 2024 · 131a158 · 131a158
1 parent a3b3ce0
commit 131a158
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -549,8 +549,8 @@ def is_api_key_exc(cls, e: Exception):
 class CohereProvider(BaseProvider, Cohere):
     id = "cohere"
     name = "Cohere"
-    # Source: https://docs.cohere.com/reference/generate
-    models = ["command", "command-nightly", "command-light", "command-light-nightly"]
+    # Source: https://docs.cohere.com/reference/generate; https://docs.cohere.com/docs/models
+    models = ["command", "command-nightly", "command-light", "command-light-nightly", "command-r-plus", "command-r"]
     model_id_key = "model"
     pypi_package_deps = ["cohere"]
     auth_strategy = EnvAuthStrategy(name="COHERE_API_KEY")