phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(utils.py): support dynamic api key for azure_ai route

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-11 17:51:29 -07:00
parent 6305d2dbcf
commit a9cbf2209c
2 changed files with 5 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/tests/test_completion.py
View file

@ -744,10 +744,7 @@ def test_completion_claude_3_function_plus_image():
@pytest.mark.parametrize( @pytest.mark.parametrize(
"provider", "provider",
[ ["azure", "azure_ai"],
"azure",
"azure_ai"
],
) )
def test_completion_azure_mistral_large_function_calling(provider): def test_completion_azure_mistral_large_function_calling(provider):
""" """
@ -780,6 +777,7 @@ def test_completion_azure_mistral_large_function_calling(provider):
"content": "What's the weather like in Boston today in Fahrenheit?", "content": "What's the weather like in Boston today in Fahrenheit?",
} }
] ]
response = completion( response = completion(
model="{}/mistral-large-latest".format(provider), model="{}/mistral-large-latest".format(provider),
api_base=os.getenv("AZURE_MISTRAL_API_BASE"), api_base=os.getenv("AZURE_MISTRAL_API_BASE"),

6
litellm/utils.py
View file

@ -6581,6 +6581,9 @@ def get_llm_provider(
or get_secret("FIREWORKSAI_API_KEY") or get_secret("FIREWORKSAI_API_KEY")
or get_secret("FIREWORKS_AI_TOKEN") or get_secret("FIREWORKS_AI_TOKEN")
) )
elif custom_llm_provider == "azure_ai":
api_base = api_base or get_secret("AZURE_AI_API_BASE") # type: ignore
dynamic_api_key = api_key or get_secret("AZURE_AI_API_KEY")
elif custom_llm_provider == "mistral": elif custom_llm_provider == "mistral":
# mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai # mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai
api_base = ( api_base = (
@ -6599,9 +6602,6 @@ def get_llm_provider(
or get_secret("MISTRAL_AZURE_API_KEY") # for Azure AI Mistral or get_secret("MISTRAL_AZURE_API_KEY") # for Azure AI Mistral
or get_secret("MISTRAL_API_KEY") or get_secret("MISTRAL_API_KEY")
) )
elif custom_llm_provider == "azure_ai":
api_base = api_base or get_secret("AZURE_AI_API_BASE") # type: ignore
dynamic_api_key = get_secret("AZURE_AI_API_KEY")
elif custom_llm_provider == "voyage": elif custom_llm_provider == "voyage":
# voyage is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.voyageai.com/v1 # voyage is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.voyageai.com/v1
api_base = "https://api.voyageai.com/v1" api_base = "https://api.voyageai.com/v1"