ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(azure): move image gen JSON helper; rename image edit finalize hook

Browse Source 
- Add image_generation/http_utils.azure_deployment_image_generation_json_body; call
  from azure.py (keeps AzureChatCompletion focused on chat).
- Rename finalize_image_edit_multipart_data to finalize_image_edit_request_data with
  docstring covering multipart and JSON POST payloads (review feedback).

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute 2026-05-05 08:49:46 +05:30
parent 4179159f0f
commit bb0e4168ad
No known key found for this signature in database
8 changed files with 37 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
litellm/llms/azure/azure.py
View File

@ -44,6 +44,7 @@ from .common_utils import (
select_azure_base_url_or_endpoint,
)
from .image_generation import get_azure_image_generation_config
from .image_generation.http_utils import azure_deployment_image_generation_json_body
class AzureOpenAIAssistantsAPIConfig:
@ -133,22 +134,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
def __init__(self) -> None:
super().__init__()
@staticmethod
def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> dict:
"""
JSON body for Azure OpenAI image generation HTTP calls.
For ``.../openai/deployments/{deployment}/images/generations``, routing uses
the deployment in the URL only; sending ``model`` in the body (especially the
deployment name) breaks some models (e.g. gpt-image-2). See LiteLLM #26316.
Provider-style URLs (e.g. ``/providers/...`` for FLUX on Azure AI) keep all
keys so non–OpenAI-deployment payloads still work.
"""
if "images/generations" in api_base and "/openai/deployments/" in api_base:
return {k: v for k, v in data.items() if k != "model"}
return data
def make_sync_azure_openai_chat_completion_request(
self,
azure_client: Union[AzureOpenAI, OpenAI],
@ -982,9 +967,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
content=json.dumps(result).encode("utf-8"),
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
)
request_json = AzureChatCompletion.azure_deployment_image_generation_json_body(
api_base, data
)
request_json = azure_deployment_image_generation_json_body(api_base, data)
return await async_handler.post(
url=api_base,
json=request_json,
@ -1104,9 +1087,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
content=json.dumps(result).encode("utf-8"),
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
)
request_json = AzureChatCompletion.azure_deployment_image_generation_json_body(
api_base, data
)
request_json = azure_deployment_image_generation_json_body(api_base, data)
return sync_handler.post(
url=api_base,
json=request_json,

2
litellm/llms/azure/image_edit/transformation.py
View File

@ -97,7 +97,7 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
return str(final_url)
def finalize_image_edit_multipart_data(
def finalize_image_edit_request_data(
self, data: dict, resolved_request_url: str
) -> dict:
return self.azure_deployment_image_edit_form_data(data, resolved_request_url)

2
litellm/llms/azure/image_generation/__init__.py
View File

@ -6,11 +6,13 @@ from litellm.llms.base_llm.image_generation.transformation import (
from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig
from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig
from .gpt_transformation import AzureGPTImageGenerationConfig
from .http_utils import azure_deployment_image_generation_json_body
__all__ = [
"AzureDallE2ImageGenerationConfig",
"AzureDallE3ImageGenerationConfig",
"AzureGPTImageGenerationConfig",
"azure_deployment_image_generation_json_body",
]

17
litellm/llms/azure/image_generation/http_utils.py Normal file
View File

@ -0,0 +1,17 @@
"""HTTP helpers for Azure OpenAI image generation (REST, not SDK)."""
def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> dict:
"""
Build the JSON body for Azure OpenAI image generation POSTs.
For ``.../openai/deployments/{deployment}/images/generations``, routing uses the
deployment in the URL only; sending ``model`` in the body (especially the deployment
name) breaks some models (e.g. gpt-image-2). See LiteLLM #26316.
Provider-style URLs (e.g. ``/providers/...`` for FLUX on Azure AI) keep all keys
so non–OpenAI-deployment payloads still work.
"""
if "images/generations" in api_base and "/openai/deployments/" in api_base:
return {k: v for k, v in data.items() if k != "model"}
return data

9
litellm/llms/base_llm/image_edit/transformation.py
View File

@ -102,12 +102,15 @@ class BaseImageEditConfig(ABC):
) -> Tuple[Dict, RequestFiles]:
pass
def finalize_image_edit_multipart_data(
def finalize_image_edit_request_data(
self, data: dict, resolved_request_url: str
) -> dict:
"""
Adjust non-file form fields after ``transform_image_edit_request`` using the
exact URL that will be used for the HTTP POST (same string as ``get_complete_url``).
Last pass on the request dict after ``transform_image_edit_request``, using the
exact URL string used for the HTTP POST (same as ``get_complete_url`` output).
The handler sends this dict as ``data=`` for multipart providers or ``json=``
for JSON-only providers; default implementation returns ``data`` unchanged.
"""
return data

4
litellm/llms/custom_httpx/llm_http_handler.py
View File

@ -5579,7 +5579,7 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
headers=headers,
)
data = image_edit_provider_config.finalize_image_edit_multipart_data(
data = image_edit_provider_config.finalize_image_edit_request_data(
data, api_base
)
@ -5680,7 +5680,7 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
headers=headers,
)
data = image_edit_provider_config.finalize_image_edit_multipart_data(
data = image_edit_provider_config.finalize_image_edit_request_data(
data, api_base
)

2
tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py
View File

@ -44,7 +44,7 @@ def test_azure_finalize_image_edit_strips_model_after_openai_transform():
api_base=litellm_params.api_base,
litellm_params=litellm_params.model_dump(exclude_none=True),
)
data_out = config.finalize_image_edit_multipart_data(data, resolved)
data_out = config.finalize_image_edit_request_data(data, resolved)
assert "model" not in data_out
assert data_out.get("prompt") == prompt
assert data_out.get("n") == 1

7
tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py
View File

@ -12,6 +12,9 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.image_generation.http_utils import (
azure_deployment_image_generation_json_body,
)
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.azure.image_generation import (
AzureDallE3ImageGenerationConfig,
@ -41,7 +44,7 @@ def test_azure_deployment_image_generation_json_body():
"images/generations?api-version=2025-04-01-preview"
)
data = {"model": "my-dep", "prompt": "x", "n": 1}
out = AzureChatCompletion.azure_deployment_image_generation_json_body(api, data)
out = azure_deployment_image_generation_json_body(api, data)
assert "model" not in out
assert out == {"prompt": "x", "n": 1}
@ -50,7 +53,7 @@ def test_azure_providers_image_generation_json_body_keeps_model():
"""Non-deployment routes (e.g. FLUX on Azure AI) keep the payload unchanged."""
api = "https://example.services.ai.azure.com/providers/blackforestlabs/v1/flux-2-pro?api-version=preview"
data = {"model": "flux.2-pro", "prompt": "x"}
out = AzureChatCompletion.azure_deployment_image_generation_json_body(api, data)
out = azure_deployment_image_generation_json_body(api, data)
assert out == data