From 72b8faf760a353f989cec4ef5683fa9b4699cb1c Mon Sep 17 00:00:00 2001
From: Chen Tao <70054568+eeee0717@users.noreply.github.com>
Date: Thu, 1 May 2025 21:03:10 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E8=B0=83=E6=95=B4=E5=9B=BE=E7=89=87?=
 =?UTF-8?q?=E9=AB=98=E5=BA=A6=EF=BC=8C=E4=BF=AE=E5=A4=8Dgemini=E7=94=BB?=
 =?UTF-8?q?=E5=9B=BE=20(#5585)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/pages/home/Messages/MessageImage.tsx  | 14 ++++++++-----
 .../providers/AiProvider/GeminiProvider.ts    | 20 ++++++++++++++-----
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/renderer/src/pages/home/Messages/MessageImage.tsx b/src/renderer/src/pages/home/Messages/MessageImage.tsx
index a8cea1d40b..646f7f3cf5 100644
--- a/src/renderer/src/pages/home/Messages/MessageImage.tsx
+++ b/src/renderer/src/pages/home/Messages/MessageImage.tsx
@@ -99,6 +99,7 @@ const MessageImage: FC<Props> = ({ block }) => {
         <Image
           src={image}
           key={`image-${index}`}
+          height={200}
           preview={{
             toolbarRender: (
               _,
@@ -130,18 +131,21 @@ const Container = styled.div`
   flex-direction: row;
   gap: 10px;
   margin-top: 8px;
-  width: 100%;
+  width: auto;
 `
 const Image = styled(AntdImage)`
-  border-radius: 10px;
-  width: 100%;
-  height: auto;
+  padding: 5px;
+  border-radius: 8px;
+
+  .ant-image-img {
+    width: auto;
+  }
 `
 const ToobarWrapper = styled(Space)`
   padding: 0px 24px;
   color: #fff;
   font-size: 20px;
-  background-color: rgba(0, 0, 0, 0.1);
+  background-color: rgba(238, 233, 233, 0.1);
   border-radius: 100px;
   .anticon {
     padding: 12px;
diff --git a/src/renderer/src/providers/AiProvider/GeminiProvider.ts b/src/renderer/src/providers/AiProvider/GeminiProvider.ts
index 7897bd762d..ef03b27dcc 100644
--- a/src/renderer/src/providers/AiProvider/GeminiProvider.ts
+++ b/src/renderer/src/providers/AiProvider/GeminiProvider.ts
@@ -40,7 +40,7 @@ import {
   Usage,
   WebSearchSource
 } from '@renderer/types'
-import { BlockCompleteChunk, ChunkType, LLMWebSearchCompleteChunk } from '@renderer/types/chunk'
+import { BlockCompleteChunk, Chunk, ChunkType, LLMWebSearchCompleteChunk } from '@renderer/types/chunk'
 import type { Message, Response } from '@renderer/types/newMessage'
 import { removeSpecialCharactersForTopicName } from '@renderer/utils'
 import { mcpToolCallResponseToGeminiMessage, parseAndCallTools } from '@renderer/utils/mcp-tools'
@@ -122,8 +122,11 @@ export default class GeminiProvider extends BaseProvider {
     // Add any generated images from previous responses
     const imageBlocks = findImageBlocks(message)
     for (const imageBlock of imageBlocks) {
-      if (imageBlock.metadata?.generateImage?.images && imageBlock.metadata.generateImage.images.length > 0) {
-        for (const imageUrl of imageBlock.metadata.generateImage.images) {
+      if (
+        imageBlock.metadata?.generateImageResponse?.images &&
+        imageBlock.metadata.generateImageResponse.images.length > 0
+      ) {
+        for (const imageUrl of imageBlock.metadata.generateImageResponse.images) {
           if (imageUrl && imageUrl.startsWith('data:')) {
             // Extract base64 data and mime type from the data URL
             const matches = imageUrl.match(/^data:(.+);base64,(.*)$/)
@@ -448,7 +451,7 @@ export default class GeminiProvider extends BaseProvider {
         }
 
         // 4. Image Generation
-        const generateImage = this.processGeminiImageResponse(chunk)
+        const generateImage = this.processGeminiImageResponse(chunk, onChunk)
         if (generateImage?.images?.length) {
           onChunk({ type: ChunkType.IMAGE_COMPLETE, image: generateImage })
         }
@@ -715,7 +718,10 @@ export default class GeminiProvider extends BaseProvider {
    * @param response - Gemini响应
    * @param onChunk - 处理生成块的回调
    */
-  private processGeminiImageResponse(chunk: GenerateContentResponse): { type: 'base64'; images: string[] } | undefined {
+  private processGeminiImageResponse(
+    chunk: GenerateContentResponse,
+    onChunk: (chunk: Chunk) => void
+  ): { type: 'base64'; images: string[] } | undefined {
     const parts = chunk.candidates?.[0]?.content?.parts
     if (!parts) {
       return
@@ -727,6 +733,10 @@ export default class GeminiProvider extends BaseProvider {
         if (!part.inlineData) {
           return null
         }
+        // onChunk的位置需要更改
+        onChunk({
+          type: ChunkType.IMAGE_CREATED
+        })
         const dataPrefix = `data:${part.inlineData.mimeType || 'image/png'};base64,`
         return part.inlineData.data?.startsWith('data:') ? part.inlineData.data : dataPrefix + part.inlineData.data
       })