From 48d0d7092cf9f9d13621ff55f2a907d789d40229 Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Sun, 20 Sep 2026 05:36:36 +0800 Subject: [PATCH] fix: normalize unsupported multimodal image formats --- .../core/provider/sources/openai_source.py | 4 +- astrbot/core/utils/media_utils.py | 73 +++++++++++++++++++ tests/test_media_utils.py | 35 +++++++++ tests/test_openai_source.py | 24 ++++++ 4 files changed, 135 insertions(+), 1 deletion(-) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index 4566441830..b6a7713f76 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -31,6 +31,7 @@ from astrbot.core.provider.entities import LLMResponse, TokenUsage, ToolCallsResult from astrbot.core.utils.media_utils import ( describe_media_ref, + normalize_image_for_provider, resolve_media_ref_to_base64_data, ) from astrbot.core.utils.network_utils import ( @@ -187,7 +188,8 @@ async def _image_ref_to_data_url( media_type="image", strict=mode == "strict", ) - return image_data.to_data_url() if image_data else None + normalized = normalize_image_for_provider(image_data) + return normalized.to_data_url() if normalized else None async def _resolve_image_part( self, diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index 321a6d17e6..2e5f89f88b 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -146,6 +146,79 @@ def to_data_url(self) -> str: return f"data:{self.mime_type};base64,{self.base64_data}" +IMAGE_PROVIDER_SUPPORTED_MIME_TYPES = frozenset( + { + "image/gif", + "image/jpeg", + "image/png", + } +) + + +def normalize_image_for_provider( + image_data: ResolvedMediaData | None, + supported_mimes: set[str] | frozenset[str] | None = None, +) -> ResolvedMediaData | None: + """Normalize image bytes to a MIME type accepted by a vision provider. + + Args: + image_data: Resolved image bytes and metadata. + supported_mimes: MIME types accepted by the provider. Defaults to JPEG, + PNG, and GIF. + + Returns: + Validated image data with corrected MIME metadata, or converted image data + when the source format is not supported. + """ + if image_data is None: + return None + + supported = supported_mimes or IMAGE_PROVIDER_SUPPORTED_MIME_TYPES + if image_data.mime_type in supported: + return image_data + + raw = image_data.to_bytes() + with PILImage.open(io.BytesIO(raw)) as image: + actual_mime = { + "GIF": "image/gif", + "JPEG": "image/jpeg", + "PNG": "image/png", + "WEBP": "image/webp", + }.get(str(image.format or "").upper()) + if actual_mime in supported: + return ResolvedMediaData( + base64_data=image_data.base64_data, + mime_type=actual_mime, + format=image_data.format, + ) + + has_alpha = image.mode in {"RGBA", "LA", "PA"} or "transparency" in image.info + if has_alpha and "image/png" in supported: + output_format = "PNG" + output_mime = "image/png" + converted = image.convert("RGBA") + elif "image/jpeg" in supported: + output_format = "JPEG" + output_mime = "image/jpeg" + converted = image.convert("RGB") + elif "image/png" in supported: + output_format = "PNG" + output_mime = "image/png" + converted = image.convert("RGB") + else: + return None + + try: + output = io.BytesIO() + converted.save(output, format=output_format) + return ResolvedMediaData( + base64_data=base64.b64encode(output.getvalue()).decode("utf-8"), + mime_type=output_mime, + ) + finally: + converted.close() + + @dataclass(slots=True) class _LocalMediaFile: path: Path diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index cf3198ebf4..45ad79226b 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -228,6 +228,41 @@ def test_detect_image_mime_type_sniffs_common_headers(): ) +def test_normalize_image_for_provider_converts_webp_bytes(): + from PIL import Image as PILImage + + image_buffer = BytesIO() + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(image_buffer, format="WEBP") + image_data = media_utils.ResolvedMediaData( + base64_data=base64.b64encode(image_buffer.getvalue()).decode("ascii"), + mime_type="image/webp", + ) + + normalized = media_utils.normalize_image_for_provider(image_data) + + assert normalized is not None + assert normalized.mime_type == "image/jpeg" + with PILImage.open(BytesIO(normalized.to_bytes())) as image: + assert image.format == "JPEG" + + +def test_normalize_image_for_provider_preserves_png_bytes(): + from PIL import Image as PILImage + + image_buffer = BytesIO() + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(image_buffer, format="PNG") + image_data = media_utils.ResolvedMediaData( + base64_data=base64.b64encode(image_buffer.getvalue()).decode("ascii"), + mime_type="image/png", + ) + + normalized = media_utils.normalize_image_for_provider(image_data) + + assert normalized is not None + assert normalized.mime_type == "image/png" + assert normalized.base64_data == image_data.base64_data + + def test_detect_image_mime_type_returns_default_for_unknown_input(): """Unknown or empty headers fall back to the provided default.""" assert ( diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index 871dbf3b38..0e9a444ee7 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -1164,6 +1164,30 @@ async def test_resolve_image_part_preserves_base64_png_mime_type(): await provider.terminate() +@pytest.mark.asyncio +async def test_resolve_image_part_normalizes_webp_to_jpeg(): + provider = _make_provider() + try: + image_buffer = BytesIO() + PILImage.new("RGB", (2, 2), (255, 0, 0)).save( + image_buffer, + format="WEBP", + ) + image_base64 = base64.b64encode(image_buffer.getvalue()).decode("ascii") + + image_part = await provider._resolve_image_part(f"base64://{image_base64}") + + assert image_part is not None + image_url = image_part["image_url"]["url"] + assert image_url.startswith("data:image/jpeg;base64,") + with PILImage.open( + BytesIO(base64.b64decode(image_url.split(",", 1)[1])) + ) as image: + assert image.format == "JPEG" + finally: + await provider.terminate() + + @pytest.mark.asyncio async def test_prepare_chat_payload_materializes_context_localhost_file_uri_image_urls( tmp_path,