fix _process_image_response

This commit is contained in:
duyalei 2025-04-23 17:52:35 +08:00
parent f5996b2f6b
commit f4ae14a490

View file

@ -23,24 +23,21 @@ def _process_image_response(response: Response, url: str) -> str:
image_bytes = response.content
base64_image = base64.b64encode(image_bytes).decode("utf-8")
image_type = response.headers.get("Content-Type")
if image_type is None:
img_type = url.split(".")[-1].lower()
_img_type = {
"jpg": "image/jpeg",
"jpeg": "image/jpeg",
"png": "image/png",
"gif": "image/gif",
"webp": "image/webp",
}.get(img_type)
if _img_type is None:
img_type_map = {
"jpg": "image/jpeg",
"jpeg": "image/jpeg",
"png": "image/png",
"gif": "image/gif",
"webp": "image/webp",
}
img_type = response.headers.get("Content-Type")
if img_type not in img_type_map:
_img_type = url.split('?')[0].split(".")[-1].lower()
img_type = img_type_map.get(img_type)
if img_type is None:
raise Exception(
f"Error: Unsupported image format. Format={_img_type}. Supported types = ['image/jpeg', 'image/png', 'image/gif', 'image/webp']"
)
img_type = _img_type
else:
img_type = image_type
result = f"data:{img_type};base64,{base64_image}"
in_memory_cache.set_cache(url, result)