From 4599a0ffd3b5591ba3ffc3fec966e616a9ae23dc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 4 Sep 2024 09:28:23 -0700 Subject: [PATCH] fix get llm provider logic --- .../litellm_core_utils/get_llm_provider_logic.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 4a615acd76b..5afa966490e 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -55,13 +55,15 @@ def handle_cohere_chat_model_custom_llm_provider( if custom_llm_provider == "cohere" and model in litellm.cohere_chat_models: return model, "cohere_chat" - _custom_llm_provider, _model = model.split("/", 1) - if ( - _custom_llm_provider - and _custom_llm_provider == "cohere" - and _model in litellm.cohere_chat_models - ): - return _model, "cohere_chat" + if "/" in model: + _custom_llm_provider, _model = model.split("/", 1) + if ( + _custom_llm_provider + and _custom_llm_provider == "cohere" + and _model in litellm.cohere_chat_models + ): + return _model, "cohere_chat" + return model, custom_llm_provider