Skip to content

Commit bd59785

Browse files
authored
fix(provider): 修复 base64:// 图片引用的 MIME 类型声明不准确问题 (#8177)
- 新增 `_detect_image_format` 方法,使用 Pillow verify() 检测图片真实格式,避免完整解码像素带来的额外开销 - 新增 `_base64_image_ref_to_data_url` 方法,将 base64:// 引用转换为携带真实 MIME 类型的 data URL,修复 PNG/GIF/WebP 等图片被错误声明为 image/jpeg 的问题 - 提取 `_IMAGE_FORMAT_MIME_TYPES` 类常量和 `_image_format_to_mime_type` 方法,统一本地文件与 base64:// 引用的格式映射逻辑,新增 TIFF/AVIF 格式支持 - 新增单元测试 `test_resolve_image_part_preserves_base64_png_mime_type`,覆盖 PNG 图片 MIME 类型正确声明的场景 Closes #8174
1 parent 95d8057 commit bd59785

2 files changed

Lines changed: 78 additions & 13 deletions

File tree

astrbot/core/provider/sources/openai_source.py

Lines changed: 55 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import asyncio
22
import base64
3+
import binascii
34
import copy
45
import inspect
56
import json
@@ -55,6 +56,17 @@
5556
)
5657
class ProviderOpenAIOfficial(Provider):
5758
_ERROR_TEXT_CANDIDATE_MAX_CHARS = 4096
59+
# 部分 OpenAI 兼容中转站会校验 data URL 的 MIME 类型是否和图片字节一致。
60+
# 这里统一维护格式映射,确保本地文件和 `base64://` 图片引用使用相同声明。
61+
_IMAGE_FORMAT_MIME_TYPES = {
62+
"JPEG": "image/jpeg",
63+
"PNG": "image/png",
64+
"GIF": "image/gif",
65+
"WEBP": "image/webp",
66+
"BMP": "image/bmp",
67+
"TIFF": "image/tiff",
68+
"AVIF": "image/avif",
69+
}
5870

5971
@classmethod
6072
def _truncate_error_text_candidate(cls, text: str) -> str:
@@ -195,25 +207,55 @@ def _encode_image_file_to_data_url(
195207
raise
196208
return None
197209

198-
try:
199-
with PILImage.open(BytesIO(image_bytes)) as image:
200-
image.verify()
201-
image_format = str(image.format or "").upper()
202-
except (OSError, UnidentifiedImageError):
210+
image_format = cls._detect_image_format(image_bytes)
211+
if image_format is None:
203212
if mode == "strict":
204213
raise ValueError(f"Invalid image file: {image_path}")
205214
return None
206215

207-
mime_type = {
208-
"JPEG": "image/jpeg",
209-
"PNG": "image/png",
210-
"GIF": "image/gif",
211-
"WEBP": "image/webp",
212-
"BMP": "image/bmp",
213-
}.get(image_format, "image/jpeg")
216+
mime_type = cls._image_format_to_mime_type(image_format)
214217
image_bs64 = base64.b64encode(image_bytes).decode("utf-8")
215218
return f"data:{mime_type};base64,{image_bs64}"
216219

220+
@classmethod
221+
def _detect_image_format(cls, image_bytes: bytes) -> str | None:
222+
"""返回 Pillow 校验后的图片格式,非法图片返回 None。"""
223+
try:
224+
# verify() 只校验图片容器,不完整解码像素。
225+
# 这里仅需要可信的格式标签,因此这种方式足够且开销较小。
226+
with PILImage.open(BytesIO(image_bytes)) as image:
227+
image.verify()
228+
return str(image.format or "").upper()
229+
except (OSError, UnidentifiedImageError):
230+
return None
231+
232+
@classmethod
233+
def _image_format_to_mime_type(cls, image_format: str | None) -> str:
234+
"""将 Pillow 图片格式映射为 data URL 使用的 MIME 类型。"""
235+
# 未识别格式保持历史 JPEG 兜底,兼容传入任意 `base64://` 内容的旧调用方。
236+
return cls._IMAGE_FORMAT_MIME_TYPES.get(
237+
str(image_format or "").upper(), "image/jpeg"
238+
)
239+
240+
@classmethod
241+
def _base64_image_ref_to_data_url(cls, image_ref: str) -> str:
242+
"""将 `base64://` 图片引用转换为带真实 MIME 的 data URL。"""
243+
raw_base64 = image_ref.removeprefix("base64://")
244+
mime_type = "image/jpeg"
245+
try:
246+
# 平台适配器可能通过 `base64://` 传入 PNG/GIF/WebP 等图片字节,
247+
# 但不会额外携带 MIME 元数据。发送 OpenAI 请求前先识别真实格式,
248+
# 避免把 PNG 等图片错误声明为 JPEG。
249+
image_bytes = base64.b64decode(raw_base64)
250+
except (binascii.Error, ValueError):
251+
# 对错误或非图片 base64 保持旧行为:继续返回 JPEG data URL,
252+
# 避免让历史调用方因为格式识别失败而直接抛异常。
253+
pass
254+
else:
255+
image_format = cls._detect_image_format(image_bytes)
256+
mime_type = cls._image_format_to_mime_type(image_format)
257+
return f"data:{mime_type};base64,{raw_base64}"
258+
217259
@staticmethod
218260
def _file_uri_to_path(file_uri: str) -> str:
219261
"""Normalize file URIs to paths.
@@ -242,7 +284,7 @@ async def _image_ref_to_data_url(
242284
mode: Literal["safe", "strict"] = "safe",
243285
) -> str | None:
244286
if image_ref.startswith("base64://"):
245-
return image_ref.replace("base64://", "data:image/jpeg;base64,")
287+
return self._base64_image_ref_to_data_url(image_ref)
246288

247289
if image_ref.startswith("http"):
248290
image_path = await download_image_by_url(image_ref)

tests/test_openai_source.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1+
import base64
12
import builtins
3+
from io import BytesIO
24
from types import SimpleNamespace
35

46
import pytest
@@ -1031,6 +1033,27 @@ async def test_resolve_image_part_supports_base64_scheme():
10311033
await provider.terminate()
10321034

10331035

1036+
@pytest.mark.asyncio
1037+
async def test_resolve_image_part_preserves_base64_png_mime_type():
1038+
provider = _make_provider()
1039+
try:
1040+
image_buffer = BytesIO()
1041+
PILImage.new("RGBA", (1, 1), (255, 0, 0, 255)).save(
1042+
image_buffer,
1043+
format="PNG",
1044+
)
1045+
image_base64 = base64.b64encode(image_buffer.getvalue()).decode("ascii")
1046+
1047+
image_part = await provider._resolve_image_part(f"base64://{image_base64}")
1048+
1049+
assert image_part == {
1050+
"type": "image_url",
1051+
"image_url": {"url": f"data:image/png;base64,{image_base64}"},
1052+
}
1053+
finally:
1054+
await provider.terminate()
1055+
1056+
10341057
@pytest.mark.asyncio
10351058
async def test_prepare_chat_payload_materializes_context_localhost_file_uri_image_urls(
10361059
tmp_path,

0 commit comments

Comments
 (0)