@lobehub/chat 1.98.2 → 1.99.0

package/src/libs/model-runtime/ModelRuntime.test.ts

@@ -8,6 +8,7 @@ import { JWTPayload } from '@/const/auth';
 import { TraceNameMap } from '@/const/trace';
 import { AgentRuntime, ChatStreamPayload, LobeOpenAI, ModelProvider } from '@/libs/model-runtime';
 import { providerRuntimeMap } from '@/libs/model-runtime/runtimeMap';
+import { CreateImagePayload } from '@/libs/model-runtime/types/image';
 import { createTraceOptions } from '@/server/modules/AgentRuntime';
 
 import { AgentChatOptions } from './ModelRuntime';
@@ -241,4 +242,83 @@ describe('AgentRuntime', () => {
       });
     });
   });
+
+  describe('AgentRuntime createImage method', () => {
+    it('should run correctly', async () => {
+      const payload: CreateImagePayload = {
+        model: 'dall-e-3',
+        params: {
+          prompt: 'A beautiful sunset over mountains',
+          width: 1024,
+          height: 1024,
+        },
+      };
+
+      const mockResponse = {
+        imageUrl: 'https://example.com/image.jpg',
+        width: 1024,
+        height: 1024,
+      };
+
+      vi.spyOn(LobeOpenAI.prototype, 'createImage').mockResolvedValue(mockResponse);
+
+      const result = await mockModelRuntime.createImage(payload);
+
+      expect(LobeOpenAI.prototype.createImage).toHaveBeenCalledWith(payload);
+      expect(result).toBe(mockResponse);
+    });
+
+    it('should handle undefined createImage method gracefully', async () => {
+      const payload: CreateImagePayload = {
+        model: 'dall-e-3',
+        params: {
+          prompt: 'A beautiful sunset over mountains',
+          width: 1024,
+          height: 1024,
+        },
+      };
+
+      // Mock runtime without createImage method
+      const runtimeWithoutCreateImage = {
+        createImage: undefined,
+      };
+
+      // @ts-ignore - testing edge case
+      mockModelRuntime['_runtime'] = runtimeWithoutCreateImage;
+
+      const result = await mockModelRuntime.createImage(payload);
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('AgentRuntime models method', () => {
+    it('should run correctly', async () => {
+      const mockModels = [
+        { id: 'gpt-4', name: 'GPT-4' },
+        { id: 'gpt-3.5-turbo', name: 'GPT-3.5 Turbo' },
+      ];
+
+      vi.spyOn(LobeOpenAI.prototype, 'models').mockResolvedValue(mockModels);
+
+      const result = await mockModelRuntime.models();
+
+      expect(LobeOpenAI.prototype.models).toHaveBeenCalled();
+      expect(result).toBe(mockModels);
+    });
+
+    it('should handle undefined models method gracefully', async () => {
+      // Mock runtime without models method
+      const runtimeWithoutModels = {
+        models: undefined,
+      };
+
+      // @ts-ignore - testing edge case
+      mockModelRuntime['_runtime'] = runtimeWithoutModels;
+
+      const result = await mockModelRuntime.models();
+
+      expect(result).toBeUndefined();
+    });
+  });
 });

package/src/libs/model-runtime/ModelRuntime.ts

@@ -1,3 +1,4 @@
+import { log } from 'debug';
 import { ClientOptions } from 'openai';
 
 import type { TracePayload } from '@/const/trace';
@@ -17,6 +18,7 @@ import {
   TextToImagePayload,
   TextToSpeechPayload,
 } from './types';
+import { CreateImagePayload } from './types/image';
 
 export interface AgentChatOptions {
   enableTrace?: boolean;
@@ -61,13 +63,17 @@ class ModelRuntime {
    * ```
    */
   async chat(payload: ChatStreamPayload, options?: ChatMethodOptions) {
-    return this._runtime.chat(payload, options);
+    return this._runtime.chat!(payload, options);
   }
 
   async textToImage(payload: TextToImagePayload) {
     return this._runtime.textToImage?.(payload);
   }
 
+  async createImage(payload: CreateImagePayload) {
+    return this._runtime.createImage?.(payload);
+  }
+
   async models() {
     return this._runtime.models?.();
   }
@@ -106,6 +112,14 @@ class ModelRuntime {
         LobeCloudflareParams & { apiKey?: string; apiVersion?: string; baseURL?: string }
     >,
   ) {
+    log('Initializing runtime with provider: %s and params: %O', provider, params);
+    // @ts-expect-error ignore
+    if (providerRuntimeMap[provider]) {
+      log('Provider runtime map found for provider: %s', provider);
+    } else {
+      log('Provider runtime map not found for provider: %s', provider);
+    }
+
     // @ts-expect-error runtime map not include vertex so it will be undefined
     const providerAI = providerRuntimeMap[provider] ?? LobeOpenAI;
     const runtimeModel: LobeRuntimeAI = new providerAI(params);

package/src/libs/model-runtime/UniformRuntime/index.ts

@@ -1,3 +1,6 @@
+/**
+ * @see https://github.com/lobehub/lobe-chat/discussions/6563
+ */
 import { LobeRuntimeAI } from '../BaseAI';
 import { LobeOpenAI } from '../openai';
 import { providerRuntimeMap } from '../runtimeMap';
@@ -74,7 +77,7 @@ class UniformRuntime {
     try {
       const runtime = this.getRuntimeByModel(payload.model);
 
-      return await runtime.chat(payload, options);
+      return await runtime.chat!(payload, options);
     } catch (e) {
       if (this._options.chat?.handleError) {
         const error = this._options.chat.handleError(e);

package/src/libs/model-runtime/fal/index.test.ts

@@ -0,0 +1,442 @@
+// @vitest-environment node
+import { fal } from '@fal-ai/client';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { CreateImagePayload } from '@/libs/model-runtime/types/image';
+
+import { LobeFalAI } from './index';
+
+// Mock the fal client
+vi.mock('@fal-ai/client', () => ({
+  fal: {
+    config: vi.fn(),
+    subscribe: vi.fn(),
+  },
+}));
+
+// Get the mocked fal instance
+const mockFal = vi.mocked(fal);
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+const provider = 'fal';
+const bizErrorType = 'ProviderBizError';
+const invalidErrorType = 'InvalidProviderAPIKey';
+
+let instance: LobeFalAI;
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  instance = new LobeFalAI({ apiKey: 'test-api-key' });
+});
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('LobeFalAI', () => {
+  describe('init', () => {
+    it('should correctly initialize with an API key', () => {
+      const instance = new LobeFalAI({ apiKey: 'test_api_key' });
+      expect(instance).toBeInstanceOf(LobeFalAI);
+      expect(mockFal.config).toHaveBeenCalledWith({
+        credentials: 'test_api_key',
+      });
+    });
+
+    it('should throw InvalidProviderAPIKey if no apiKey is provided', () => {
+      expect(() => {
+        new LobeFalAI({});
+      }).toThrow();
+    });
+
+    it('should throw InvalidProviderAPIKey if apiKey is undefined', () => {
+      expect(() => {
+        new LobeFalAI({ apiKey: undefined });
+      }).toThrow();
+    });
+  });
+
+  describe('createImage', () => {
+    it('should create image successfully with basic parameters', async () => {
+      // Arrange
+      const mockImageResponse = {
+        requestId: 'test-request-id',
+        data: {
+          images: [
+            {
+              url: 'https://example.com/image.jpg',
+              width: 1024,
+              height: 1024,
+            },
+          ],
+        },
+      };
+      mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+      const payload: CreateImagePayload = {
+        model: 'flux/dev',
+        params: {
+          prompt: 'A beautiful landscape',
+          width: 1024,
+          height: 1024,
+        },
+      };
+
+      // Act
+      const result = await instance.createImage(payload);
+
+      // Assert
+      expect(mockFal.subscribe).toHaveBeenCalledWith('fal-ai/flux/dev', {
+        input: {
+          enable_safety_checker: false,
+          num_images: 1,
+          prompt: 'A beautiful landscape',
+          image_size: {
+            width: 1024,
+            height: 1024,
+          },
+        },
+      });
+      expect(result).toEqual({
+        imageUrl: 'https://example.com/image.jpg',
+        width: 1024,
+        height: 1024,
+      });
+    });
+
+    it('should map standard parameters to fal-specific parameters', async () => {
+      // Arrange
+      const mockImageResponse = {
+        requestId: 'test-request-id',
+        data: {
+          images: [
+            {
+              url: 'https://example.com/image.jpg',
+              width: 512,
+              height: 512,
+            },
+          ],
+        },
+      };
+      mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+      const payload: CreateImagePayload = {
+        model: 'flux/dev',
+        params: {
+          prompt: 'Test image',
+          width: 512,
+          height: 512,
+          steps: 20,
+          cfg: 7.5,
+          imageUrl: 'https://example.com/input.jpg',
+        },
+      };
+
+      // Act
+      await instance.createImage(payload);
+
+      // Assert
+      expect(mockFal.subscribe).toHaveBeenCalledWith('fal-ai/flux/dev', {
+        input: {
+          enable_safety_checker: false,
+          num_images: 1,
+          prompt: 'Test image',
+          image_size: {
+            width: 512,
+            height: 512,
+          },
+          num_inference_steps: 20,
+          guidance_scale: 7.5,
+          image_url: 'https://example.com/input.jpg',
+        },
+      });
+    });
+
+    it('should handle parameters without width and height', async () => {
+      // Arrange
+      const mockImageResponse = {
+        requestId: 'test-request-id',
+        data: {
+          images: [
+            {
+              url: 'https://example.com/image.jpg',
+              width: 1024,
+              height: 1024,
+            },
+          ],
+        },
+      };
+      mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+      const payload: CreateImagePayload = {
+        model: 'flux/schnell',
+        params: {
+          prompt: 'Simple test',
+          steps: 10,
+        },
+      };
+
+      // Act
+      await instance.createImage(payload);
+
+      // Assert
+      expect(mockFal.subscribe).toHaveBeenCalledWith('fal-ai/flux/schnell', {
+        input: {
+          enable_safety_checker: false,
+          num_images: 1,
+          prompt: 'Simple test',
+          num_inference_steps: 10,
+        },
+      });
+    });
+
+    it('should handle custom parameters that are not in the mapping', async () => {
+      // Arrange
+      const mockImageResponse = {
+        requestId: 'test-request-id',
+        data: {
+          images: [
+            {
+              url: 'https://example.com/image.jpg',
+              width: 768,
+              height: 768,
+            },
+          ],
+        },
+      };
+      mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+      const payload: CreateImagePayload = {
+        model: 'flux/dev',
+        params: {
+          prompt: 'Custom test',
+          width: 768,
+          height: 768,
+          seed: 12345,
+        } as any, // Use any to allow custom parameters
+      };
+
+      // Act
+      await instance.createImage(payload);
+
+      // Assert
+      expect(mockFal.subscribe).toHaveBeenCalledWith('fal-ai/flux/dev', {
+        input: {
+          enable_safety_checker: false,
+          num_images: 1,
+          prompt: 'Custom test',
+          image_size: {
+            width: 768,
+            height: 768,
+          },
+          seed: 12345,
+        },
+      });
+    });
+
+    it('should return only imageUrl when width and height are not provided in response', async () => {
+      // Arrange
+      const mockImageResponse = {
+        requestId: 'test-request-id',
+        data: {
+          images: [
+            {
+              url: 'https://example.com/image.jpg',
+            },
+          ],
+        },
+      };
+      mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+      const payload: CreateImagePayload = {
+        model: 'flux/dev',
+        params: {
+          prompt: 'Test without dimensions',
+        },
+      };
+
+      // Act
+      const result = await instance.createImage(payload);
+
+      // Assert
+      expect(result).toEqual({
+        imageUrl: 'https://example.com/image.jpg',
+      });
+    });
+
+    describe('Error handling', () => {
+      it('should throw InvalidProviderAPIKey on 401 error', async () => {
+        // Arrange
+        const apiError = new Error('Unauthorized') as Error & { status: number };
+        apiError.status = 401;
+        mockFal.subscribe.mockRejectedValue(apiError);
+
+        const payload: CreateImagePayload = {
+          model: 'flux/dev',
+          params: {
+            prompt: 'Test image',
+          },
+        };
+
+        // Act & Assert
+        await expect(instance.createImage(payload)).rejects.toEqual({
+          error: { error: apiError },
+          errorType: invalidErrorType,
+        });
+      });
+
+      it('should throw ProviderBizError on other errors', async () => {
+        // Arrange
+        const apiError = new Error('Some other error');
+        mockFal.subscribe.mockRejectedValue(apiError);
+
+        const payload: CreateImagePayload = {
+          model: 'flux/dev',
+          params: {
+            prompt: 'Test image',
+          },
+        };
+
+        // Act & Assert
+        await expect(instance.createImage(payload)).rejects.toEqual({
+          error: { error: apiError },
+          errorType: bizErrorType,
+        });
+      });
+
+      it('should throw ProviderBizError on non-401 status errors', async () => {
+        // Arrange
+        const apiError = new Error('Server error') as Error & { status: number };
+        apiError.status = 500;
+        mockFal.subscribe.mockRejectedValue(apiError);
+
+        const payload: CreateImagePayload = {
+          model: 'flux/dev',
+          params: {
+            prompt: 'Test image',
+          },
+        };
+
+        // Act & Assert
+        await expect(instance.createImage(payload)).rejects.toEqual({
+          error: { error: apiError },
+          errorType: bizErrorType,
+        });
+      });
+    });
+
+    describe('Edge cases', () => {
+      it('should handle empty params object', async () => {
+        // Arrange
+        const mockImageResponse = {
+          requestId: 'test-request-id',
+          data: {
+            images: [
+              {
+                url: 'https://example.com/image.jpg',
+              },
+            ],
+          },
+        };
+        mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+        const payload: CreateImagePayload = {
+          model: 'flux/dev',
+          params: {
+            prompt: 'Empty params test',
+          },
+        };
+
+        // Act
+        const result = await instance.createImage(payload);
+
+        // Assert
+        expect(mockFal.subscribe).toHaveBeenCalledWith('fal-ai/flux/dev', {
+          input: {
+            enable_safety_checker: false,
+            num_images: 1,
+            prompt: 'Empty params test',
+          },
+        });
+        expect(result).toEqual({
+          imageUrl: 'https://example.com/image.jpg',
+        });
+      });
+
+      it('should handle model with different format', async () => {
+        // Arrange
+        const mockImageResponse = {
+          requestId: 'test-request-id',
+          data: {
+            images: [
+              {
+                url: 'https://example.com/image.jpg',
+              },
+            ],
+          },
+        };
+        mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+        const payload: CreateImagePayload = {
+          model: 'some-custom-model',
+          params: {
+            prompt: 'Test with custom model',
+          },
+        };
+
+        // Act
+        await instance.createImage(payload);
+
+        // Assert
+        expect(mockFal.subscribe).toHaveBeenCalledWith('fal-ai/some-custom-model', {
+          input: {
+            enable_safety_checker: false,
+            num_images: 1,
+            prompt: 'Test with custom model',
+          },
+        });
+      });
+
+      it('should handle response with multiple images (take first one)', async () => {
+        // Arrange
+        const mockImageResponse = {
+          requestId: 'test-request-id',
+          data: {
+            images: [
+              {
+                url: 'https://example.com/image1.jpg',
+                width: 1024,
+                height: 1024,
+              },
+              {
+                url: 'https://example.com/image2.jpg',
+                width: 512,
+                height: 512,
+              },
+            ],
+          },
+        };
+        mockFal.subscribe.mockResolvedValue(mockImageResponse as any);
+
+        const payload: CreateImagePayload = {
+          model: 'flux/dev',
+          params: {
+            prompt: 'Multiple images test',
+          },
+        };
+
+        // Act
+        const result = await instance.createImage(payload);
+
+        // Assert
+        expect(result).toEqual({
+          imageUrl: 'https://example.com/image1.jpg',
+          width: 1024,
+          height: 1024,
+        });
+      });
+    });
+  });
+});

package/src/libs/model-runtime/fal/index.ts

@@ -0,0 +1,88 @@
+import { fal } from '@fal-ai/client';
+import debug from 'debug';
+import { pick } from 'lodash-es';
+import { ClientOptions } from 'openai';
+
+import { RuntimeImageGenParamsValue } from '@/libs/standard-parameters/meta-schema';
+
+import { LobeRuntimeAI } from '../BaseAI';
+import { AgentRuntimeErrorType } from '../error';
+import { CreateImagePayload, CreateImageResponse } from '../types/image';
+import { AgentRuntimeError } from '../utils/createError';
+
+// Create debug logger
+const log = debug('lobe-image:fal');
+
+type FluxDevOutput = Awaited<ReturnType<typeof fal.subscribe<'fal-ai/flux/dev'>>>['data'];
+
+export class LobeFalAI implements LobeRuntimeAI {
+  constructor({ apiKey }: ClientOptions = {}) {
+    if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
+
+    fal.config({
+      credentials: apiKey,
+    });
+    log('FalAI initialized with apiKey: %s', apiKey ? '*****' : 'Not set');
+  }
+
+  async createImage(payload: CreateImagePayload): Promise<CreateImageResponse> {
+    const { model, params } = payload;
+    log('Creating image with model: %s and params: %O', model, params);
+
+    const paramsMap = new Map<RuntimeImageGenParamsValue, string>([
+      ['steps', 'num_inference_steps'],
+      ['cfg', 'guidance_scale'],
+      ['imageUrl', 'image_url'],
+    ]);
+
+    const defaultInput = {
+      enable_safety_checker: false,
+      num_images: 1,
+    };
+    const userInput = Object.fromEntries(
+      (Object.entries(params) as [keyof typeof params, any][]).map(([key, value]) => [
+        paramsMap.get(key) ?? key,
+        value,
+      ]),
+    );
+
+    if ('width' in userInput && 'height' in userInput) {
+      userInput.image_size = {
+        height: userInput.height,
+        width: userInput.width,
+      };
+      delete userInput.width;
+      delete userInput.height;
+    }
+
+    const endpoint = `fal-ai/${model}`;
+    log('Calling fal.subscribe with endpoint: %s and input: %O', endpoint, {
+      ...defaultInput,
+      ...userInput,
+    });
+
+    try {
+      const { data } = await fal.subscribe(endpoint, {
+        input: {
+          ...defaultInput,
+          ...userInput,
+        },
+      });
+      const image = (data as FluxDevOutput).images[0];
+      log('Received image data: %O', image);
+
+      return {
+        imageUrl: image.url,
+        ...pick(image, ['width', 'height']),
+      };
+    } catch (error) {
+      if (error instanceof Error && 'status' in error && (error as any).status === 401) {
+        throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey, {
+          error,
+        });
+      }
+
+      throw AgentRuntimeError.createError(AgentRuntimeErrorType.ProviderBizError, { error });
+    }
+  }
+}