diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py new file mode 100644 index 00000000000..84a04405483 --- /dev/null +++ b/tests/unit/images/test_main.py @@ -0,0 +1,55 @@ +from datetime import datetime +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx import llm_http_handler as llm_http_handler_module +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +PNG_BYTES: Final = b"\x89PNG\r\n\x1a\nfakepng" + + +def _edit_ok(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ("openai/dall-e-2", "openai/unpriced-image-model")) +async def test_router_image_edit_bills_the_deployment_price_with_a_logger_built_before_routing( + monkeypatch: pytest.MonkeyPatch, model: str +): + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_edit_ok)) + monkeypatch.setattr(llm_http_handler_module, "get_async_httpx_client", lambda **_kwargs: client) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "priced-edit-deployment", + "litellm_params": { + "model": model, + "api_base": "https://edit.example/v1", + "api_key": "sk-test", + "output_cost_per_image": 0.5, + }, + } + ] + ) + request: Final = { + "model": "priced-edit-deployment", + "prompt": "add a hat", + "image": PNG_BYTES, + "size": "1024x1024", + "litellm_call_id": "proxy-call-id", + } + logging_obj, routed_request = litellm.utils.function_setup( + original_function="aimage_edit", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **request, + ) + + response: Final = await router.aimage_edit(**routed_request, litellm_logging_obj=logging_obj) + + assert response._hidden_params["response_cost"] == pytest.approx(0.5) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index b9818d0d347..9f361ecb19b 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -24,6 +24,8 @@ import litellm.constants from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( + MAX_JPEG_FILL_BYTES, + MAX_JPEG_HEADER_SEGMENTS, _get_exact_count_function, _get_extrapolating_count_function, _get_tiktoken_count_function, @@ -1517,26 +1519,77 @@ def test_image_dimensions_from_bytes_returns_none_for_unreadable_headers(image: assert image_dimensions_from_bytes(image) is None -def _jpeg_sof(width: int, height: int) -> bytes: - return b"\xff\xc0" + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" +def _jpeg_sof(width: int, height: int, marker: bytes = b"\xc0") -> bytes: + return b"\xff" + marker + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" + + +def _jpeg_segment(marker: bytes, payload: bytes) -> bytes: + return b"\xff" + marker + struct.pack(">H", len(payload) + 2) + payload + + +_EMPTY_JPEG_SEGMENT: Final = _jpeg_segment(b"\xe0", b"") @pytest.mark.parametrize( - "image", + "marker", + [pytest.param(bytes([marker]), id=hex(marker)) for marker in range(0xC0, 0xD0) if marker not in (0xC4, 0xC8, 0xCC)], +) +def test_image_dimensions_from_bytes_reads_every_start_of_frame_marker(marker: bytes) -> None: + assert image_dimensions_from_bytes(b"\xff\xd8" + _jpeg_sof(800, 600, marker)) == (800, 600) + + +@pytest.mark.parametrize( + "marker", [pytest.param(b"\xc4", id="dht"), pytest.param(b"\xc8", id="jpg"), pytest.param(b"\xcc", id="dac")] +) +def test_image_dimensions_from_bytes_skips_the_non_frame_markers_in_the_start_of_frame_range(marker: bytes) -> None: + lookalike: Final = _jpeg_segment(marker, struct.pack(">BHH", 8, 1, 1) + b"\x00" * 8) + + assert image_dimensions_from_bytes(b"\xff\xd8" + lookalike + _jpeg_sof(800, 600)) == (800, 600) + + +def test_image_dimensions_from_bytes_skips_each_segment_by_its_own_length() -> None: + segments: Final = ( + _jpeg_segment(b"\xe0", b"JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00"), + _jpeg_segment(b"\xe1", b"Exif\x00\x00" + _jpeg_sof(1, 1) + b"\x00" * 7), + _jpeg_segment(b"\xfe", b"c"), + _jpeg_segment(b"\xdb", b"\x00" * 65), + ) + + assert image_dimensions_from_bytes(b"\xff\xd8" + b"".join(segments) + _jpeg_sof(800, 600)) == (800, 600) + + +@pytest.mark.parametrize( + ("segments_before_frame", "expected"), [ - pytest.param(b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1025 + _jpeg_sof(800, 600), id="too-many-segments"), - pytest.param(b"\xff\xd8" + b"\xff" * 2000 + _jpeg_sof(800, 600)[1:], id="too-many-fill-bytes"), - pytest.param(b"\xff\xd8\xff\xe0\x00\x00\x02" + _jpeg_sof(800, 600), id="segment-length-below-two"), + pytest.param(MAX_JPEG_HEADER_SEGMENTS - 1, (800, 600), id="frame-is-the-last-segment-read"), + pytest.param(MAX_JPEG_HEADER_SEGMENTS, None, id="frame-is-past-the-segment-limit"), ], ) -def test_image_dimensions_from_bytes_gives_up_on_pathological_jpeg_headers(image: bytes) -> None: - assert image_dimensions_from_bytes(image) is None +def test_image_dimensions_from_bytes_reads_at_most_the_segment_limit( + segments_before_frame: int, expected: tuple[int, int] | None +) -> None: + image: Final = b"\xff\xd8" + _EMPTY_JPEG_SEGMENT * segments_before_frame + _jpeg_sof(800, 600) + + assert image_dimensions_from_bytes(image) == expected -def test_image_dimensions_from_bytes_still_reads_a_jpeg_with_many_real_segments() -> None: - image: Final = b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1000 + b"\xff" * 64 + _jpeg_sof(800, 600)[1:] +@pytest.mark.parametrize( + ("fill_bytes", "expected"), + [ + pytest.param(MAX_JPEG_FILL_BYTES - 1, (800, 600), id="marker-is-the-last-byte-read"), + pytest.param(MAX_JPEG_FILL_BYTES, None, id="marker-is-past-the-fill-byte-limit"), + ], +) +def test_image_dimensions_from_bytes_reads_at_most_the_fill_byte_limit( + fill_bytes: int, expected: tuple[int, int] | None +) -> None: + image: Final = b"\xff\xd8" + b"\xff" * fill_bytes + _jpeg_sof(800, 600)[1:] - assert image_dimensions_from_bytes(image) == (800, 600) + assert image_dimensions_from_bytes(image) == expected + + +def test_image_dimensions_from_bytes_gives_up_on_a_segment_length_below_two() -> None: + assert image_dimensions_from_bytes(b"\xff\xd8\xff\xe0\x00\x00\x02" + _jpeg_sof(800, 600)) is None @pytest.mark.parametrize( @@ -1555,7 +1608,14 @@ def test_get_image_dimensions_still_raises_for_a_truncated_header(header: bytes) "image", [ pytest.param(b"BM" + b"\x00" * 30, id="unknown-format"), - pytest.param(b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1025 + _jpeg_sof(800, 600), id="pathological-jpeg"), + pytest.param( + b"\xff\xd8" + _EMPTY_JPEG_SEGMENT * MAX_JPEG_HEADER_SEGMENTS + _jpeg_sof(800, 600), + id="too-many-jpeg-segments", + ), + pytest.param( + b"\xff\xd8" + b"\xff" * MAX_JPEG_FILL_BYTES + _jpeg_sof(800, 600)[1:], id="too-many-jpeg-fill-bytes" + ), + pytest.param(b"\xff\xd8\xff\xe0\x00\x01\x02" + _jpeg_sof(800, 600), id="jpeg-segment-length-one"), ], ) def test_get_image_dimensions_falls_back_to_the_default_size_for_a_header_it_cannot_read(image: bytes) -> None: diff --git a/tests/unit/llms/azure_ai/conftest.py b/tests/unit/llms/azure_ai/conftest.py new file mode 100644 index 00000000000..9f80641f940 --- /dev/null +++ b/tests/unit/llms/azure_ai/conftest.py @@ -0,0 +1,14 @@ +import logging +from collections.abc import Iterator + +import pytest + +from litellm._logging import verbose_logger + + +@pytest.fixture +def litellm_warnings(caplog: pytest.LogCaptureFixture) -> Iterator[pytest.LogCaptureFixture]: + verbose_logger.addHandler(caplog.handler) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + yield caplog + verbose_logger.removeHandler(caplog.handler) diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index ed547b69ab2..0517d693b32 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -419,12 +419,13 @@ def test_flux2_image_edit_measures_a_stream_reference(): assert response._hidden_params["response_cost"] == pytest.approx(rate * 1024 * 1024 + rate * 2048 * 2048) -# Billable megapixels for a lone reference as Azure's FLUX.2-pro request_meta reported them on 2026-09-25 +# Billable megapixels for a lone reference follow Azure's FLUX.2-pro request_meta as reported on 2026-09-25 @pytest.mark.parametrize( ("reference", "billed_megapixels"), ( pytest.param(_png(640, 640), 1, id="small-reference-rounds-up"), pytest.param(_png(1024, 1280), 2, id="fractional-reference-rounds-up"), + pytest.param(_png(2048, 1024), 2, id="landscape-reference-is-measured-by-area"), pytest.param(_jpeg(4032, 3024), 4, id="photo-reference-caps-at-four-megapixels"), ), ) @@ -453,7 +454,9 @@ def test_flux2_image_edit_bills_a_lone_reference_in_whole_megapixels(reference: pytest.param(_png(0, 640), id="zero-width-header"), ), ) -def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel(reference: bytes): +def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel( + reference: bytes, litellm_warnings: pytest.LogCaptureFixture +): response: Final = litellm.image_edit( model="azure_ai/FLUX.2-flex", image=[reference], @@ -467,6 +470,7 @@ def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel(r assert response._hidden_params["reference_image_pixels"] == (UNMEASURED_REFERENCE_IMAGE_PIXELS,) assert response._hidden_params["response_cost"] == pytest.approx(rate * 1024 * 1024 + rate * 1024 * 1024) + assert "billing it as one megapixel" in litellm_warnings.text def test_flux2_image_edit_still_bills_every_reference_when_one_header_reports_zero_pixels(): diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 734308c9122..91c84cbd6cf 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -405,6 +405,32 @@ def test_flux2_pro_bills_one_megapixel_for_a_size_it_cannot_measure(size: str): assert cost == pytest.approx(first) +@pytest.mark.parametrize( + ("optional_params", "megapixels"), + ( + pytest.param({"width": 1, "height": 1}, 1, id="smallest-mapped-dimensions-win"), + pytest.param({"width": 0, "height": 1024}, 2, id="zero-width-falls-back-to-size"), + pytest.param({"width": 2048, "height": 0}, 2, id="zero-height-falls-back-to-size"), + pytest.param({"width": True, "height": 1024}, 2, id="bool-width-falls-back-to-size"), + ), +) +def test_flux2_pro_bills_mapped_dimensions_only_when_both_are_positive_integers( + optional_params: dict[str, int | bool], megapixels: int +): + first, additional = _pro_megapixel_rates() + + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model="flux.2-pro", + completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]), + custom_llm_provider="azure_ai", + size="2048x1024", + optional_params=optional_params, + call_type="image_generation", + ) + + assert cost == pytest.approx(first + additional * (megapixels - 1)) + + def test_flux2_pro_bills_the_requested_size_when_the_returned_image_is_not_valid_base64(): first, additional = _pro_megapixel_rates() @@ -435,14 +461,15 @@ def test_flux2_pro_bills_the_requested_image_count_when_the_response_lists_none( assert cost == pytest.approx(first * billed_images) -def test_flux2_pro_reads_an_uppercase_size(): +@pytest.mark.parametrize("size", ("2048X1024", "2048-x-1024")) +def test_flux2_pro_reads_each_size_spelling(size: str): first, additional = _pro_megapixel_rates() cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( model="flux.2-pro", completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]), custom_llm_provider="azure_ai", - size="2048X1024", + size=size, call_type="image_generation", ) @@ -511,6 +538,7 @@ def test_flux2_flex_edit_bills_reference_pixels_on_top_of_generated_pixels(): assert _flex_edit_cost(_edit_response((3 * 1024 * 1024,))) - generated_only == pytest.approx( _flex_pixel_rate() * 3 * 1024 * 1024 ) + assert _flex_edit_cost(_edit_response((1,))) - generated_only == pytest.approx(_flex_pixel_rate() * 1024 * 1024) def test_flux2_flex_edit_reference_cost_does_not_scale_with_image_count(): @@ -535,6 +563,15 @@ def test_flux2_flex_edit_bills_only_positive_integer_reference_counts(reference_ assert _flex_edit_cost(_edit_response(reference_image_pixels)) == generated_only +@pytest.mark.parametrize("reference_image_pixels", ((0,), ("1048576",), 1048576)) +def test_flux2_flex_edit_warns_when_it_ignores_malformed_reference_counts( + reference_image_pixels: object, litellm_warnings: pytest.LogCaptureFixture +): + _flex_edit_cost(_edit_response(reference_image_pixels)) + + assert "Ignoring malformed FLUX.2 reference pixel counts" in litellm_warnings.text + + def test_flux2_flex_edit_ignores_reference_pixels_when_provider_reports_token_usage(): response: Final = ImageResponse( data=[ImageObject(b64_json="aW1n")], @@ -634,19 +671,20 @@ def test_flux2_flex_cost_prefers_deployment_input_cost_per_pixel() -> None: assert cost == pytest.approx(2e-07 * 2048 * 1024 * 2) -def test_unlisted_azure_ai_model_bills_deployment_input_cost_per_pixel() -> None: +@pytest.mark.parametrize(("width", "height"), ((1024, 1024), (1536, 1024)), ids=("whole-megapixel", "fractional")) +def test_unlisted_azure_ai_model_bills_deployment_input_cost_per_pixel(width: int, height: int) -> None: response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")]) cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( model="unlisted-flux-deployment", completion_response=response, custom_llm_provider="azure_ai", - size="1024x1024", + size=f"{width}x{height}", call_type="image_generation", model_info={"input_cost_per_pixel": 1e-07}, ) - assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2) + assert cost == pytest.approx(1e-07 * width * height * 2) def test_unlisted_azure_ai_deployment_without_a_generated_image_price_bills_nothing() -> None: