scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

Fix: base_model name for body and deplyment name in URL

Browse source
This commit is contained in:
Sameer Kankute 2026-02-09 16:25:46 +05:30
parent 4f96a3b126
commit d35691aa0c
2 changed files with 214 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/llms/azure/azure.py
View file

@ -1060,6 +1060,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
headers: dict,
client=None,
timeout=None,
model: Optional[str] = None,
) -> ImageResponse:
response: Optional[dict] = None
@ -1071,8 +1072,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if api_base.endswith("/"):
api_base = api_base.rstrip("/")
api_version: str = azure_client_params.get("api_version", "")
# Use the deployment name (model) for URL construction, not the base_model from data
img_gen_api_base = self.create_azure_base_url(
azure_client_params=azure_client_params, model=data.get("model", "")
azure_client_params=azure_client_params, model=model or data.get("model", "")
)
## LOGGING
@ -1159,21 +1161,18 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
model = model
else:
model = None
## BASE MODEL CHECK
if (
model_response is not None
and optional_params.get("base_model", None) is not None
and litellm_params.get("base_model", None) is not None
):
model_response._hidden_params["model"] = optional_params.pop(
"base_model"
)
model_response._hidden_params["model"] = litellm_params.get("base_model", None)
# Azure image generation API doesn't support extra_body parameter
extra_body = optional_params.pop("extra_body", {})
flattened_params = {**optional_params, **extra_body}
data = {"model": model, "prompt": prompt, **flattened_params}
data = {"model": litellm_params.get("base_model", None) or model, "prompt": prompt, **flattened_params}
max_retries = data.pop("max_retries", 2)
if not isinstance(max_retries, int):
raise AzureOpenAIError(
@ -1196,10 +1195,11 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
is_async=False,
)
if aimg_generation is True:
return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers) # type: ignore
return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers, model=model) # type: ignore
# Use the deployment name (model) for URL construction, not the base_model from data
img_gen_api_base = self.create_azure_base_url(
azure_client_params=azure_client_params, model=data.get("model", "")
azure_client_params=azure_client_params, model=model
)
## LOGGING

205
tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
View file

@ -251,3 +251,208 @@ def test_azure_image_generation_drop_params_false_raises_error():
# Verify the error message mentions the unsupported parameter
assert "response_format" in str(exc_info.value)
def test_azure_image_generation_base_model_vs_deployment_name():
"""
Test that Azure image generation correctly uses base_model in request body
but deployment name in the URL.
When base_model is specified in litellm_params, the request should:
1. Use base_model (e.g., "gpt-image-1.5") in the JSON request body
2. Use the deployment name (e.g., "gpt-image-15") in the URL path
This is important because Azure expects:
- URL: /openai/deployments/{deployment_name}/images/generations
- Body: {"model": "{base_model}", ...}
Example config:
model: azure/gpt-image-15 # deployment name
base_model: gpt-image-1.5 # actual model name
"""
from unittest.mock import MagicMock
# Setup test parameters
azure_chat_completion = AzureChatCompletion()
prompt = "A beautiful image of a cat"
model = "gpt-image-15" # This is the deployment name
base_model = "gpt-image-1.5" # This is the actual model name
api_base = "https://openai-gpt-image-1-5-test-v-1.openai.azure.com/"
api_version = "2024-07-01-preview"
api_key = "test-api-key"
litellm_params = {
"base_model": base_model,
"api_base": api_base,
"api_version": api_version,
}
optional_params = {
"n": 1,
"size": "1024x1024"
}
# Mock the HTTP request to capture what gets sent
with patch.object(
azure_chat_completion,
"make_sync_azure_httpx_request",
return_value=MagicMock(
json=lambda: {
"created": 1234567890,
"data": [
{
"url": "https://example.com/image.png",
"revised_prompt": prompt
}
]
}
)
) as mock_request:
# Mock logging object
logging_obj = MagicMock()
logging_obj.pre_call = MagicMock()
logging_obj.post_call = MagicMock()
# Call the image_generation method
try:
response = azure_chat_completion.image_generation(
prompt=prompt,
timeout=60.0,
optional_params=optional_params,
logging_obj=logging_obj,
headers={},
model=model,
api_key=api_key,
api_base=api_base,
api_version=api_version,
litellm_params=litellm_params,
)
except Exception as e:
# If there's an error, we still want to check the mock calls
pass
# Verify the mock was called
assert mock_request.called, "HTTP request should have been made"
# Get the call arguments
call_kwargs = mock_request.call_args.kwargs
# Verify the URL uses the deployment name (not base_model)
api_base_used = call_kwargs.get("api_base", "")
assert model in api_base_used, (
f"URL should contain deployment name '{model}', "
f"but got: {api_base_used}"
)
assert base_model not in api_base_used or base_model == model, (
f"URL should NOT contain base_model '{base_model}' when it differs from deployment name, "
f"but got: {api_base_used}"
)
# Verify the request body uses base_model (not deployment name)
request_data = call_kwargs.get("data", {})
assert request_data.get("model") == base_model, (
f"Request body 'model' field should be base_model '{base_model}', "
f"but got: {request_data.get('model')}"
)
# Verify other fields are correct
assert request_data.get("prompt") == prompt
assert request_data.get("n") == 1
assert request_data.get("size") == "1024x1024"
@pytest.mark.asyncio
async def test_azure_aimage_generation_base_model_vs_deployment_name():
"""
Test that Azure async image generation correctly uses base_model in request body
but deployment name in the URL.
This is the async version of test_azure_image_generation_base_model_vs_deployment_name.
"""
from unittest.mock import MagicMock
# Setup test parameters
azure_chat_completion = AzureChatCompletion()
prompt = "A beautiful image of a cat"
model = "gpt-image-15" # This is the deployment name
base_model = "gpt-image-1.5" # This is the actual model name
api_base = "https://openai-gpt-image-1-5-test-v-1.openai.azure.com/"
api_version = "2024-07-01-preview"
api_key = "test-api-key"
data = {
"model": base_model,
"prompt": prompt,
"n": 1,
"size": "1024x1024"
}
azure_client_params = {
"api_base": api_base,
"api_version": api_version,
}
# Mock the HTTP request to capture what gets sent
with patch.object(
azure_chat_completion,
"make_async_azure_httpx_request",
new_callable=AsyncMock,
return_value=MagicMock(
json=lambda: {
"created": 1234567890,
"data": [
{
"url": "https://example.com/image.png",
"revised_prompt": prompt
}
]
}
)
) as mock_request:
# Mock logging object
logging_obj = MagicMock()
logging_obj.pre_call = MagicMock()
logging_obj.post_call = MagicMock()
# Call the aimage_generation method
try:
response = await azure_chat_completion.aimage_generation(
data=data,
model_response=None,
azure_client_params=azure_client_params,
api_key=api_key,
input=[],
logging_obj=logging_obj,
headers={},
model=model, # Pass the deployment name
timeout=60.0,
)
except Exception as e:
# If there's an error, we still want to check the mock calls
pass
# Verify the mock was called
assert mock_request.called, "HTTP request should have been made"
# Get the call arguments
call_kwargs = mock_request.call_args.kwargs
# Verify the URL uses the deployment name (not base_model)
api_base_used = call_kwargs.get("api_base", "")
assert model in api_base_used, (
f"URL should contain deployment name '{model}', "
f"but got: {api_base_used}"
)
assert base_model not in api_base_used or base_model == model, (
f"URL should NOT contain base_model '{base_model}' when it differs from deployment name, "
f"but got: {api_base_used}"
)
# Verify the request body uses base_model (not deployment name)
request_data = call_kwargs.get("data", {})
assert request_data.get("model") == base_model, (
f"Request body 'model' field should be base_model '{base_model}', "
f"but got: {request_data.get('model')}"
)