From 594bf768defd88a4de58c38d3dd9b94e12e74cef Mon Sep 17 00:00:00 2001 From: unknowbug Date: Sun, 23 Aug 2026 04:14:03 +0800 Subject: [PATCH 1/2] fix: normalize unsupported images before sending to vision providers --- .../agent/runners/tool_loop_agent_runner.py | 21 ++- .../core/provider/sources/openai_source.py | 31 ++++- astrbot/core/utils/media_utils.py | 86 ++++++++++++ tests/test_media_utils.py | 39 ++++++ tests/test_openai_source.py | 122 ++++++++++++++++-- 5 files changed, 279 insertions(+), 20 deletions(-) diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 8c91adbbfd..e43121ddf5 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -33,6 +33,10 @@ from astrbot.core.message.message_event_result import ( MessageChain, ) +from astrbot.core.utils.media_utils import ( + ResolvedMediaData, + normalize_image_for_provider, +) from astrbot.core.persona_error_reply import ( extract_persona_custom_error_message_from_event, ) @@ -1039,10 +1043,25 @@ async def step(self): text=f"[Image from tool '{cached_img.tool_name}', path='{cached_img.file_path}']" ) ) + normalized = normalize_image_for_provider( + ResolvedMediaData( + base64_data=base64_data, + mime_type=mime_type, + format=None, + ) + ) + if normalized is None: + logger.warning( + "Skip cached image for provider review: unsupported " + "tool image mime_type=%s path=%s", + mime_type, + cached_img.file_path, + ) + continue image_parts.append( ImageURLPart( image_url=ImageURLPart.ImageURL( - url=f"data:{mime_type};base64,{base64_data}", + url=normalized.to_data_url(), id=cached_img.file_path, ) ) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index f7870b7137..024a6cfa34 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 ( @@ -174,6 +175,8 @@ def _is_invalid_attachment_error(self, error: Exception) -> bool: return True if "download attachment" in error_text and "404" in error_text: return True + if "unsupported image" in error_text: + return True return False async def _image_ref_to_data_url( @@ -195,11 +198,18 @@ async def _resolve_image_part( *, image_detail: str | None = None, ) -> dict | None: - image_data = await self._image_ref_to_data_url(image_url, mode="safe") - if not image_data: - logger.warning("图片预处理结果为空,将忽略。") + image_data = await resolve_media_ref_to_base64_data( + image_url, + media_type="image", + strict=False, + ) + normalized = normalize_image_for_provider(image_data) + if normalized is None: + logger.warning( + "Image preprocessing returned no usable image; skipping image_url part." + ) return None - image_payload = {"url": image_data} + image_payload = {"url": normalized.to_data_url()} if image_detail: image_payload["detail"] = image_detail @@ -281,13 +291,20 @@ async def _transform_content_part(self, part: dict) -> dict: ) except Exception as exc: logger.warning( - "图片 %s 预处理失败,将保留原始内容。错误: %s", + "Image %s preprocessing failed; replacing with text placeholder: %s", url, exc, ) - return part + return {"type": "text", "text": "[image omitted]"} - return resolved_part or part + if resolved_part is None: + logger.warning( + "Image %s cannot be converted to a supported format; " + "replacing with text placeholder.", + url, + ) + return {"type": "text", "text": "[image omitted]"} + return resolved_part if part.get("type") == "audio_url": audio_ref = self._extract_audio_part_info(part) diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index 6cd58b92ec..43395637e3 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -131,6 +131,92 @@ def to_data_url(self) -> str: return f"data:{self.mime_type};base64,{self.base64_data}" +IMAGE_PROVIDER_SUPPORTED_MIME_TYPES = frozenset( + { + "image/jpeg", + "image/jpg", + "image/png", + "image/webp", + "image/gif", + } +) + + +def normalize_image_for_provider( + image_data: ResolvedMediaData | None, + supported_mimes: set[str] | frozenset[str] | None = None, +) -> ResolvedMediaData | None: + """Normalize image data to a MIME type accepted by vision providers. + + Args: + image_data: Resolved image data. + supported_mimes: MIME types accepted by the provider. Defaults to common + vision provider formats (webp/png/jpeg/gif). + + Returns: + Normalized image data, or None if the image cannot be converted. + """ + if image_data is None: + return None + + supported = supported_mimes or IMAGE_PROVIDER_SUPPORTED_MIME_TYPES + mime = (getattr(image_data, "mime_type", "") or "").lower() + if mime in supported: + return image_data + + try: + raw = image_data.to_bytes() + with PILImage.open(io.BytesIO(raw)) as img: + actual_fmt = str(img.format or "").upper() + # Re-label already supported formats so providers that validate the + # MIME header do not reject a mislabeled payload. + if actual_fmt == "JPEG": + return ResolvedMediaData( + base64_data=image_data.base64_data, + mime_type="image/jpeg", + format=image_data.format, + ) + if actual_fmt == "PNG": + return ResolvedMediaData( + base64_data=image_data.base64_data, + mime_type="image/png", + format=image_data.format, + ) + if actual_fmt == "WEBP": + return ResolvedMediaData( + base64_data=image_data.base64_data, + mime_type="image/webp", + format=image_data.format, + ) + if actual_fmt == "GIF": + return ResolvedMediaData( + base64_data=image_data.base64_data, + mime_type="image/gif", + format=image_data.format, + ) + + # Convert genuinely unsupported formats to a lossless or lossy + # provider-safe representation. Preserve transparency with PNG. + output = io.BytesIO() + has_alpha = img.mode in ("RGBA", "LA", "P") or "transparency" in img.info + if has_alpha: + img = img.convert("RGBA") + img.save(output, format="PNG") + converted_mime = "image/png" + else: + img = img.convert("RGB") + img.save(output, format="JPEG", quality=95) + converted_mime = "image/jpeg" + + return ResolvedMediaData( + base64_data=base64.b64encode(output.getvalue()).decode("utf-8"), + mime_type=converted_mime, + format=None, + ) + except Exception: + return None + + @dataclass(slots=True) class _LocalMediaFile: path: Path diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index efe5e65f02..41656316c8 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -871,3 +871,42 @@ async def test_wav_to_tencent_silk_skips_resample_for_supported_rate( assert len(fake.calls) == 1 assert fake.calls[0]["sample_rate"] == 24000 + + +def test_normalize_image_for_provider_converts_bmp_to_jpeg(): + from PIL import Image as PILImage + + buffer = BytesIO() + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(buffer, format="BMP") + data = media_utils.ResolvedMediaData( + base64_data=base64.b64encode(buffer.getvalue()).decode("ascii"), + mime_type="image/bmp", + format=None, + ) + + normalized = media_utils.normalize_image_for_provider(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_returns_none_for_unparseable_image(): + data = media_utils.ResolvedMediaData( + base64_data=base64.b64encode(b"not-an-image").decode("ascii"), + mime_type="image/svg+xml", + format=None, + ) + + assert media_utils.normalize_image_for_provider(data) is None + + +def test_normalize_image_for_provider_preserves_supported_mime(): + data = media_utils.ResolvedMediaData( + base64_data="abcd", + mime_type="image/png", + format=None, + ) + + assert media_utils.normalize_image_for_provider(data) is data diff --git a/tests/test_openai_source.py b/tests/test_openai_source.py index 911b76131f..c2c96a273b 100644 --- a/tests/test_openai_source.py +++ b/tests/test_openai_source.py @@ -1074,13 +1074,11 @@ async def test_encode_image_bs64_supports_file_uri(tmp_path): @pytest.mark.asyncio -async def test_resolve_image_part_supports_base64_scheme(): +async def test_resolve_image_part_rejects_non_image_base64_payload(): provider = _make_provider() try: - assert await provider._resolve_image_part("base64://abcd") == { - "type": "image_url", - "image_url": {"url": "data:image/jpeg;base64,abcd"}, - } + # "abcd" decodes to bytes that cannot be identified as an image. + assert await provider._resolve_image_part("base64://abcd") is None finally: await provider.terminate() @@ -1225,7 +1223,7 @@ def fake_warning(message, *args, **kwargs): @pytest.mark.asyncio -async def test_prepare_chat_payload_keeps_original_context_image_when_materialization_fails( +async def test_prepare_chat_payload_replaces_unresolvable_context_image_with_text( monkeypatch, ): provider = _make_provider() @@ -1268,12 +1266,7 @@ async def fake_resolve_media_ref_to_base64_data( assert payloads["messages"][0]["content"] == [ {"type": "text", "text": "look"}, - { - "type": "image_url", - "image_url": { - "url": "https://example.com/expired.png", - }, - }, + {"type": "text", "text": "[image omitted]"}, ] finally: await provider.terminate() @@ -2201,3 +2194,108 @@ async def fake_create(**kwargs): assert messages[1] == {"role": "user", "content": "again"} finally: await provider.terminate() + + +@pytest.mark.asyncio +async def test_handle_api_error_unsupported_image_removes_images_and_retries_text_only(): + provider = _make_provider() + try: + payloads = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,abcd"}, + }, + ], + } + ] + } + context_query = payloads["messages"] + err = _ErrorWithBody( + "upstream error", + { + "error": { + "code": "INVALID_REQUEST_ERROR", + "message": ( + ".messages[67].image[0]: You have uploaded an " + "unsupported image. Please make sure your image is valid " + "and has one of the following formats: webp, png, jpeg, " + "and gif." + ), + } + }, + ) + + success, *_rest = await provider._handle_api_error( + err, + payloads=payloads, + context_query=context_query, + func_tool=None, + chosen_key="test-key", + available_api_keys=["test-key"], + retry_cnt=0, + max_retries=10, + ) + + assert success is False + assert payloads["messages"][0]["content"] == [{"type": "text", "text": "hello"}] + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_resolve_image_part_converts_bmp_to_jpeg(): + provider = _make_provider() + try: + image_buffer = BytesIO() + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(image_buffer, format="BMP") + 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 + assert image_part["type"] == "image_url" + url = image_part["image_url"]["url"] + assert url.startswith("data:image/jpeg;base64,") + raw = base64.b64decode(url.split(",", 1)[1]) + assert PILImage.open(BytesIO(raw)).format == "JPEG" + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_resolve_image_part_returns_none_for_unsupported_svg(): + provider = _make_provider() + try: + svg = b"" + svg_base64 = base64.b64encode(svg).decode("ascii") + + assert await provider._resolve_image_part(f"base64://{svg_base64}") is None + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_transform_content_part_replaces_unresolvable_image_with_text(monkeypatch): + provider = _make_provider() + try: + + async def fake_resolve(image_url: str, *, image_detail: str | None = None): + return None + + monkeypatch.setattr(provider, "_resolve_image_part", fake_resolve) + + part = await provider._transform_content_part( + { + "type": "image_url", + "image_url": {"url": "data:image/bmp;base64,abcd"}, + } + ) + + assert part == {"type": "text", "text": "[image omitted]"} + finally: + await provider.terminate() From e373c90f0110bc22b5eade10377317bc08442942 Mon Sep 17 00:00:00 2001 From: unknowbug Date: Sun, 23 Aug 2026 04:18:49 +0800 Subject: [PATCH 2/2] fix: validate image bytes even when MIME is provider-safe --- astrbot/core/utils/media_utils.py | 62 +++++++++++++++---------------- tests/test_media_utils.py | 24 ++++++++++-- 2 files changed, 50 insertions(+), 36 deletions(-) diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index 43395637e3..3d533ebc6a 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -160,54 +160,50 @@ def normalize_image_for_provider( return None supported = supported_mimes or IMAGE_PROVIDER_SUPPORTED_MIME_TYPES - mime = (getattr(image_data, "mime_type", "") or "").lower() - if mime in supported: - return image_data try: raw = image_data.to_bytes() with PILImage.open(io.BytesIO(raw)) as img: actual_fmt = str(img.format or "").upper() - # Re-label already supported formats so providers that validate the - # MIME header do not reject a mislabeled payload. - if actual_fmt == "JPEG": + actual_mime_by_fmt = { + "JPEG": "image/jpeg", + "PNG": "image/png", + "WEBP": "image/webp", + "GIF": "image/gif", + } + actual_mime = actual_mime_by_fmt.get(actual_fmt) + if actual_mime in supported: + # Still validate bytes; return the re-labeled payload so a + # mislabeled provider-safe MIME header is corrected. return ResolvedMediaData( base64_data=image_data.base64_data, - mime_type="image/jpeg", - format=image_data.format, - ) - if actual_fmt == "PNG": - return ResolvedMediaData( - base64_data=image_data.base64_data, - mime_type="image/png", - format=image_data.format, - ) - if actual_fmt == "WEBP": - return ResolvedMediaData( - base64_data=image_data.base64_data, - mime_type="image/webp", - format=image_data.format, - ) - if actual_fmt == "GIF": - return ResolvedMediaData( - base64_data=image_data.base64_data, - mime_type="image/gif", + mime_type=actual_mime, format=image_data.format, ) - # Convert genuinely unsupported formats to a lossless or lossy - # provider-safe representation. Preserve transparency with PNG. - output = io.BytesIO() + # Convert unsupported formats to a provider-safe representation. + # Preserve transparency with PNG; otherwise use JPEG when available. has_alpha = img.mode in ("RGBA", "LA", "P") or "transparency" in img.info - if has_alpha: + if has_alpha and "image/png" in supported: + converted_mime = "image/png" + output_format = "PNG" img = img.convert("RGBA") - img.save(output, format="PNG") + elif "image/jpeg" in supported: + converted_mime = "image/jpeg" + output_format = "JPEG" + img = img.convert("RGB") + elif "image/png" in supported: converted_mime = "image/png" - else: + output_format = "PNG" img = img.convert("RGB") - img.save(output, format="JPEG", quality=95) - converted_mime = "image/jpeg" + else: + return None + output = io.BytesIO() + save_kwargs = {} + if output_format == "JPEG": + save_kwargs["quality"] = 95 + img.save(output, format=output_format, **save_kwargs) return ResolvedMediaData( base64_data=base64.b64encode(output.getvalue()).decode("utf-8"), mime_type=converted_mime, diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index 41656316c8..2f6d3cbf7d 100644 --- a/tests/test_media_utils.py +++ b/tests/test_media_utils.py @@ -902,11 +902,29 @@ def test_normalize_image_for_provider_returns_none_for_unparseable_image(): assert media_utils.normalize_image_for_provider(data) is None -def test_normalize_image_for_provider_preserves_supported_mime(): +def test_normalize_image_for_provider_rejects_invalid_supported_mime(): data = media_utils.ResolvedMediaData( - base64_data="abcd", + base64_data=base64.b64encode(b"not-an-image").decode("ascii"), mime_type="image/png", format=None, ) - assert media_utils.normalize_image_for_provider(data) is data + assert media_utils.normalize_image_for_provider(data) is None + + +def test_normalize_image_for_provider_preserves_supported_mime(): + from PIL import Image as PILImage + + buffer = BytesIO() + PILImage.new("RGB", (2, 2), (1, 2, 3)).save(buffer, format="PNG") + data = media_utils.ResolvedMediaData( + base64_data=base64.b64encode(buffer.getvalue()).decode("ascii"), + mime_type="image/jpg", # Mislabeled alias should be corrected. + format=None, + ) + + normalized = media_utils.normalize_image_for_provider(data) + + assert normalized is not None + assert normalized.mime_type == "image/png" + assert normalized.base64_data == data.base64_data