From f40d47942feddc9dc60483e3528008794f9fe838 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 11 Jun 2024 14:11:02 -0700 Subject: [PATCH] test(test_completion.py): add testing to make sure tool calling works as expected for 'azure_ai/' route --- litellm/tests/test_completion.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 3c3ba564ae5..bb607e679ca 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -741,8 +741,13 @@ def test_completion_claude_3_function_plus_image(): print(response) - -def test_completion_azure_mistral_large_function_calling(): +@pytest.mark.parametrize( + "provider", [ + "azure", + "azure_ai" + ] +) +def test_completion_azure_mistral_large_function_calling(provider): """ This primarily tests if the 'Function()' pydantic object correctly handles argument param passed in as a dict vs. string """ @@ -774,7 +779,7 @@ def test_completion_azure_mistral_large_function_calling(): } ] response = completion( - model="azure/mistral-large-latest", + model="{}/mistral-large-latest".format(provider), api_base=os.getenv("AZURE_MISTRAL_API_BASE"), api_key=os.getenv("AZURE_MISTRAL_API_KEY"), messages=messages,