From b51ca06c7cf4094cd6fb241e43e6988fcd1e912d Mon Sep 17 00:00:00 2001
From: ligen131 <i@ligen131.com>
Date: Mon, 19 Aug 2024 00:00:29 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=8F=91=E9=80=81=E6=AD=A3=E7=A1=AE?=
 =?UTF-8?q?=E7=9A=84=E5=9B=BE=E7=89=87=E6=A0=BC=E5=BC=8F=E8=80=8C=E4=B8=8D?=
 =?UTF-8?q?=E6=98=AF=E9=BB=98=E8=AE=A4=E7=9A=84=20`image/jpeg`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/provider/modelmgr/apis/anthropicmsgs.py |  5 +++--
 pkg/provider/modelmgr/apis/chatcmpl.py      |  6 ++----
 pkg/provider/modelmgr/apis/ollamachat.py    |  4 ++--
 pkg/utils/image.py                          | 13 +++++++++----
 4 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/pkg/provider/modelmgr/apis/anthropicmsgs.py b/pkg/provider/modelmgr/apis/anthropicmsgs.py
index a9a3f05..8c5360e 100644
--- a/pkg/provider/modelmgr/apis/anthropicmsgs.py
+++ b/pkg/provider/modelmgr/apis/anthropicmsgs.py
@@ -72,12 +72,13 @@ class AnthropicMessages(api.LLMAPIRequester):
 
                 for i, ce in enumerate(m.content):
                     if ce.type == "image_url":
+                        base64_image, image_format = await image.qq_image_url_to_base64(ce.image_url.url)
                         alter_image_ele = {
                             "type": "image",
                             "source": {
                                 "type": "base64",
-                                "media_type": "image/jpeg",
-                                "data": await image.qq_image_url_to_base64(ce.image_url.url)
+                                "media_type": f"image/{image_format}",
+                                "data": base64_image
                             }
                         }
                         msg_dict["content"][i] = alter_image_ele
diff --git a/pkg/provider/modelmgr/apis/chatcmpl.py b/pkg/provider/modelmgr/apis/chatcmpl.py
index 3f1cfb3..222553a 100644
--- a/pkg/provider/modelmgr/apis/chatcmpl.py
+++ b/pkg/provider/modelmgr/apis/chatcmpl.py
@@ -136,7 +136,5 @@ class OpenAIChatCompletions(api.LLMAPIRequester):
         self,
         original_url: str,
     ) -> str:
-        
-        base64_image = await image.qq_image_url_to_base64(original_url)
-
-        return f"data:image/jpeg;base64,{base64_image}"
+        base64_image, image_format = await image.qq_image_url_to_base64(original_url)
+        return f"data:image/{image_format};base64,{base64_image}"
diff --git a/pkg/provider/modelmgr/apis/ollamachat.py b/pkg/provider/modelmgr/apis/ollamachat.py
index 88edfe7..65295dc 100644
--- a/pkg/provider/modelmgr/apis/ollamachat.py
+++ b/pkg/provider/modelmgr/apis/ollamachat.py
@@ -101,5 +101,5 @@ class OllamaChatCompletions(api.LLMAPIRequester):
             self,
             original_url: str,
     ) -> str:
-        base64_image: str = await image.qq_image_url_to_base64(original_url)
-        return f"data:image/jpeg;base64,{base64_image}"
+        base64_image, image_format = await image.qq_image_url_to_base64(original_url)
+        return f"data:image/{image_format};base64,{base64_image}"
diff --git a/pkg/utils/image.py b/pkg/utils/image.py
index 34acc2f..539337d 100644
--- a/pkg/utils/image.py
+++ b/pkg/utils/image.py
@@ -8,14 +8,14 @@ import aiohttp
 
 async def qq_image_url_to_base64(
     image_url: str
-) -> str:
-    """将QQ图片URL转为base64
+) -> typing.Tuple[str, str]:
+    """将QQ图片URL转为base64，并返回图片格式
 
     Args:
         image_url (str): QQ图片URL
 
     Returns:
-        str: base64编码
+        typing.Tuple[str, str]: base64编码和图片格式
     """
     parsed = urlparse(image_url)
     query = parse_qs(parsed.query)
@@ -35,7 +35,12 @@ async def qq_image_url_to_base64(
         ) as resp:
             resp.raise_for_status()  # 检查HTTP错误
             file_bytes = await resp.read()
+            content_type = resp.headers.get('Content-Type')
+            if not content_type or not content_type.startswith('image/'):
+                image_format = 'jpeg'
+            else: 
+                image_format = content_type.split('/')[-1]
 
     base64_str = base64.b64encode(file_bytes).decode()
 
-    return base64_str
+    return base64_str, image_format