fix(fal_ai): price non-canonical image sizes from the nearest row and honour dump options (#42336)

* fix(fal_ai): price non-canonical image sizes from the nearest row and honour dump options

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(fal_ai): drop monkeypatched mixed pricing case

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(fal_ai): use the default dimensions constant directly

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(fal_ai): forward nested include and exclude when dumping image data

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(fal_ai): honour pydantic item selectors in image data serializer

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(fal_ai): match negative item selectors in image data serializer

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: kerry <kerry@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/fal_ai/cost_calculator.py

@@ -9,7 +9,8 @@ import litellm
 from litellm.types.utils import ImageObject, ImageResponse
 
 FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high"
-FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768"
+_DEFAULT_KEYED_DIMENSIONS: Final[tuple[int, int]] = (1024, 768)
+FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = f"{_DEFAULT_KEYED_DIMENSIONS[0]}-x-{_DEFAULT_KEYED_DIMENSIONS[1]}"
 FAL_PIXELS_PER_MEGAPIXEL: Final[int] = 1_048_576
 FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType(
     {
@@ -55,37 +56,56 @@ def _image_dimensions(image: object) -> tuple[int, int] | None:
     return width, height
 
 
-def _response_size(image: object) -> str | None:
-    dimensions: Final = _image_dimensions(image)
-    if dimensions is None:
-        return None
-    width, height = dimensions
-    return f"{width}-x-{height}"
-
-
 def _keyed_quality(optional_params: Mapping[str, object]) -> str:
     raw_quality: Final = optional_params.get("quality")
     return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY
 
 
+def _parse_keyed_dimensions(size: str | None) -> tuple[int, int] | None:
+    if size is None:
+        return None
+    parts: Final = tuple(size.split("-x-"))
+    if len(parts) != 2:
+        return None
+    try:
+        width, height = (int(part) for part in parts)
+    except ValueError:
+        return None
+    return (width, height) if width > 0 and height > 0 else None
+
+
+def _keyed_rows(model: str, quality: str) -> tuple[tuple[int, int, float], ...]:
+    prefix: Final = f"fal_ai/{quality}/"
+    suffix: Final = f"/{model}"
+    return tuple(
+        (width, height, float(raw_cost))
+        for key in litellm.model_cost
+        if isinstance(key, str) and key.startswith(prefix) and key.endswith(suffix)
+        for size in (key[len(prefix) : -len(suffix)],)
+        for dimensions in (_parse_keyed_dimensions(size),)
+        if dimensions is not None
+        for entry in (_entry(key),)
+        if entry is not None
+        for raw_cost in (entry.get("output_cost_per_image"),)
+        if isinstance(raw_cost, (int, float))
+        for width, height in (dimensions,)
+    )
+
+
 def _keyed_cost_per_image(
     model: str,
     image: object,
     optional_params: Mapping[str, object],
 ) -> float | None:
     quality: Final = _keyed_quality(optional_params)
-    request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE
-    sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE)
-    for size in sizes:
-        if size is None:
-            continue
-        keyed_entry = _entry(f"fal_ai/{quality}/{size}/{model}")
-        if keyed_entry is None:
-            continue
-        keyed_cost = keyed_entry.get("output_cost_per_image")
-        if isinstance(keyed_cost, (int, float)):
-            return float(keyed_cost)
-    return None
+    rows: Final = _keyed_rows(model, quality)
+    if not rows:
+        return None
+    target_dimensions: Final = (
+        _image_dimensions(image) or _parse_keyed_dimensions(_keyed_size(optional_params)) or _DEFAULT_KEYED_DIMENSIONS
+    )
+    target_pixels: Final = target_dimensions[0] * target_dimensions[1]
+    return min(rows, key=lambda row: (abs(row[0] * row[1] - target_pixels), row[0] * row[1]))[2]
 
 
 def _flat_cost_per_image(
@@ -129,7 +149,7 @@ def cost_calculator(
         )
         for image in images
     )
-    if all(cost is not None for cost in keyed_costs):
+    if not any(cost is None for cost in keyed_costs):
         return sum(cost for cost in keyed_costs if cost is not None)
     model_info: Final = litellm.get_model_info(
         model=normalized_model,
@@ -144,10 +164,12 @@ def cost_calculator(
         float(raw_output_cost_per_pixel) if isinstance(raw_output_cost_per_pixel, (int, float)) else None
     )
     return sum(
-        _flat_cost_per_image(
+        keyed_cost
+        if keyed_cost is not None
+        else _flat_cost_per_image(
             image=image,
             output_cost_per_image=output_cost_per_image,
             output_cost_per_pixel=output_cost_per_pixel,
         )
-        for image in images
+        for image, keyed_cost in zip(images, keyed_costs)
     )

litellm/types/utils.py

@@ -1,7 +1,7 @@
 import json
 import re
 import time
-from collections.abc import Mapping, Sequence
+from collections.abc import Collection, Mapping, Sequence
 from enum import Enum
 from types import MappingProxyType
 from typing import (
@@ -36,12 +36,14 @@ from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
+    FieldSerializationInfo,
     JsonValue,
     PrivateAttr,
     SkipValidation,
     field_serializer,
     field_validator,
 )
+from pydantic.main import IncEx
 from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
 
 from litellm._logging import verbose_logger
@@ -80,6 +82,27 @@ from .llms.openai import (
 )
 from .rerank import RerankResponse as RerankResponse
 
+
+def _nested_selector(
+    selector: IncEx | None,
+    index: int,
+    count: int,
+    is_include: bool,
+) -> tuple[bool, IncEx | None]:
+    if selector is None:
+        return True, None
+    if isinstance(selector, Mapping):
+        value: Final = selector.get(index, selector.get(index - count, selector.get("__all__")))
+        keep: Final = value is not None if is_include else value is not True
+        per_item_selector: Final = None if value is True or value is None else value
+        return keep, per_item_selector
+    if isinstance(selector, Collection) and not isinstance(selector, (str, bytes)):
+        if all(isinstance(item, int) for item in selector):
+            addressed: Final = index in selector or index - count in selector
+            return (addressed if is_include else not addressed), None
+    return True, selector
+
+
 if TYPE_CHECKING:
     from .vector_stores import VectorStoreSearchResponse
 else:
@@ -2557,8 +2580,35 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject):
     model_config = ConfigDict(extra="allow", protected_namespaces=())
 
     @field_serializer("data")
-    def _serialize_image_data(self, data: Sequence[OpenAIImage] | None) -> Sequence[Mapping[str, object]] | None:
-        return None if data is None else [image.model_dump() for image in data]
+    def _serialize_image_data(
+        self,
+        data: Sequence[OpenAIImage] | None,
+        info: FieldSerializationInfo,
+    ) -> Sequence[Mapping[str, object]] | None:
+        if data is None:
+            return None
+        include: Final = info.include
+        exclude: Final = info.exclude
+
+        def _serialize_image(index: int, image: OpenAIImage) -> Mapping[str, object] | None:
+            include_keep, include_selector = _nested_selector(include, index, len(data), is_include=True)
+            exclude_keep, exclude_selector = _nested_selector(exclude, index, len(data), is_include=False)
+            if not include_keep or not exclude_keep:
+                return None
+            return image.model_dump(
+                mode=info.mode,
+                include=include_selector,
+                exclude=exclude_selector,
+                context=info.context,
+                exclude_none=info.exclude_none,
+                exclude_unset=info.exclude_unset,
+                exclude_defaults=info.exclude_defaults,
+                round_trip=info.round_trip,
+                by_alias=info.by_alias,
+            )
+
+        serialized_images: Final = tuple(_serialize_image(index, image) for index, image in enumerate(data))
+        return [image for image in serialized_images if image is not None]
 
     def __init__(
         self,

tests/integration/contracts.json

@@ -172,6 +172,12 @@
     "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_row": [
       "other.provider_wire.fal_ai.gpt_image_generation_quality_size_wire_and_keyed_pricing"
     ],
+    "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_generation_prices_non_canonical_size_from_nearest_row": [
+      "other.provider_wire.fal_ai.gpt_image_generation_noncanonical_size_uses_nearest_keyed_row"
+    ],
+    "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_sdk_response_honors_dump_options": [
+      "other.provider_wire.fal_ai.sdk_image_response_dump_options"
+    ],
     "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image": [
       "other.provider_wire.fal_ai.flux_dev_endpoint_and_per_image_pricing"
     ],

tests/integration/providers/test_fal_ai_image_wire.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Final
 
 import httpx
+import litellm
 import pytest
 from integration._support.client import Gateway
 from integration._support.wire import Reply, Request, wire_server
@@ -120,6 +121,63 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro
         ]
 
 
+@pytest.mark.covers("other.provider_wire.fal_ai.gpt_image_generation_noncanonical_size_uses_nearest_keyed_row")
+def test_fal_gpt_image_25_generation_prices_non_canonical_size_from_nearest_row(gateway: Gateway) -> None:
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.headers["authorization"] == "Key synthetic-fal-key"
+        assert request.target == "/openai/gpt-image-2.5/flare/text-to-image"
+        body: Final = _JSON_OBJECT.validate_json(request.body)
+        assert body == {"prompt": _PROMPT, "quality": "low", "image_size": {"width": 1536, "height": 1024}}
+        return Reply(body=_image_response(((f"{wire_url}/files/noncanonical.png", 1536, 1024),), _PROMPT))
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        wire_url: Final = wire.url
+        model: Final = scenario.model(
+            model=f"fal_ai/{_GPT_IMAGE_MODEL}", api_base=wire.url, api_key="synthetic-fal-key"
+        )
+        response: Final = gateway.request(
+            "POST",
+            "/v1/images/generations",
+            {"model": model, "prompt": _PROMPT, "quality": "low", "size": "1536x1024"},
+        )
+        assert response.status_code == 200, response.text
+        cost: Final = _response_cost(response)
+        assert cost == _approx(_catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image"))
+        assert [(request.method, request.target) for request in wire.drain()] == [
+            ("POST", "/openai/gpt-image-2.5/flare/text-to-image")
+        ]
+
+
+@pytest.mark.covers("other.provider_wire.fal_ai.sdk_image_response_dump_options")
+def test_fal_gpt_image_sdk_response_honors_dump_options() -> None:
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.headers["authorization"] == "Key synthetic-fal-key"
+        assert request.target == "/openai/gpt-image-2.5/flare/text-to-image"
+        assert _JSON_OBJECT.validate_json(request.body) == {"prompt": _PROMPT, "quality": "low"}
+        return Reply(body=_image_response((("https://example.com/fal.png", 1024, 1536),), _PROMPT))
+
+    with wire_server(respond) as wire:
+        response: Final = litellm.image_generation(
+            model=_GPT_IMAGE_MODEL,
+            prompt=_PROMPT,
+            quality="low",
+            api_base=wire.url,
+            api_key="synthetic-fal-key",
+            custom_llm_provider="fal_ai",
+        )
+        assert response.model_dump(exclude_none=True)["data"] == [
+            {
+                "url": "https://example.com/fal.png",
+                "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"},
+            }
+        ]
+        assert [(request.method, request.target) for request in wire.drain()] == [
+            ("POST", "/openai/gpt-image-2.5/flare/text-to-image")
+        ]
+
+
 @pytest.mark.covers("other.provider_wire.fal_ai.flux_dev_endpoint_and_per_image_pricing")
 def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gateway: Gateway) -> None:
     def respond(request: Request) -> Reply:

tests/test_litellm/llms/fal_ai/test_cost_calculator.py

@@ -1,3 +1,5 @@
+from typing import Final
+
 import pytest
 
 import litellm
@@ -5,7 +7,6 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
 from litellm.llms.fal_ai.cost_calculator import cost_calculator
 from litellm.types.utils import ImageObject, ImageResponse
 
-
 @pytest.fixture(autouse=True)
 def _use_local_model_cost_map(monkeypatch):
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
@@ -78,14 +79,27 @@ def test_gpt_image_response_dimensions_override_request_size():
     assert cost == expected
 
 
-def test_gpt_image_response_dimensions_fall_back_to_request_size_when_unpriced():
+def test_gpt_image_response_dimensions_use_nearest_keyed_row_when_unpriced():
     model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image"
     cost = cost_calculator(
         model=model,
         image_response=_image_response_with_dimensions(((777, 888),)),
         optional_params={"quality": "low", "image_size": {"width": 1024, "height": 1536}},
     )
-    expected = litellm.model_cost[f"fal_ai/low/1024-x-1536/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"]
+    expected = litellm.model_cost[f"fal_ai/low/1024-x-768/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"]
+    assert cost == expected
+
+
+def test_gpt_image_25_noncanonical_response_uses_nearest_keyed_row():
+    model: Final = "fal_ai/openai/gpt-image-2.5/flare/text-to-image"
+    cost: Final = cost_calculator(
+        model=model,
+        image_response=_image_response_with_dimensions(((1536, 1024),)),
+        optional_params={"quality": "low", "image_size": {"width": 1536, "height": 1024}},
+    )
+    expected: Final = litellm.model_cost[
+        "fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image"
+    ]["output_cost_per_image"]
     assert cost == expected
 
 

tests/test_litellm/types/test_types_utils.py

@@ -1,9 +1,16 @@
+import json
 from typing import Final
 
 import pytest
 
 
-from litellm.types.utils import HiddenParams, all_litellm_params, text_tokens_without_nested_reasoning
+from litellm.types.utils import (
+    HiddenParams,
+    ImageObject,
+    ImageResponse,
+    all_litellm_params,
+    text_tokens_without_nested_reasoning,
+)
 
 
 def test_rust_is_a_known_litellm_param():
@@ -763,13 +770,70 @@ def test_delta_function_tool_call_unchanged_by_custom_support():
 
 def test_image_response_keeps_background():
     """https://github.com/BerriAI/litellm/issues/38649"""
-    from litellm.types.utils import ImageResponse
-
     response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png")
     assert response.background == "transparent"
     assert response.model_dump()["background"] == "transparent"
 
 
+def test_image_response_serialization_honors_dump_options():
+    response: Final = ImageResponse(
+        data=[
+            ImageObject(
+                url="https://example.com/image.png",
+                provider_specific_fields={"width": 1024, "height": 1536, "content_type": "image/png"},
+            )
+        ]
+    )
+    expected: Final = [
+        {
+            "url": "https://example.com/image.png",
+            "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"},
+        }
+    ]
+    assert response.model_dump(exclude_none=True)["data"] == expected
+    assert json.loads(response.model_dump_json(exclude_none=True))["data"] == expected
+    assert response.model_dump()["data"][0]["provider_specific_fields"] == expected[0]["provider_specific_fields"]
+    assert "url" not in response.model_dump(exclude={"data": {0: {"url"}}})["data"][0]
+    assert response.model_dump(include={"data": {"__all__": {"url"}}})["data"] == [
+        {"url": "https://example.com/image.png"}
+    ]
+    assert response.model_dump(include={"data": {0: True}})["data"] == [
+        {
+            "b64_json": None,
+            "revised_prompt": None,
+            "url": "https://example.com/image.png",
+            "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"},
+        }
+    ]
+    assert response.model_dump(exclude={"data": {0: True}})["data"] == []
+
+    two_image_response: Final = ImageResponse(
+        data=[
+            ImageObject(url="https://example.com/image.png"),
+            ImageObject(url="https://example.com/second-image.png"),
+        ]
+    )
+    assert two_image_response.model_dump(exclude={"data": {1}})["data"] == [
+        {
+            "b64_json": None,
+            "revised_prompt": None,
+            "url": "https://example.com/image.png",
+            "provider_specific_fields": None,
+        }
+    ]
+    assert two_image_response.model_dump(exclude={"data": {-1}})["data"] == [
+        {
+            "b64_json": None,
+            "revised_prompt": None,
+            "url": "https://example.com/image.png",
+            "provider_specific_fields": None,
+        }
+    ]
+    assert two_image_response.model_dump(include={"data": {-1: {"url"}}})["data"] == [
+        {"url": "https://example.com/second-image.png"}
+    ]
+
+
 @pytest.mark.parametrize(
     ("completion_tokens", "text_tokens", "reasoning_tokens", "other_modality_tokens", "expected_text_tokens"),
     (