From 30858937b50a984eef7311af73e9d495630ce046 Mon Sep 17 00:00:00 2001
From: yinpeng <2291314224@qq.com>
Date: Tue, 11 Feb 2025 05:13:36 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E5=9B=BE=E7=89=87?=
 =?UTF-8?q?=E7=94=9F=E6=88=90=E5=93=8D=E5=BA=94=E6=A0=BC=E5=BC=8F=E5=88=87?=
 =?UTF-8?q?=E6=8D=A2=E5=B9=B6=E4=BC=98=E5=8C=96Markdown=E6=B8=B2=E6=9F=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/schemas/openai_models.py         |  2 +-
 app/services/image_create_service.py | 23 ++++++++++++++++-------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/app/schemas/openai_models.py b/app/schemas/openai_models.py
index fb1594f..52b143e 100644
--- a/app/schemas/openai_models.py
+++ b/app/schemas/openai_models.py
@@ -27,4 +27,4 @@ class ImageGenerationRequest(BaseModel):
     size: Optional[str] = "1024x1024"
     quality: Optional[str] = ""
     style: Optional[str] = ""
-    response_format: Optional[str] = "b64_json"
+    response_format: Optional[str] = "url"
diff --git a/app/services/image_create_service.py b/app/services/image_create_service.py
index 5f9b086..4498df0 100644
--- a/app/services/image_create_service.py
+++ b/app/services/image_create_service.py
@@ -56,11 +56,17 @@ class ImageCreateService:
                     filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
                     upload_response = image_uploader.upload(image_data,filename)
                     
-                # base64_image = base64.b64encode(image_data).decode('utf-8')
-                images_data.append({
-                    "url": f"{upload_response.data.url}",
-                    "revised_prompt": request.prompt
-                })
+                if request.response_format == "b64_json":
+                    base64_image = base64.b64encode(image_data).decode('utf-8')
+                    images_data.append({
+                        "b64_json": base64_image,
+                        "revised_prompt": request.prompt
+                    })
+                else:
+                    images_data.append({
+                        "url": f"{upload_response.data.url}",
+                        "revised_prompt": request.prompt
+                    })
 
             response_data = {
                 "created": int(time.time()),  # Current timestamp
@@ -76,6 +82,9 @@ class ImageCreateService:
         if image_datas:
             markdown_images = []
             for index, image_data in enumerate(image_datas):
-                markdown_images.append(f"![Generated Image {index+1}]({image_data['url']})")
+                if 'url' in image_data:
+                    markdown_images.append(f"![Generated Image {index+1}]({image_data['url']})")
+                else:
+                    # 如果是base64格式，创建data URL
+                    markdown_images.append(f"![Generated Image {index+1}](data:image/png;base64,{image_data['b64_json']})")
             return "\n".join(markdown_images)
-            
\ No newline at end of file