From bb8e5321e872b8edab5e9bc2a308cd2bd32cff4d Mon Sep 17 00:00:00 2001
From: Doekin <105162544+Doekin@users.noreply.github.com>
Date: Sat, 14 Jun 2025 22:39:32 +0800
Subject: [PATCH] fix(ImageGenerationMiddleware): correctly process image URLs
 (#7198)

---
 .../middleware/feat/ImageGenerationMiddleware.ts   | 14 ++++++++++++--
 src/renderer/src/types/chunk.ts                    |  2 +-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts b/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts
index 560a9e0aac..324382b918 100644
--- a/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts
+++ b/src/renderer/src/aiCore/middleware/feat/ImageGenerationMiddleware.ts
@@ -97,11 +97,21 @@ export const ImageGenerationMiddleware: CompletionsMiddleware =
             )
           }
 
-          const b64_json_array = response.data?.map((item) => `data:image/png;base64,${item.b64_json}`) || []
+          let imageType: 'url' | 'base64' = 'base64'
+          const imageList =
+            response.data?.reduce((acc: string[], image) => {
+              if (image.url) {
+                acc.push(image.url)
+                imageType = 'url'
+              } else if (image.b64_json) {
+                acc.push(`data:image/png;base64,${image.b64_json}`)
+              }
+              return acc
+            }, []) || []
 
           enqueue({
             type: ChunkType.IMAGE_COMPLETE,
-            image: { type: 'base64', images: b64_json_array }
+            image: { type: imageType, images: imageList }
           })
 
           const usage = (response as any).usage || { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }
diff --git a/src/renderer/src/types/chunk.ts b/src/renderer/src/types/chunk.ts
index 4cb4755382..746c8999cb 100644
--- a/src/renderer/src/types/chunk.ts
+++ b/src/renderer/src/types/chunk.ts
@@ -137,7 +137,7 @@ export interface ImageCompleteChunk {
   /**
    * The image content of the chunk
    */
-  image?: { type: 'base64'; images: string[] }
+  image?: { type: 'url' | 'base64'; images: string[] }
 }
 
 export interface ThinkingDeltaChunk {