Merge pull request #39510 from BerriAI/litellm_coerce_multipart_numeric_fields

fix(proxy): parse numeric multipart fields on /v1/images/edits back into numbers

litellm/proxy/common_utils/http_parsing_utils.py

@@ -1,10 +1,12 @@
 import json
 import re
-from collections.abc import Collection
-from typing import Any, Final
+from collections.abc import Collection, Mapping
+from types import MappingProxyType, UnionType
+from typing import Any, Final, Union, get_args, get_origin
 
 import orjson
 from fastapi import Request, UploadFile, status
+from typing_extensions import ReadOnly
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
@@ -40,6 +42,65 @@ def _is_json_content_type(content_type: str) -> bool:
     return _normalize_media_type(content_type) == "application/json"
 
 
+def _numeric_form_type(annotation: object) -> type[int] | type[float] | None:
+    """The scalar to parse an ``int``/``float``-typed field as, else ``None``."""
+    unwrapped: Final = get_args(annotation)[0] if get_origin(annotation) is ReadOnly else annotation
+    candidates: Final = (
+        tuple(arg for arg in get_args(unwrapped) if arg is not type(None))
+        if get_origin(unwrapped) in (Union, UnionType)
+        else (unwrapped,)
+    )
+    if len(candidates) != 1:
+        return None
+    if candidates[0] is int:
+        return int
+    if candidates[0] is float:
+        return float
+    return None
+
+
+def numeric_form_fields(annotations: Mapping[str, object]) -> Mapping[str, type[int] | type[float]]:
+    """
+    The numeric fields of a request schema, mapped to the scalar to parse them as.
+
+    Only a bare ``int``/``float`` or an optional one qualifies, so container and
+    literal fields are left alone and ``bool`` is excluded on purpose.
+    """
+    return MappingProxyType(
+        {
+            name: scalar
+            for name, annotation in annotations.items()
+            if (scalar := _numeric_form_type(annotation)) is not None
+        }
+    )
+
+
+def _numeric_form_value(value: object, scalar: type[int] | type[float]) -> object:
+    if not isinstance(value, str):
+        return value
+    try:
+        return scalar(value)
+    except ValueError:
+        return value
+
+
+def coerce_numeric_form_fields(
+    parsed_body: Mapping[str, object],
+    numeric_fields: Mapping[str, type[int] | type[float]],
+) -> Mapping[str, object]:
+    """
+    Parse the numeric fields of a form-encoded body back into numbers.
+
+    ``request.form()`` yields every field as a string, so a provider that puts the
+    value in a JSON body would send a string where its API requires a number. A
+    value that will not parse is left as-is for the provider to reject as before.
+    """
+    return {
+        name: _numeric_form_value(value, numeric_fields[name]) if name in numeric_fields else value
+        for name, value in parsed_body.items()
+    }
+
+
 async def _read_request_body(request: Request | None) -> dict:
     """
     Safely read the request body and parse it as JSON.

litellm/proxy/image_endpoints/endpoints.py

@@ -2,7 +2,7 @@ import asyncio
 import io
 import traceback
 from collections.abc import Sequence
-from typing import Final
+from typing import Final, get_type_hints
 
 import orjson
 from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status
@@ -16,11 +16,18 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_utils.http_parsing_utils import (
+    coerce_numeric_form_fields,
+    numeric_form_fields,
+)
 from litellm.proxy.route_llm_request import route_request
+from litellm.types.images.main import ImageEditRequestParams
 from litellm.types.llms.openai import ChatCompletionUserMessage
 
 router: Final = APIRouter()
 
+IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams))
+
 
 async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO:
     """
@@ -279,7 +286,12 @@ async def image_edit_api(
     #########################################################
     # Read request body and convert UploadFiles to BytesIO
     #########################################################
-    data: Final = await _read_request_body(request=request)
+    data: Final = dict(
+        coerce_numeric_form_fields(
+            parsed_body=await _read_request_body(request=request),
+            numeric_fields=IMAGE_EDIT_NUMERIC_FORM_FIELDS,
+        )
+    )
     image_files: Final = await batch_to_bytesio(image)
     mask_files: Final = await batch_to_bytesio(mask)
     if image_files:

tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py

@@ -1,4 +1,6 @@
+import io
 import json
+from typing import get_type_hints
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import orjson
@@ -18,9 +20,11 @@ from litellm.proxy.common_utils.http_parsing_utils import (
     _safe_get_request_parsed_body,
     _safe_get_request_query_params,
     _safe_set_request_parsed_body,
+    coerce_numeric_form_fields,
     get_form_data,
     get_request_body,
     get_tags_from_request_body,
+    numeric_form_fields,
     populate_request_with_path_params,
 )
 
@@ -1029,3 +1033,79 @@ class TestGetRequestBody:
         mock_request = MagicMock()
         mock_request.method = "GET"
         assert await get_request_body(mock_request) == {}
+
+
+class TestNumericFormFields:
+    def test_image_edit_schema_yields_only_n(self):
+        from litellm.types.images.main import ImageEditRequestParams
+
+        assert dict(numeric_form_fields(get_type_hints(ImageEditRequestParams))) == {"n": int}
+
+    def test_qualifiers_and_optionality_are_unwrapped(self):
+        from typing import Optional
+
+        from typing_extensions import Annotated, NotRequired, ReadOnly, Required, TypedDict
+
+        class Schema(TypedDict, total=False):
+            plain: int
+            optional: Optional[int]
+            piped: int | None
+            read_only: ReadOnly[int | None]
+            not_required: NotRequired[ReadOnly[int]]
+            required: Required[ReadOnly[Annotated[float, "meta"]]]
+
+        assert dict(numeric_form_fields(get_type_hints(Schema))) == {
+            "plain": int,
+            "optional": int,
+            "piped": int,
+            "read_only": int,
+            "not_required": int,
+            "required": float,
+        }
+
+    def test_non_scalar_and_bool_fields_are_skipped(self):
+        from typing import Any, Literal, Optional, Union
+
+        from typing_extensions import TypedDict
+
+        class Schema(TypedDict, total=False):
+            flag: bool
+            optional_flag: Optional[bool]
+            text: str
+            choice: Optional[Literal["high", "low"]]
+            numbers: list[int]
+            mapping: Optional[dict[str, Any]]
+            ambiguous: Union[int, str]
+
+        assert dict(numeric_form_fields(get_type_hints(Schema))) == {}
+
+
+class TestCoerceNumericFormFields:
+    numeric_fields = {"n": int, "temperature": float}
+
+    def test_numeric_strings_are_parsed(self):
+        assert coerce_numeric_form_fields(
+            parsed_body={"n": "2", "temperature": "0.5"},
+            numeric_fields=self.numeric_fields,
+        ) == {"n": 2, "temperature": 0.5}
+
+    def test_other_fields_keep_their_string_values(self):
+        result = coerce_numeric_form_fields(
+            parsed_body={"size": "1024x1024", "prompt": "2", "quality": "high"},
+            numeric_fields=self.numeric_fields,
+        )
+        assert result == {"size": "1024x1024", "prompt": "2", "quality": "high"}
+
+    def test_unparseable_value_is_left_for_the_provider_to_reject(self):
+        assert coerce_numeric_form_fields(
+            parsed_body={"n": "two", "temperature": ""},
+            numeric_fields=self.numeric_fields,
+        ) == {"n": "two", "temperature": ""}
+
+    def test_already_typed_and_non_string_values_pass_through(self):
+        buffer = io.BytesIO(b"png")
+        result = coerce_numeric_form_fields(
+            parsed_body={"n": 3, "temperature": None, "image": buffer},
+            numeric_fields=self.numeric_fields,
+        )
+        assert result == {"n": 3, "temperature": None, "image": buffer}

tests/test_litellm/proxy/image_endpoints/test_endpoints.py

@@ -5,10 +5,13 @@ from typing import Any, Dict
 
 import orjson
 import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
 from starlette.requests import Request
 from starlette.responses import Response
 
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.image_endpoints import endpoints
 
 
@@ -115,3 +118,52 @@ async def test_image_generation_prompt_rerouting(monkeypatch):
     assert captured_route_request_data["prompt"] == "sanitized prompt"
     assert "messages" not in captured_route_request_data
     assert response.headers.get("x-callback-test") == "value"
+
+
+def _image_edit_client(monkeypatch, captured: Dict[str, Any]) -> TestClient:
+    class CaptureProcessing:
+        def __init__(self, data: Dict[str, Any]) -> None:
+            captured.update(data)
+
+        async def base_process_llm_request(self, **_: Any) -> Dict[str, Any]:
+            return {"data": [{"b64_json": "aGk="}]}
+
+    monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", CaptureProcessing)
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+    monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
+
+    app = FastAPI()
+    app.include_router(endpoints.router)
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth()
+    return TestClient(app)
+
+
+def test_image_edit_multipart_n_reaches_the_provider_as_an_int(monkeypatch):
+    """A multipart `n` must not arrive as the string Starlette parsed it into."""
+    captured: Dict[str, Any] = {}
+
+    response = _image_edit_client(monkeypatch, captured).post(
+        "/v1/images/edits",
+        files={"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")},
+        data={"model": "nova-canvas", "prompt": "add a hat", "n": "2", "size": "1024x1024"},
+    )
+
+    assert response.status_code == 200
+    assert captured["n"] == 2
+    assert isinstance(captured["n"], int)
+    assert captured["size"] == "1024x1024"
+    assert captured["prompt"] == "add a hat"
+
+
+def test_image_edit_multipart_n_that_is_not_a_number_is_left_alone(monkeypatch):
+    """An unparseable `n` still reaches the provider, which rejects it as before."""
+    captured: Dict[str, Any] = {}
+
+    response = _image_edit_client(monkeypatch, captured).post(
+        "/v1/images/edits",
+        files={"image": ("tree.png", b"\x89PNG\r\n\x1a\n", "image/png")},
+        data={"model": "nova-canvas", "prompt": "add a hat", "n": "two"},
+    )
+
+    assert response.status_code == 200
+    assert captured["n"] == "two"