diff --git a/litellm/__init__.py b/litellm/__init__.py index cb07c6a389..2ed12e7f02 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -260,6 +260,7 @@ open_ai_chat_completion_models: List = [] open_ai_text_completion_models: List = [] cohere_models: List = [] cohere_chat_models: List = [] +mistral_chat_models: List = [] anthropic_models: List = [] openrouter_models: List = [] vertex_language_models: List = [] @@ -284,6 +285,8 @@ for key, value in model_cost.items(): cohere_models.append(key) elif value.get("litellm_provider") == "cohere_chat": cohere_chat_models.append(key) + elif value.get("litellm_provider") == "mistral": + mistral_chat_models.append(key) elif value.get("litellm_provider") == "anthropic": anthropic_models.append(key) elif value.get("litellm_provider") == "openrouter": diff --git a/litellm/utils.py b/litellm/utils.py index d642fbbb95..9da560947a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5575,13 +5575,15 @@ def get_llm_provider( # AZURE AI-Studio Logic - Azure AI Studio supports AZURE/Cohere # If User passes azure/command-r-plus -> we should send it to cohere_chat/command-r-plus - if ( - model.split("/", 1)[0] == "azure" - and model.split("/", 1)[1] in litellm.cohere_chat_models - ): - custom_llm_provider = "openai" - model = model.split("/", 1)[1] - return model, custom_llm_provider, dynamic_api_key, api_base + if model.split("/", 1)[0] == "azure": + model_name = model.split("/", 1)[1] + if ( + model_name in litellm.cohere_chat_models + or model_name in litellm.mistral_chat_models + ): + custom_llm_provider = "openai" + model = model_name + return model, custom_llm_provider, dynamic_api_key, api_base if custom_llm_provider: return model, custom_llm_provider, dynamic_api_key, api_base