fix(azure_ai/transformation.py): support passing api version to azure ai services endpoint

Fixes https://github.com/BerriAI/litellm/issues/7275
This commit is contained in:
Krrish Dholakia 2025-03-12 15:16:42 -07:00
parent 91949be3bd
commit c76cf6ad6c
6 changed files with 40 additions and 7 deletions

View file

@ -58,6 +58,7 @@ def get_litellm_params(
async_call: Optional[bool] = None, async_call: Optional[bool] = None,
ssl_verify: Optional[bool] = None, ssl_verify: Optional[bool] = None,
merge_reasoning_content_in_choices: Optional[bool] = None, merge_reasoning_content_in_choices: Optional[bool] = None,
api_version: Optional[str] = None,
**kwargs, **kwargs,
) -> dict: ) -> dict:
litellm_params = { litellm_params = {
@ -99,5 +100,6 @@ def get_litellm_params(
"async_call": async_call, "async_call": async_call,
"ssl_verify": ssl_verify, "ssl_verify": ssl_verify,
"merge_reasoning_content_in_choices": merge_reasoning_content_in_choices, "merge_reasoning_content_in_choices": merge_reasoning_content_in_choices,
"api_version": api_version,
} }
return litellm_params return litellm_params

View file

@ -67,6 +67,7 @@ class AzureAIStudioConfig(OpenAIConfig):
api_base: Optional[str], api_base: Optional[str],
model: str, model: str,
optional_params: dict, optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None, stream: Optional[bool] = None,
) -> str: ) -> str:
""" """
@ -92,12 +93,14 @@ class AzureAIStudioConfig(OpenAIConfig):
original_url = httpx.URL(api_base) original_url = httpx.URL(api_base)
# Extract api_version or use default # Extract api_version or use default
api_version = cast(Optional[str], optional_params.get("api_version")) api_version = cast(Optional[str], litellm_params.get("api_version"))
# Check if 'api-version' is already present # Create a new dictionary with existing params
if "api-version" not in original_url.params and api_version: query_params = dict(original_url.params)
# Add api_version to optional_params
original_url.params["api-version"] = api_version # Add api_version if needed
if "api-version" not in query_params and api_version:
query_params["api-version"] = api_version
# Add the path to the base URL # Add the path to the base URL
if "services.ai.azure.com" in api_base: if "services.ai.azure.com" in api_base:
@ -109,8 +112,7 @@ class AzureAIStudioConfig(OpenAIConfig):
api_base=api_base, ending_path="/chat/completions" api_base=api_base, ending_path="/chat/completions"
) )
# Convert optional_params to query parameters # Use the new query_params dictionary
query_params = original_url.params
final_url = httpx.URL(new_url).copy_with(params=query_params) final_url = httpx.URL(new_url).copy_with(params=query_params)
return str(final_url) return str(final_url)

View file

@ -270,6 +270,7 @@ class BaseConfig(ABC):
api_base: Optional[str], api_base: Optional[str],
model: str, model: str,
optional_params: dict, optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None, stream: Optional[bool] = None,
) -> str: ) -> str:
""" """

View file

@ -234,6 +234,7 @@ class BaseLLMHTTPHandler:
model=model, model=model,
optional_params=optional_params, optional_params=optional_params,
stream=stream, stream=stream,
litellm_params=litellm_params,
) )
data = provider_config.transform_request( data = provider_config.transform_request(

View file

@ -1162,6 +1162,7 @@ def completion( # type: ignore # noqa: PLR0915
merge_reasoning_content_in_choices=kwargs.get( merge_reasoning_content_in_choices=kwargs.get(
"merge_reasoning_content_in_choices", None "merge_reasoning_content_in_choices", None
), ),
api_version=api_version,
) )
logging.update_environment_variables( logging.update_environment_variables(
model=model, model=model,

View file

@ -159,6 +159,32 @@ def test_azure_ai_services_handler(api_base, expected_url):
assert mock_client.call_args.kwargs["url"] == expected_url assert mock_client.call_args.kwargs["url"] == expected_url
def test_azure_ai_services_with_api_version():
from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler
client = HTTPHandler()
with patch.object(client, "post") as mock_client:
try:
response = litellm.completion(
model="azure_ai/Meta-Llama-3.1-70B-Instruct",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_key="my-fake-api-key",
api_version="2024-05-01-preview",
api_base="https://litellm8397336933.services.ai.azure.com/models",
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_client.assert_called_once()
assert mock_client.call_args.kwargs["headers"]["api-key"] == "my-fake-api-key"
assert (
mock_client.call_args.kwargs["url"]
== "https://litellm8397336933.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview"
)
def test_completion_azure_ai_command_r(): def test_completion_azure_ai_command_r():
try: try:
import os import os