Fix: Add image generation display to markdown message

Author: ShenKSPZ

src/components/chat/ImageGenerationButton.tsx

@@ -141,7 +141,6 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
     });
   };
 
-  const isButtonEnabled = !disabled && providers.length > 0 && isEnabled;
   const buttonClass = `flex items-center justify-center w-8 h-8 rounded-full focus:outline-none ${
     isEnabled ? 'image-generation-button' : 'text-gray-400 bg-gray-100'
   }`;
@@ -152,10 +151,10 @@ const ImageGenerationButton: React.FC<ImageGenerationButtonProps> = ({
         ref={buttonRef}
         type="button"
         onClick={togglePopup}
-        disabled={!providers.length > 0 || disabled}
+        disabled={providers.length === 0 || disabled}
         className={buttonClass}
         title={
-          !providers.length > 0 
+          providers.length === 0 
             ? t('chat.imageGenerationNotAvailable') 
             : isEnabled 
               ? t('chat.generateImage') 

src/components/chat/MarkdownContent.tsx

@@ -11,6 +11,8 @@ import type { HTMLProps } from 'react';
 import { MessageContent, MessageContentType } from '../../types/chat';
 import { MessageHelper } from '../../services/message-helper';
 import FileAttachmentDisplay from './FileAttachmentDisplay';
+import { Loader2 } from 'lucide-react';
+import { useTranslation } from '../../hooks/useTranslation';
 
 interface MarkdownContentProps {
   content: MessageContent[];
@@ -23,28 +25,38 @@ type CodeProps = React.ClassAttributes<HTMLElement> &
   };
 
 export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUserMessage = false }) => {
+  const { t } = useTranslation();
   const [processedContent, setProcessedContent] = useState('');
   const [thinkContent, setThinkContent] = useState<string | null>(null);
   const [isThinkExpanded, setIsThinkExpanded] = useState(true);
   const [fileContents, setFileContents] = useState<MessageContent[]>([]);
+  const [imageContents, setImageContents] = useState<MessageContent[]>([]);
+  const [isProcessingImage, setIsProcessingImage] = useState(false);
+  const [imageGenerationError, setImageGenerationError] = useState<string | null>(null);
 
-  // Process content and check for thinking blocks and files
+  // Process content and check for thinking blocks, files, and image generation
   useEffect(() => {
-    // Extract text and file contents
+    // Extract text, file, and image contents
     const textContents: MessageContent[] = [];
     const files: MessageContent[] = [];
+    const images: MessageContent[] = [];
 
     content.forEach(item => {
       if (item.type === MessageContentType.Text) {
         textContents.push(item);
       } else if (item.type === MessageContentType.File) {
         files.push(item);
+      } else if (item.type === MessageContentType.Image) {
+        images.push(item);
       }
     });
 
     // Save file contents for rendering
     setFileContents(files);
 
+    // Save image contents for rendering
+    setImageContents(images);
+    
     // Create a function for a safer replacement
     function safeReplace(str: string, search: string, replace: string): string {
       // Split the string by the search term
@@ -55,6 +67,32 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
 
     let processed = MessageHelper.MessageContentToText(textContents);
 
+    // Detect image generation in progress
+    const imageGenInProgressMatch = processed.match(
+      /(?:generating|creating|processing)\s+(?:an\s+)?image(?:s)?\s+(?:with|using|for|from)?(?:\s+prompt)?(?::|;)?\s*["']?([^"']+)["']?/i
+    ) || processed.match(/\bimage\s+generation\s+in\s+progress\b/i);
+    
+    if ((imageGenInProgressMatch && images.length === 0) || 
+        (processed.includes('generate_image') && processed.includes('tool call') && images.length === 0)) {
+      setIsProcessingImage(true);
+      setImageGenerationError(null);
+    } else {
+      setIsProcessingImage(false);
+    }
+    
+    // Detect image generation errors
+    const imageGenErrorMatch = processed.match(
+      /(?:error|failed|couldn't|unable)\s+(?:in\s+)?(?:generating|creating|processing)\s+(?:an\s+)?image(?:s)?(?::|;)?\s*["']?([^"']+)["']?/i
+    ) || processed.match(/\bimage\s+generation\s+(?:error|failed)\b:?\s*["']?([^"']+)["']?/i);
+    
+    if (imageGenErrorMatch || (processed.includes('error') && processed.includes('generate_image'))) {
+      const errorMessage = imageGenErrorMatch ? (imageGenErrorMatch[1] || "Unknown error occurred") : "Failed to generate image";
+      setImageGenerationError(errorMessage);
+      setIsProcessingImage(false);
+    } else if (images.length > 0) {
+      setImageGenerationError(null);
+    }
+    
     // Check if content contains thinking block
     const thinkMatch = processed.match(/<think>([\s\S]*?)<\/think>([\s\S]*)/);
 
@@ -134,6 +172,46 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
         </div>
       )}
 
+      {/* Image Generation In Progress */}
+      {isProcessingImage && (
+        <div className="p-4 mb-3 border border-gray-200 rounded-md">
+          <div className="flex items-center gap-2 mb-2">
+            <Loader2 size={20} className="text-blue-500 animate-spin" />
+            <span className="font-medium">{t('imageGeneration.generating')}</span>
+          </div>
+          <div className="flex items-center justify-center w-full h-40 bg-gray-100 rounded-md">
+            <span className="text-sm text-gray-400">{t('imageGeneration.creatingImage')}</span>
+          </div>
+        </div>
+      )}
+      
+      {/* Image Generation Error */}
+      {imageGenerationError && (
+        <div className="p-4 mb-3 border border-red-200 rounded-md bg-red-50">
+          <div className="flex items-center gap-2 mb-2">
+            <span className="font-medium text-red-600">{t('imageGeneration.generationFailed')}</span>
+          </div>
+          <div className="text-sm text-red-600">
+            {imageGenerationError}
+          </div>
+        </div>
+      )}
+      
+      {/* Generated Images */}
+      {imageContents.length > 0 && (
+        <div className="grid grid-cols-1 gap-4 mb-3 md:grid-cols-2">
+          {imageContents.map((image, index) => (
+            <div key={index} className="overflow-hidden border border-gray-200 rounded-md">
+              <img 
+                src={`data:image/png;base64,${image.content}`} 
+                alt={t('imageGeneration.generatedImage')} 
+                className="w-full h-auto"
+              />
+            </div>
+          ))}
+        </div>
+      )}
+      
       {thinkContent && (
         <div className="mb-4">
           <div 

src/locales/en/translation.json

@@ -103,7 +103,10 @@
     "generationCount": "Generation Count",
     "randomSeed": "Random Seed",
     "generateButton": "Generate",
-    "generating": "Generating...",
+    "generating": "Generating image...",
+    "creatingImage": "AI is creating your image",
+    "generationFailed": "Image generation failed",
+    "generatedImage": "Generated image",
     "prompt": "Prompt",
     "promptPlaceholder": "Describe the image you want to create, e.g.: a peaceful lake at sunset with mountains in the background",
     "results": "Generated Results",

src/locales/es/translation.json

@@ -96,16 +96,19 @@
     "selectModel_search_placeholder": "Buscar modelos..."
   },
   "imageGeneration": {
-    "title": "Imagen Generación",
+    "title": "Generación de Imágenes",
     "provider": "Proveedor",
     "model": "Modelo",
     "imageSize": "Tamaño de Imagen",
     "generationCount": "Cantidad de Generaciones",
     "randomSeed": "Semilla Aleatoria",
     "generateButton": "Generar",
-    "generating": "Generando...",
+    "generating": "Generando imagen...",
+    "creatingImage": "La IA está creando tu imagen",
+    "generationFailed": "Falló la generación de imagen",
+    "generatedImage": "Imagen generada",
     "prompt": "Prompt",
-    "promptPlaceholder": "Describe la imagen que quieres crear, ej.: un lago tranquilo al atardecer con montañas en el fondo",
+    "promptPlaceholder": "Describe la imagen que quieres crear, p.ej.: un lago tranquilo al atardecer con montañas en el fondo",
     "results": "Resultados Generados",
     "placeholderText": "Ingresa un prompt y haz clic en generar para crear imágenes",
     "apiKeyMissing": "Por favor, configura tu clave API para el proveedor seleccionado en la configuración.",

src/locales/ja/translation.json

@@ -103,13 +103,16 @@
     "generationCount": "生成数",
     "randomSeed": "ランダムシード",
     "generateButton": "生成",
-    "generating": "生成中...",
+    "generating": "画像を生成中...",
+    "creatingImage": "AIが画像を作成しています",
+    "generationFailed": "画像生成に失敗しました",
+    "generatedImage": "生成された画像",
     "prompt": "プロンプト",
-    "promptPlaceholder": "作成したい画像を説明してください。例：夕日が沈む静かな湖、背景には山々",
+    "promptPlaceholder": "作成したい画像を説明してください。例：夕暮れの穏やかな湖と背景の山々",
     "results": "生成結果",
-    "placeholderText": "プロンプトを入力して生成ボタンをクリックして画像を作成",
-    "apiKeyMissing": "選択したプロバイダーのAPIキーを設定で設定してください。",
-    "seedHelp": "再現可能な結果のためのシード値"
+    "placeholderText": "プロンプトを入力して生成ボタンをクリックすると画像が作成されます",
+    "apiKeyMissing": "設定で選択したプロバイダーのAPIキーを設定してください。",
+    "seedHelp": "再現可能な結果のためのシード"
   },
   "mcpServer": {
     "title": "MCPサーバー",

src/locales/ko/translation.json

@@ -100,15 +100,18 @@
     "provider": "제공자",
     "model": "모델",
     "imageSize": "이미지 크기",
-    "generationCount": "생성 수량",
+    "generationCount": "생성 수",
     "randomSeed": "랜덤 시드",
     "generateButton": "생성",
-    "generating": "생성 중...",
+    "generating": "이미지 생성 중...",
+    "creatingImage": "AI가 이미지를 만들고 있습니다",
+    "generationFailed": "이미지 생성 실패",
+    "generatedImage": "생성된 이미지",
     "prompt": "프롬프트",
-    "promptPlaceholder": "만들고 싶은 이미지를 설명하세요. 예: 석양이 지는 고요한 호수, 배경에는 산맥",
-    "results": "생성 결과",
-    "placeholderText": "프롬프트를 입력하고 생성 버튼을 클릭하여 이미지 생성",
-    "apiKeyMissing": "선택한 제공자의 API 키를 설정에서 설정하세요.",
+    "promptPlaceholder": "원하는 이미지를 설명하세요. 예: 일몰 시 평화로운 호수와 배경에 산이 있는 풍경",
+    "results": "생성된 결과",
+    "placeholderText": "프롬프트를 입력하고 생성 버튼을 클릭하여 이미지를 만드세요",
+    "apiKeyMissing": "설정에서 선택한 제공자의 API 키를 설정해 주세요.",
     "seedHelp": "재현 가능한 결과를 위한 시드"
   },
   "mcpServer": {

src/locales/zh-CN/translation.json

@@ -96,19 +96,22 @@
     "selectModel_search_placeholder": "搜索模型..."
   },
   "imageGeneration": {
-    "title": "图片生成",
+    "title": "图像生成",
     "provider": "提供商",
     "model": "模型",
-    "imageSize": "图片尺寸",
+    "imageSize": "图像尺寸",
     "generationCount": "生成数量",
     "randomSeed": "随机种子",
     "generateButton": "生成",
-    "generating": "生成中...",
+    "generating": "正在生成图像...",
+    "creatingImage": "AI正在创建您的图像",
+    "generationFailed": "图像生成失败",
+    "generatedImage": "生成的图像",
     "prompt": "提示词",
-    "promptPlaceholder": "描述你想创建的图片，例如：一个宁静的湖泊，夕阳西下，远处是群山",
+    "promptPlaceholder": "描述您想要创建的图像，例如：日落时分的平静湖泊，背景是山脉",
     "results": "生成结果",
-    "placeholderText": "输入提示词并点击生成按钮来创建图片",
-    "apiKeyMissing": "请在设置中为所选提供商设置您的 API 密钥。",
+    "placeholderText": "输入提示词并点击生成按钮创建图像",
+    "apiKeyMissing": "请在设置中为所选提供商设置API密钥。",
     "seedHelp": "用于可重现结果的种子"
   },
   "mcpServer": {

src/locales/zh-TW/translation.json

@@ -96,19 +96,22 @@
     "selectModel_search_placeholder": "搜尋模型..."
   },
   "imageGeneration": {
-    "title": "圖片生成",
-    "provider": "提供商",
+    "title": "圖像生成",
+    "provider": "提供者",
     "model": "模型",
-    "imageSize": "圖片尺寸",
+    "imageSize": "圖像大小",
     "generationCount": "生成數量",
     "randomSeed": "隨機種子",
     "generateButton": "生成",
-    "generating": "生成中...",
+    "generating": "正在生成圖像...",
+    "creatingImage": "AI正在創建您的圖像",
+    "generationFailed": "圖像生成失敗",
+    "generatedImage": "生成的圖像",
     "prompt": "提示詞",
-    "promptPlaceholder": "描述你想創建的圖片，例如：一個寧靜的湖泊，夕陽西下，遠處是群山",
+    "promptPlaceholder": "描述您想要創建的圖像，例如：日落時分的平靜湖泊，背景是山脈",
     "results": "生成結果",
-    "placeholderText": "輸入提示詞並點擊生成按鈕來創建圖片",
-    "apiKeyMissing": "請在設定中為所選提供商設置您的 API 金鑰。",
+    "placeholderText": "輸入提示詞並點擊生成按鈕創建圖像",
+    "apiKeyMissing": "請在設定中為所選提供者設置API金鑰。",
     "seedHelp": "用於可重現結果的種子"
   },
   "mcpServer": {

src/services/providers/common-provider-service.ts

@@ -1,13 +1,13 @@
-import { generateText, LanguageModelV1, LanguageModelUsage, Provider, streamText, ToolSet, ToolChoice } from 'ai';
-import { v4 as uuidv4 } from 'uuid';
+import { generateText, streamText, Provider, type LanguageModelUsage } from 'ai';
 import { Message, MessageRole } from '../../types/chat';
 import { AiServiceProvider, CompletionOptions } from '../core/ai-service-provider';
 import { SettingsService } from '../settings-service';
 import { StreamControlHandler } from '../streaming-control';
-import { AIServiceCapability } from '../../types/capabilities';
-import { mapModelCapabilities } from '../../types/capabilities';
-import { ModelSettings } from '../../types/settings';
+import { v4 as uuidv4 } from 'uuid';
 import { MessageHelper } from '../message-helper';
+import { AIServiceCapability, mapModelCapabilities } from '../../types/capabilities';
+import { ModelSettings } from '../../types/settings';
+import { LanguageModelV1, ToolChoice, ToolSet } from 'ai';
 /**
  * Implementation of OpenAI service provider using the AI SDK
  */
@@ -165,6 +165,18 @@ export class CommonProviderHelper implements AiServiceProvider {
           presencePenalty: options.presence_penalty,
           tools: tools,
           toolChoice: toolChoice,
+          onToolCall: (toolCall) => {
+            console.log('Tool call:', toolCall);
+            streamController.onToolCall(toolCall);
+          },
+          onToolCallResult: (toolCallId, result) => {
+            console.log('Tool call result:', toolCallId, result);
+            streamController.onToolCallResult(toolCallId, result);
+          },
+          onToolCallError: (toolCallId, error) => {
+            console.error('Tool call error:', toolCallId, error);
+            streamController.onToolCallError(toolCallId, error);
+          },
           onFinish: (result: { usage: LanguageModelUsage }) => {
             console.log('OpenAI streaming chat completion finished');
             streamController.onFinish(result.usage);
@@ -192,9 +204,26 @@ export class CommonProviderHelper implements AiServiceProvider {
           presencePenalty: options.presence_penalty,
           tools: tools,
           toolChoice: toolChoice,
+          maxSteps: 3, // Allow multiple steps for tool calls
         });
 
         console.log('toolResults: ', toolResults);
+        
+        // Process tool results for images
+        if (toolResults && toolResults.length > 0) {
+          for (const toolResult of toolResults) {
+            if (toolResult.name === 'generate_image' && toolResult.result?.images) {
+              const images = toolResult.result.images;
+              if (Array.isArray(images)) {
+                for (const imageData of images) {
+                  if (typeof imageData === 'string') {
+                    streamController.onToolCallResult(toolResult.id, { images: [imageData] });
+                  }
+                }
+              }
+            }
+          }
+        }
 
         fullText = text;
         streamController.onChunk(fullText);