diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index dfe662cc165..d9c5bea1a3f 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -759,4 +759,4 @@ class BaseAzureLLM(BaseOpenAILLM): def _is_azure_v1_api_version(api_version: Optional[str]) -> bool: if api_version is None: return False - return api_version == "preview" or api_version == "latest" + return api_version in {"preview", "latest", "v1"} diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index d4f6e3d8fb3..114ab3603d6 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -1504,3 +1504,23 @@ def test_get_azure_ad_token_fallback_to_default_azure_credential(setup_mocks, mo # Verify the token is what we expect from our DefaultAzureCredential mock assert token == "mock-default-azure-credential-token" + + +@pytest.mark.parametrize( + "api_version,expected", + [ + ("preview", True), + ("latest", True), + ("v1", True), + (None, False), + ("2023-05-15", False), + ("2024-01-01", False), + ("", False), + ], +) +def test_is_azure_v1_api_version(api_version, expected): + """ + Test that _is_azure_v1_api_version correctly identifies v1 API versions. + """ + result = BaseAzureLLM._is_azure_v1_api_version(api_version=api_version) + assert result == expected