proxy - add azure/command r

This commit is contained in:
Ishaan Jaff 2024-04-05 14:35:31 -07:00
parent 8df13306ea
commit 9055a071e6
3 changed files with 62 additions and 6 deletions

View file

@ -5563,6 +5563,19 @@ def get_formatted_prompt(
return prompt
def _is_non_openai_azure_model(model: str) -> bool:
try:
model_name = model.split("/", 1)[1]
if (
model_name in litellm.cohere_chat_models
or f"mistral/{model_name}" in litellm.mistral_chat_models
):
return True
except:
return False
return False
def get_llm_provider(
model: str,
custom_llm_provider: Optional[str] = None,
@ -5576,13 +5589,8 @@ 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":
model_name = model.split("/", 1)[1]
if (
model_name in litellm.cohere_chat_models
or f"mistral/{model_name}" in litellm.mistral_chat_models
):
if _is_non_openai_azure_model(model):
custom_llm_provider = "openai"
model = model_name
return model, custom_llm_provider, dynamic_api_key, api_base
if custom_llm_provider: