mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
fix(images): carry deployment pricing kwargs to logging and harden dimension parsers
- forward CustomPricingLiteLLMParams kwargs through the image_generation and aimage_edit litellm_params dicts so deployment-scoped rates reach the spend ledger (council H2: litellm_params-slot rates were dropped) - require PNG IHDR, WebP VP8 start code, VP8L signature, and known BMP DIB sizes so signature-matching malformed files fall back instead of fabricating dimensions (council H3) - make read_image_dimensions and total_reference_pixels never raise so odd streams and short tuples degrade to the generated-pixel fallback
This commit is contained in:
shrey kharbanda
parent
dc09609e55
commit
793371270b
5 changed files with 157 additions and 9 deletions
|
|
@ -18,6 +18,12 @@ _JPEG_FIRST_SEGMENT_OFFSET: Final = 2
|
|||
_JPEG_SOF_PAYLOAD_SIZE: Final = 5
|
||||
_JPEG_FILL_CHUNK: Final = 64 * 1024
|
||||
|
||||
_PNG_IHDR: Final = b"IHDR"
|
||||
_WEBP_VP8_START_CODE: Final = b"\x9d\x01\x2a"
|
||||
_WEBP_VP8L_SIGNATURE: Final = 0x2F
|
||||
_BMP_CORE_DIB_SIZE: Final = 12
|
||||
_BMP_KNOWN_DIB_SIZES: Final = frozenset({12, 40, 52, 56, 64, 108, 124})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ImageDimensions:
|
||||
|
|
@ -30,7 +36,11 @@ class ImageDimensions:
|
|||
|
||||
|
||||
def total_reference_pixels(images: Sequence[FileTypes]) -> int | None:
|
||||
measured: Final = tuple(read_image_dimensions(image) for image in images)
|
||||
"""All-or-nothing sum of measured reference pixels; never raises."""
|
||||
try:
|
||||
measured: Final = tuple(read_image_dimensions(image) for image in images)
|
||||
except Exception: # noqa: BLE001 # a billing helper must fall back, never fail the request
|
||||
return None
|
||||
dimensions: Final = tuple(size for size in measured if size is not None)
|
||||
if len(dimensions) != len(measured):
|
||||
verbose_logger.debug(
|
||||
|
|
@ -44,10 +54,13 @@ def total_reference_pixels(images: Sequence[FileTypes]) -> int | None:
|
|||
|
||||
def read_image_dimensions(image: FileTypes) -> ImageDimensions | None:
|
||||
try:
|
||||
content = image[1] if isinstance(image, tuple) else image
|
||||
if isinstance(content, (str, os.PathLike)):
|
||||
content: Final = cast(
|
||||
"IO[bytes] | bytes | str | os.PathLike[str] | None",
|
||||
image[1] if isinstance(image, tuple) else image,
|
||||
)
|
||||
if content is None or isinstance(content, (str, os.PathLike)):
|
||||
return None
|
||||
stream = BytesIO(content) if isinstance(content, bytes) else content
|
||||
stream = BytesIO(content) if isinstance(content, (bytes, bytearray, memoryview)) else content
|
||||
if not stream.seekable():
|
||||
return None
|
||||
position = stream.tell()
|
||||
|
|
@ -58,13 +71,15 @@ def read_image_dimensions(image: FileTypes) -> ImageDimensions | None:
|
|||
if dimensions is None or dimensions.width <= 0 or dimensions.height <= 0:
|
||||
return None
|
||||
return dimensions
|
||||
except (OSError, ValueError, AttributeError, struct.error):
|
||||
except Exception: # noqa: BLE001 # an odd stream or malformed file must fall back, never raise
|
||||
return None
|
||||
|
||||
|
||||
def _header_dimensions(stream: IO[bytes], position: int) -> ImageDimensions | None:
|
||||
stream.seek(position)
|
||||
head: Final = stream.read(_HEADER_READ_SIZE)
|
||||
head: Final = cast("bytes | None", stream.read(_HEADER_READ_SIZE))
|
||||
if head is None:
|
||||
return None
|
||||
match get_image_type(head): # pyright: ignore[reportMatchNotExhaustive] # unmatched types fall through to the None below
|
||||
case "png":
|
||||
return _png_dimensions(head)
|
||||
|
|
@ -80,7 +95,7 @@ def _header_dimensions(stream: IO[bytes], position: int) -> ImageDimensions | No
|
|||
|
||||
|
||||
def _png_dimensions(head: bytes) -> ImageDimensions | None:
|
||||
if len(head) < 24:
|
||||
if len(head) < 24 or head[12:16] != _PNG_IHDR:
|
||||
return None
|
||||
width, height = cast(tuple[int, int], struct.unpack(">II", head[16:24]))
|
||||
return ImageDimensions(width=width, height=height)
|
||||
|
|
@ -97,7 +112,10 @@ def _gif_dimensions(head: bytes) -> ImageDimensions | None:
|
|||
def _bmp_dimensions(head: bytes) -> ImageDimensions | None:
|
||||
if len(head) < 26:
|
||||
return None
|
||||
if int.from_bytes(head[14:18], "little") == 12:
|
||||
dib_size: Final = int.from_bytes(head[14:18], "little")
|
||||
if dib_size not in _BMP_KNOWN_DIB_SIZES:
|
||||
return None
|
||||
if dib_size == _BMP_CORE_DIB_SIZE:
|
||||
core_width: Final = int.from_bytes(head[18:20], "little")
|
||||
core_height: Final = int.from_bytes(head[20:22], "little")
|
||||
return ImageDimensions(width=core_width, height=core_height)
|
||||
|
|
@ -116,9 +134,13 @@ def _webp_dimensions(head: bytes) -> ImageDimensions | None:
|
|||
height=int.from_bytes(head[27:30], "little") + 1,
|
||||
)
|
||||
case b"VP8 ":
|
||||
if head[23:26] != _WEBP_VP8_START_CODE:
|
||||
return None
|
||||
width, height = cast(tuple[int, int], struct.unpack("<HH", head[26:30]))
|
||||
return ImageDimensions(width=width & 0x3FFF, height=height & 0x3FFF)
|
||||
case b"VP8L":
|
||||
if head[20] != _WEBP_VP8L_SIGNATURE:
|
||||
return None
|
||||
bits: Final = int.from_bytes(head[21:25], "little")
|
||||
return ImageDimensions(width=(bits & 0x3FFF) + 1, height=((bits >> 14) & 0x3FFF) + 1)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.types.llms.openai import ImageGenerationRequestQuality
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
LITELLM_IMAGE_VARIATION_PROVIDERS,
|
||||
CustomPricingLiteLLMParams,
|
||||
LlmProviders,
|
||||
all_litellm_params,
|
||||
)
|
||||
|
|
@ -292,6 +293,7 @@ def image_generation(
|
|||
"model_info": model_info,
|
||||
"preset_cache_key": None,
|
||||
"stream_response": {},
|
||||
**{k: v for k, v in kwargs.items() if k in CustomPricingLiteLLMParams.model_fields},
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
|
@ -885,6 +887,7 @@ def image_edit(
|
|||
**image_edit_request_params,
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"model_info": model_info,
|
||||
**{k: v for k, v in kwargs.items() if k in CustomPricingLiteLLMParams.model_fields},
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -129,7 +129,24 @@ def _dims(width: int, height: int) -> ImageDimensions:
|
|||
pytest.param(_webp_chunk(b"VP8X", VP8X_BODY), _dims(700, 350), id="webp-vp8x-extended"),
|
||||
pytest.param(_webp_chunk(b"VP8Z", bytes(10)), None, id="webp-unknown-fourcc"),
|
||||
pytest.param(_webp_chunk(b"VP8X", VP8X_BODY)[:29], None, id="webp-shorter-than-30-bytes"),
|
||||
pytest.param(
|
||||
_webp_chunk(b"VP8 ", bytes(3) + bytes(3) + struct.pack("<HH", 0xC000 | 500, 0xC000 | 250) + bytes(4)),
|
||||
None,
|
||||
id="webp-vp8-missing-start-code",
|
||||
),
|
||||
pytest.param(
|
||||
_webp_chunk(b"VP8L", bytes(1) + (299 | (199 << 14)).to_bytes(4, "little") + bytes(5)),
|
||||
None,
|
||||
id="webp-vp8l-missing-signature",
|
||||
),
|
||||
pytest.param(b"\x89PNG\r\n\x1a\n" + bytes(10), None, id="png-shorter-than-24-bytes"),
|
||||
pytest.param(
|
||||
b"\x89PNG\r\n\x1a\n"
|
||||
+ _png_chunk(b"CgBI", struct.pack(">II", 4000, 3000) + bytes(8))
|
||||
+ _png_chunk(b"IHDR", struct.pack(">IIBBBBB", 1024, 1024, 8, 0, 0, 0, 0)),
|
||||
None,
|
||||
id="png-cgbi-first-chunk",
|
||||
),
|
||||
pytest.param(_gif(b"GIF89a", 1024, 768), _dims(1024, 768), id="gif89a"),
|
||||
pytest.param(_gif(b"GIF87a", 320, 200), _dims(320, 200), id="gif87a"),
|
||||
pytest.param(b"GIF89a" + bytes(2), None, id="gif-shorter-than-10-bytes"),
|
||||
|
|
@ -138,6 +155,11 @@ def _dims(width: int, height: int) -> ImageDimensions:
|
|||
pytest.param(_bmp(-1024, 768), None, id="bmp-negative-width"),
|
||||
pytest.param(_bmp(0, 768), None, id="bmp-zero-width"),
|
||||
pytest.param(b"BM" + bytes(18), None, id="bmp-shorter-than-26-bytes"),
|
||||
pytest.param(
|
||||
b"BM" + bytes(12) + struct.pack("<I", 999) + struct.pack("<ii", 640, 480) + bytes(6),
|
||||
None,
|
||||
id="bmp-unknown-dib-size",
|
||||
),
|
||||
pytest.param(
|
||||
b"BM" + bytes(12) + struct.pack("<I", 12) + struct.pack("<HH", 640, 480) + bytes(4),
|
||||
_dims(640, 480),
|
||||
|
|
@ -231,6 +253,25 @@ class _ReadOnlyObject:
|
|||
return b""
|
||||
|
||||
|
||||
class _NonBlockingStream(io.BytesIO):
|
||||
def read(self, size: int = -1) -> bytes | None:
|
||||
return None
|
||||
|
||||
|
||||
def test_read_image_dimensions_returns_none_when_read_returns_none():
|
||||
assert read_image_dimensions(cast(Any, _NonBlockingStream(_png(64, 64)))) is None
|
||||
|
||||
|
||||
def test_read_image_dimensions_returns_none_for_short_and_none_tuples():
|
||||
assert read_image_dimensions(("ref.png",)) is None
|
||||
assert read_image_dimensions(("ref.png", None)) is None
|
||||
|
||||
|
||||
def test_total_reference_pixels_returns_none_instead_of_raising():
|
||||
assert total_reference_pixels([("ref.png",)]) is None
|
||||
assert total_reference_pixels(cast(Any, 123)) is None
|
||||
|
||||
|
||||
def test_total_reference_pixels_returns_none_for_closed_and_read_only_streams():
|
||||
closed: Final = io.BytesIO(_png(64, 64))
|
||||
closed.close()
|
||||
|
|
|
|||
|
|
@ -268,6 +268,55 @@ class TestImageEditCustomPricing:
|
|||
litellm_params = {"litellm_call_id": "test-call-id"}
|
||||
assert use_custom_pricing_for_model(litellm_params) is False
|
||||
|
||||
def test_image_edit_forwards_custom_pricing_kwargs_to_logging(self):
|
||||
"""Pricing a deployment declares under litellm_params (e.g. input_cost_per_pixel)
|
||||
reaches self.litellm_params so use_custom_pricing_for_model fires on image edits."""
|
||||
from litellm.images.main import image_edit
|
||||
|
||||
captured_litellm_params = {}
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {}
|
||||
|
||||
original_update = mock_logging_obj.update_from_kwargs
|
||||
|
||||
def capturing_update(**update_kwargs):
|
||||
captured_litellm_params.update(update_kwargs.get("litellm_params", {}))
|
||||
return original_update(**update_kwargs)
|
||||
|
||||
mock_logging_obj.update_from_kwargs = capturing_update
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.images.main.get_llm_provider",
|
||||
return_value=("test-model", "openai", None, None),
|
||||
),
|
||||
patch(
|
||||
"litellm.images.main.ProviderConfigManager.get_provider_image_edit_config",
|
||||
return_value=MagicMock(),
|
||||
),
|
||||
patch(
|
||||
"litellm.images.main._get_ImageEditRequestUtils",
|
||||
return_value=MagicMock(
|
||||
get_requested_image_edit_optional_param=MagicMock(return_value={}),
|
||||
get_optional_params_image_edit=MagicMock(return_value={}),
|
||||
),
|
||||
),
|
||||
patch("litellm.images.main.base_llm_http_handler") as mock_handler,
|
||||
):
|
||||
mock_handler.image_edit_handler.return_value = MagicMock()
|
||||
|
||||
image_edit(
|
||||
image=b"fake-image-data",
|
||||
prompt="test prompt",
|
||||
model="openai/test-model",
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
input_cost_per_pixel=4.76837158203125e-08,
|
||||
)
|
||||
|
||||
assert captured_litellm_params["input_cost_per_pixel"] == pytest.approx(4.76837158203125e-08)
|
||||
assert use_custom_pricing_for_model(captured_litellm_params) is True
|
||||
|
||||
|
||||
class TestImageEditHandlerCredentialsForwarding:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -10,9 +10,9 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.images.main import image_generation
|
||||
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
|
||||
|
||||
|
||||
class TestImageGenerationExtraHeaders:
|
||||
|
|
@ -73,3 +73,36 @@ class TestImageGenerationExtraHeaders:
|
|||
)
|
||||
|
||||
assert "extra_headers" not in optional_params
|
||||
|
||||
@patch("litellm.images.main.openai_chat_completions")
|
||||
def test_image_generation_forwards_custom_pricing_kwargs_to_logging(
|
||||
self, mock_openai_chat_completions
|
||||
):
|
||||
"""Pricing a deployment declares under litellm_params (e.g. input_cost_per_pixel)
|
||||
reaches self.litellm_params so use_custom_pricing_for_model fires on generations."""
|
||||
mock_openai_chat_completions.image_generation.return_value = litellm.utils.ImageResponse(
|
||||
created=1234567890,
|
||||
data=[{"url": "https://example.com/image.png"}],
|
||||
)
|
||||
captured_litellm_params = {}
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {}
|
||||
|
||||
original_update = mock_logging_obj.update_from_kwargs
|
||||
|
||||
def capturing_update(**update_kwargs):
|
||||
captured_litellm_params.update(update_kwargs.get("litellm_params", {}))
|
||||
return original_update(**update_kwargs)
|
||||
|
||||
mock_logging_obj.update_from_kwargs = capturing_update
|
||||
|
||||
image_generation(
|
||||
model="openai/dall-e-3",
|
||||
prompt="A red circle",
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
input_cost_per_pixel=4.76837158203125e-08,
|
||||
)
|
||||
|
||||
assert captured_litellm_params["input_cost_per_pixel"] == pytest.approx(4.76837158203125e-08)
|
||||
assert use_custom_pricing_for_model(captured_litellm_params) is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue