From a18c94162133c3219727c1fdb538d703afbb6120 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 12 Mar 2024 12:42:12 -0700 Subject: [PATCH] (fix) failing cohere test --- litellm/model_prices_and_context_window_backup.json | 13 ++++++------- litellm/utils.py | 3 +++ model_prices_and_context_window.json | 13 ++++++------- 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 55762982f83..276b74adeba 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -981,8 +981,7 @@ "litellm_provider": "gemini", "mode": "chat" }, - - "cohere_chat/command-r": { + "command-r": { "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 4096, @@ -991,35 +990,35 @@ "litellm_provider": "cohere_chat", "mode": "chat" }, - "cohere_chat/command-light": { + "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": { + "command-nightly": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, "litellm_provider": "cohere", "mode": "completion" }, - "cohere/command": { + "command": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, "litellm_provider": "cohere", "mode": "completion" }, - "cohere/command-medium-beta": { + "command-medium-beta": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, "litellm_provider": "cohere", "mode": "completion" }, - "cohere/command-xlarge-beta": { + "command-xlarge-beta": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, diff --git a/litellm/utils.py b/litellm/utils.py index 3c1bf989a7d..3749a6af7de 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5371,6 +5371,9 @@ def get_llm_provider( ## cohere elif model in litellm.cohere_models or model in litellm.cohere_embedding_models: custom_llm_provider = "cohere" + ## cohere chat models + elif model in litellm.cohere_chat_models: + custom_llm_provider = "cohere_chat" ## replicate elif model in litellm.replicate_models or (":" in model and len(model) > 64): model_parts = model.split(":") diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 55762982f83..276b74adeba 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -981,8 +981,7 @@ "litellm_provider": "gemini", "mode": "chat" }, - - "cohere_chat/command-r": { + "command-r": { "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 4096, @@ -991,35 +990,35 @@ "litellm_provider": "cohere_chat", "mode": "chat" }, - "cohere_chat/command-light": { + "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": { + "command-nightly": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, "litellm_provider": "cohere", "mode": "completion" }, - "cohere/command": { + "command": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, "litellm_provider": "cohere", "mode": "completion" }, - "cohere/command-medium-beta": { + "command-medium-beta": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015, "litellm_provider": "cohere", "mode": "completion" }, - "cohere/command-xlarge-beta": { + "command-xlarge-beta": { "max_tokens": 4096, "input_cost_per_token": 0.000015, "output_cost_per_token": 0.000015,