@lobehub/chat 1.116.4 → 1.117.1

package/packages/model-runtime/src/google/createImage.ts

@@ -0,0 +1,152 @@
+import { Content, GoogleGenAI, Part } from '@google/genai';
+import { imageUrlToBase64 } from '@lobechat/utils';
+
+import { CreateImagePayload, CreateImageResponse } from '../types/image';
+import { AgentRuntimeError } from '../utils/createError';
+import { parseGoogleErrorMessage } from '../utils/googleErrorParser';
+import { parseDataUri } from '../utils/uriParser';
+
+/**
+ * Extract image data from generateContent response
+ */
+function extractImageFromResponse(response: any): CreateImageResponse {
+  const candidate = response.candidates?.[0];
+  if (!candidate?.content?.parts) {
+    throw new Error('No image generated');
+  }
+
+  for (const part of candidate.content.parts) {
+    if (part.inlineData?.data) {
+      const imageUrl = `data:${part.inlineData.mimeType || 'image/png'};base64,${part.inlineData.data}`;
+      return { imageUrl };
+    }
+  }
+
+  throw new Error('No image data found in response');
+}
+
+/**
+ * Generate images using traditional Imagen models with generateImages API
+ */
+async function generateByImageModel(
+  client: GoogleGenAI,
+  payload: CreateImagePayload,
+): Promise<CreateImageResponse> {
+  const { model, params } = payload;
+
+  const response = await client.models.generateImages({
+    config: {
+      aspectRatio: params.aspectRatio,
+      numberOfImages: 1,
+    },
+    model,
+    prompt: params.prompt,
+  });
+
+  if (!response.generatedImages || response.generatedImages.length === 0) {
+    throw new Error('No images generated');
+  }
+
+  const generatedImage = response.generatedImages[0];
+  if (!generatedImage.image || !generatedImage.image.imageBytes) {
+    throw new Error('Invalid image data');
+  }
+
+  const { imageBytes } = generatedImage.image;
+  // 1. official doc use png as example
+  // 2. no responseType param support like openai now.
+  // I think we can just hard code png now
+  const imageUrl = `data:image/png;base64,${imageBytes}`;
+
+  return { imageUrl };
+}
+
+/**
+ * Generate images using Gemini Chat Models with generateContent
+ */
+async function generateImageByChatModel(
+  client: GoogleGenAI,
+  payload: CreateImagePayload,
+): Promise<CreateImageResponse> {
+  const { model, params } = payload;
+  const actualModel = model.replace(':image', '');
+
+  // Build content parts
+  const parts: Part[] = [{ text: params.prompt }];
+
+  // Add image for editing if provided
+  if (params.imageUrl && params.imageUrl !== null) {
+    const { mimeType, base64, type } = parseDataUri(params.imageUrl);
+
+    if (type === 'base64') {
+      if (!base64) {
+        throw new TypeError("Image URL doesn't contain base64 data");
+      }
+
+      parts.push({
+        inlineData: {
+          data: base64,
+          mimeType: mimeType || 'image/png',
+        },
+      });
+    } else if (type === 'url') {
+      const { base64: urlBase64, mimeType: urlMimeType } = await imageUrlToBase64(params.imageUrl);
+
+      parts.push({
+        inlineData: {
+          data: urlBase64,
+          mimeType: urlMimeType,
+        },
+      });
+    } else {
+      throw new TypeError(`currently we don't support image url: ${params.imageUrl}`);
+    }
+  }
+
+  const contents: Content[] = [
+    {
+      parts,
+      role: 'user',
+    },
+  ];
+
+  const response = await client.models.generateContent({
+    config: {
+      responseModalities: ['Image'],
+    },
+    contents,
+    model: actualModel,
+  });
+
+  return extractImageFromResponse(response);
+}
+
+/**
+ * Create image using Google AI models
+ */
+export async function createGoogleImage(
+  client: GoogleGenAI,
+  provider: string,
+  payload: CreateImagePayload,
+): Promise<CreateImageResponse> {
+  try {
+    const { model } = payload;
+
+    // Handle Gemini 2.5 Flash Image models that use generateContent
+    if (model.endsWith(':image')) {
+      return await generateImageByChatModel(client, payload);
+    }
+
+    // Handle traditional Imagen models that use generateImages
+    return await generateByImageModel(client, payload);
+  } catch (error) {
+    const err = error as Error;
+
+    const { errorType, error: parsedError } = parseGoogleErrorMessage(err.message);
+    throw AgentRuntimeError.createImage({
+      error: parsedError,
+      errorType,
+      provider,
+    });
+  }
+}

package/packages/model-runtime/src/google/index.test.ts

@@ -5,7 +5,6 @@ import OpenAI from 'openai';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { OpenAIChatMessage } from '@/libs/model-runtime';
-import { CreateImagePayload } from '@/libs/model-runtime/types/image';
 import { ChatStreamPayload } from '@/types/openai/chat';
 
 import * as debugStreamModule from '../utils/debugStream';
@@ -827,331 +826,4 @@ describe('LobeGoogleAI', () => {
       });
     });
   });
-
-  describe('createImage', () => {
-    it('should create image successfully with basic parameters', async () => {
-      // Arrange - Use real base64 image data (5x5 red pixel PNG)
-      const realBase64ImageData =
-        'iVBORw0KGgoAAAANSUhEUgAAAAUAAAAFCAYAAACNbyblAAAAHElEQVQI12P4//8/w38GIAXDIBKE0DHxgljNBAAO9TXL0Y4OHwAAAABJRU5ErkJggg==';
-      const mockImageResponse = {
-        generatedImages: [
-          {
-            image: {
-              imageBytes: realBase64ImageData,
-            },
-          },
-        ],
-      };
-      vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-        mockImageResponse as any,
-      );
-
-      const payload: CreateImagePayload = {
-        model: 'imagen-4.0-generate-preview-06-06',
-        params: {
-          prompt: 'A beautiful landscape with mountains and trees',
-          aspectRatio: '1:1',
-        },
-      };
-
-      // Act
-      const result = await instance.createImage(payload);
-
-      // Assert
-      expect(instance['client'].models.generateImages).toHaveBeenCalledWith({
-        model: 'imagen-4.0-generate-preview-06-06',
-        prompt: 'A beautiful landscape with mountains and trees',
-        config: {
-          aspectRatio: '1:1',
-          numberOfImages: 1,
-        },
-      });
-      expect(result).toEqual({
-        imageUrl: `data:image/png;base64,${realBase64ImageData}`,
-      });
-    });
-
-    it('should support different aspect ratios like 16:9 for widescreen images', async () => {
-      // Arrange - Use real base64 data
-      const realBase64Data =
-        'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
-      const mockImageResponse = {
-        generatedImages: [
-          {
-            image: {
-              imageBytes: realBase64Data,
-            },
-          },
-        ],
-      };
-      vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-        mockImageResponse as any,
-      );
-
-      const payload: CreateImagePayload = {
-        model: 'imagen-4.0-ultra-generate-preview-06-06',
-        params: {
-          prompt: 'Cinematic landscape shot with dramatic lighting',
-          aspectRatio: '16:9',
-        },
-      };
-
-      // Act
-      await instance.createImage(payload);
-
-      // Assert
-      expect(instance['client'].models.generateImages).toHaveBeenCalledWith({
-        model: 'imagen-4.0-ultra-generate-preview-06-06',
-        prompt: 'Cinematic landscape shot with dramatic lighting',
-        config: {
-          aspectRatio: '16:9',
-          numberOfImages: 1,
-        },
-      });
-    });
-
-    it('should work with only prompt when aspect ratio is not specified', async () => {
-      // Arrange
-      const realBase64Data =
-        'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
-      const mockImageResponse = {
-        generatedImages: [
-          {
-            image: {
-              imageBytes: realBase64Data,
-            },
-          },
-        ],
-      };
-      vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-        mockImageResponse as any,
-      );
-
-      const payload: CreateImagePayload = {
-        model: 'imagen-4.0-generate-preview-06-06',
-        params: {
-          prompt: 'A cute cat sitting in a garden',
-        },
-      };
-
-      // Act
-      await instance.createImage(payload);
-
-      // Assert
-      expect(instance['client'].models.generateImages).toHaveBeenCalledWith({
-        model: 'imagen-4.0-generate-preview-06-06',
-        prompt: 'A cute cat sitting in a garden',
-        config: {
-          aspectRatio: undefined,
-          numberOfImages: 1,
-        },
-      });
-    });
-
-    describe('Error handling', () => {
-      it('should throw InvalidProviderAPIKey error when API key is invalid', async () => {
-        // Arrange - Use real Google AI error format
-        const message = `[GoogleGenerativeAI Error]: Error fetching from https://generativelanguage.googleapis.com/v1/models/imagen-4.0:generateImages: [400 Bad Request] API key not valid. Please pass a valid API key. [{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"API_KEY_INVALID","domain":"googleapis.com","metadata":{"service":"generativelanguage.googleapis.com"}}]`;
-        const apiError = new Error(message);
-        vi.spyOn(instance['client'].models, 'generateImages').mockRejectedValue(apiError);
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-4.0-generate-preview-06-06',
-          params: {
-            prompt: 'A realistic landscape photo',
-          },
-        };
-
-        // Act & Assert - Test error type rather than specific text
-        await expect(instance.createImage(payload)).rejects.toEqual(
-          expect.objectContaining({
-            errorType: invalidErrorType,
-            provider,
-          }),
-        );
-      });
-
-      it('should throw ProviderBizError for network and API errors', async () => {
-        // Arrange
-        const apiError = new Error('Network connection failed');
-        vi.spyOn(instance['client'].models, 'generateImages').mockRejectedValue(apiError);
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-4.0-generate-preview-06-06',
-          params: {
-            prompt: 'A digital art portrait',
-          },
-        };
-
-        // Act & Assert - Test error type and basic structure
-        await expect(instance.createImage(payload)).rejects.toEqual(
-          expect.objectContaining({
-            errorType: bizErrorType,
-            provider,
-            error: expect.objectContaining({
-              message: expect.any(String),
-            }),
-          }),
-        );
-      });
-
-      it('should throw error when API response is malformed - missing generatedImages', async () => {
-        // Arrange
-        const mockImageResponse = {};
-        vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-          mockImageResponse as any,
-        );
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-4.0-generate-preview-06-06',
-          params: {
-            prompt: 'Abstract geometric patterns',
-          },
-        };
-
-        // Act & Assert - Test error behavior rather than specific text
-        await expect(instance.createImage(payload)).rejects.toEqual(
-          expect.objectContaining({
-            errorType: bizErrorType,
-            provider,
-          }),
-        );
-      });
-
-      it('should throw error when API response contains empty image array', async () => {
-        // Arrange
-        const mockImageResponse = {
-          generatedImages: [],
-        };
-        vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-          mockImageResponse as any,
-        );
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-4.0-generate-preview-06-06',
-          params: {
-            prompt: 'Minimalist design poster',
-          },
-        };
-
-        // Act & Assert
-        await expect(instance.createImage(payload)).rejects.toEqual(
-          expect.objectContaining({
-            errorType: bizErrorType,
-            provider,
-          }),
-        );
-      });
-
-      it('should throw error when generated image lacks required data', async () => {
-        // Arrange
-        const mockImageResponse = {
-          generatedImages: [
-            {
-              image: {}, // Missing imageBytes
-            },
-          ],
-        };
-        vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-          mockImageResponse as any,
-        );
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-4.0-generate-preview-06-06',
-          params: {
-            prompt: 'Watercolor painting style',
-          },
-        };
-
-        // Act & Assert
-        await expect(instance.createImage(payload)).rejects.toEqual(
-          expect.objectContaining({
-            errorType: bizErrorType,
-            provider,
-          }),
-        );
-      });
-    });
-
-    describe('Edge cases', () => {
-      it('should return first image when API returns multiple generated images', async () => {
-        // Arrange - Use two different real base64 image data
-        const firstImageData =
-          'iVBORw0KGgoAAAANSUhEUgAAAAUAAAAFCAYAAACNbyblAAAAHElEQVQI12P4//8/w38GIAXDIBKE0DHxgljNBAAO9TXL0Y4OHwAAAABJRU5ErkJggg==';
-        const secondImageData =
-          'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
-        const mockImageResponse = {
-          generatedImages: [
-            {
-              image: {
-                imageBytes: firstImageData,
-              },
-            },
-            {
-              image: {
-                imageBytes: secondImageData,
-              },
-            },
-          ],
-        };
-        vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-          mockImageResponse as any,
-        );
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-4.0-generate-preview-06-06',
-          params: {
-            prompt: 'Generate multiple variations of a sunset',
-          },
-        };
-
-        // Act
-        const result = await instance.createImage(payload);
-
-        // Assert - Should return the first image
-        expect(result).toEqual({
-          imageUrl: `data:image/png;base64,${firstImageData}`,
-        });
-      });
-
-      it('should work with custom future Imagen model versions', async () => {
-        // Arrange
-        const realBase64Data =
-          'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
-        const mockImageResponse = {
-          generatedImages: [
-            {
-              image: {
-                imageBytes: realBase64Data,
-              },
-            },
-          ],
-        };
-        vi.spyOn(instance['client'].models, 'generateImages').mockResolvedValue(
-          mockImageResponse as any,
-        );
-
-        const payload: CreateImagePayload = {
-          model: 'imagen-5.0-future-model',
-          params: {
-            prompt: 'Photorealistic portrait with soft lighting',
-            aspectRatio: '4:3',
-          },
-        };
-
-        // Act
-        await instance.createImage(payload);
-
-        // Assert
-        expect(instance['client'].models.generateImages).toHaveBeenCalledWith({
-          model: 'imagen-5.0-future-model',
-          prompt: 'Photorealistic portrait with soft lighting',
-          config: {
-            aspectRatio: '4:3',
-            numberOfImages: 1,
-          },
-        });
-      });
-    });
-  });
 });

package/packages/model-runtime/src/google/index.ts

@@ -27,6 +27,7 @@ import { parseGoogleErrorMessage } from '../utils/googleErrorParser';
 import { StreamingResponse } from '../utils/response';
 import { GoogleGenerativeAIStream, VertexAIStream } from '../utils/streams';
 import { parseDataUri } from '../utils/uriParser';
+import { createGoogleImage } from './createImage';
 
 const modelsOffSafetySettings = new Set(['gemini-2.0-flash-exp']);
 
@@ -258,49 +259,11 @@ export class LobeGoogleAI implements LobeRuntimeAI {
   }
 
   /**
-   * Generate images using Google AI Imagen API
+   * Generate images using Google AI Imagen API or Gemini Chat Models
    * @see https://ai.google.dev/gemini-api/docs/image-generation#imagen
    */
   async createImage(payload: CreateImagePayload): Promise<CreateImageResponse> {
-    try {
-      const { model, params } = payload;
-
-      const response = await this.client.models.generateImages({
-        config: {
-          aspectRatio: params.aspectRatio,
-          numberOfImages: 1,
-        },
-        model,
-        prompt: params.prompt,
-      });
-
-      if (!response.generatedImages || response.generatedImages.length === 0) {
-        throw new Error('No images generated');
-      }
-
-      const generatedImage = response.generatedImages[0];
-      if (!generatedImage.image || !generatedImage.image.imageBytes) {
-        throw new Error('Invalid image data');
-      }
-
-      const { imageBytes } = generatedImage.image;
-      // 1. official doc use png as example
-      // 2. no responseType param support like openai now.
-      // I think we can just hard code png now
-      const imageUrl = `data:image/png;base64,${imageBytes}`;
-
-      return { imageUrl };
-    } catch (error) {
-      const err = error as Error;
-      console.error('Google AI image generation error:', err);
-
-      const { errorType, error: parsedError } = parseGoogleErrorMessage(err.message);
-      throw AgentRuntimeError.createImage({
-        error: parsedError,
-        errorType,
-        provider: this.provider,
-      });
-    }
+    return createGoogleImage(this.client, this.provider, payload);
   }
 
   private createEnhancedStream(originalStream: any, signal: AbortSignal): ReadableStream {

package/packages/model-runtime/src/utils/modelParse.ts

@@ -111,6 +111,7 @@ export const IMAGE_MODEL_KEYWORDS = [
   'wanxiang',
   'DESCRIBE',
   'UPSCALE',
+  '!gemini', // 排除 gemini 模型，即使包含 -image 也是 chat 模型
   '-image',
   '^V3',
   '^V_2',
@@ -169,7 +170,7 @@ const findKnownModelByProvider = async (
 
   try {
     // 尝试动态导入对应的配置文件
-    const moduleImport = await import(`@/config/aiModels/${provider}`);
+    const moduleImport = await import(`@/config/aiModels/${provider}.ts`);
     const providerModels = moduleImport.default;
 
     // 如果导入成功且有数据，进行查找