From aa8b5e9768669ba5af0e8f219e5a78306fa5c3d0 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 12 Mar 2024 10:51:33 -0700
Subject: [PATCH] (feat) add cohere_chat to model_prices

---
 ...odel_prices_and_context_window_backup.json | 32 ++++++++++++-------
 litellm/tests/test_completion.py              |  1 +
 model_prices_and_context_window.json          | 32 ++++++++++++-------
 3 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 18c4b0d9a0f..55762982f83 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -981,35 +981,45 @@
         "litellm_provider": "gemini",
         "mode": "chat"
     },
-    "command-nightly": {
+
+    "cohere_chat/command-r": {
+        "max_tokens": 128000, 
+        "max_input_tokens": 128000,
+        "max_output_tokens": 4096,
+        "input_cost_per_token": 0.00000050,
+        "output_cost_per_token": 0.0000015,
+        "litellm_provider": "cohere_chat",
+        "mode": "chat"
+    },
+    "cohere_chat/command-light": {
+        "max_tokens": 4096,
+        "input_cost_per_token": 0.000015,
+        "output_cost_per_token": 0.000015,
+        "litellm_provider": "cohere_chat",
+        "mode": "chat"
+    },
+    "cohere/command-nightly": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
         "litellm_provider": "cohere",
         "mode": "completion"
     },
-     "command": {
+     "cohere/command": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
         "litellm_provider": "cohere",
         "mode": "completion"
     },
-     "command-light": {
+     "cohere/command-medium-beta": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
         "litellm_provider": "cohere",
         "mode": "completion"
     },
-     "command-medium-beta": {
-        "max_tokens": 4096,
-        "input_cost_per_token": 0.000015,
-        "output_cost_per_token": 0.000015,
-        "litellm_provider": "cohere",
-        "mode": "completion"
-    },
-     "command-xlarge-beta": {
+     "cohere/command-xlarge-beta": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index b298cec4abd..978007c6077 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -1960,6 +1960,7 @@ def test_completion_cohere():
         pytest.fail(f"Error occurred: {e}")
 
 
+# FYI - cohere_chat looks quite unstable, even when testing locally
 def test_chat_completion_cohere():
     try:
         litellm.set_verbose = True
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 18c4b0d9a0f..55762982f83 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -981,35 +981,45 @@
         "litellm_provider": "gemini",
         "mode": "chat"
     },
-    "command-nightly": {
+
+    "cohere_chat/command-r": {
+        "max_tokens": 128000, 
+        "max_input_tokens": 128000,
+        "max_output_tokens": 4096,
+        "input_cost_per_token": 0.00000050,
+        "output_cost_per_token": 0.0000015,
+        "litellm_provider": "cohere_chat",
+        "mode": "chat"
+    },
+    "cohere_chat/command-light": {
+        "max_tokens": 4096,
+        "input_cost_per_token": 0.000015,
+        "output_cost_per_token": 0.000015,
+        "litellm_provider": "cohere_chat",
+        "mode": "chat"
+    },
+    "cohere/command-nightly": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
         "litellm_provider": "cohere",
         "mode": "completion"
     },
-     "command": {
+     "cohere/command": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
         "litellm_provider": "cohere",
         "mode": "completion"
     },
-     "command-light": {
+     "cohere/command-medium-beta": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,
         "litellm_provider": "cohere",
         "mode": "completion"
     },
-     "command-medium-beta": {
-        "max_tokens": 4096,
-        "input_cost_per_token": 0.000015,
-        "output_cost_per_token": 0.000015,
-        "litellm_provider": "cohere",
-        "mode": "completion"
-    },
-     "command-xlarge-beta": {
+     "cohere/command-xlarge-beta": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000015,
         "output_cost_per_token": 0.000015,