🏞️ feat: Modifiable OpenAI Image Gen Model Environment Variable (#11082)

api/test/app/clients/tools/structured/OpenAIImageTools.test.js

@@ -0,0 +1,162 @@
+const OpenAI = require('openai');
+const createOpenAIImageTools = require('~/app/clients/tools/structured/OpenAIImageTools');
+
+jest.mock('openai');
+jest.mock('@librechat/data-schemas', () => ({
+  logger: {
+    warn: jest.fn(),
+    error: jest.fn(),
+    debug: jest.fn(),
+  },
+}));
+
+jest.mock('@librechat/api', () => ({
+  logAxiosError: jest.fn(),
+  oaiToolkit: {
+    image_gen_oai: {
+      name: 'image_gen_oai',
+      description: 'Generate an image',
+      schema: {},
+    },
+    image_edit_oai: {
+      name: 'image_edit_oai',
+      description: 'Edit an image',
+      schema: {},
+    },
+  },
+  extractBaseURL: jest.fn((url) => url),
+}));
+
+jest.mock('~/server/services/Files/strategies', () => ({
+  getStrategyFunctions: jest.fn(),
+}));
+
+jest.mock('~/models', () => ({
+  getFiles: jest.fn().mockResolvedValue([]),
+}));
+
+describe('OpenAIImageTools - IMAGE_GEN_OAI_MODEL environment variable', () => {
+  let originalEnv;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    originalEnv = { ...process.env };
+
+    process.env.IMAGE_GEN_OAI_API_KEY = 'test-api-key';
+
+    OpenAI.mockImplementation(() => ({
+      images: {
+        generate: jest.fn().mockResolvedValue({
+          data: [
+            {
+              b64_json: 'base64-encoded-image-data',
+            },
+          ],
+        }),
+      },
+    }));
+  });
+
+  afterEach(() => {
+    process.env = originalEnv;
+  });
+
+  it('should use default model "gpt-image-1" when IMAGE_GEN_OAI_MODEL is not set', async () => {
+    delete process.env.IMAGE_GEN_OAI_MODEL;
+
+    const [imageGenTool] = createOpenAIImageTools({
+      isAgent: true,
+      override: false,
+      req: { user: { id: 'test-user' } },
+    });
+
+    const mockGenerate = jest.fn().mockResolvedValue({
+      data: [
+        {
+          b64_json: 'base64-encoded-image-data',
+        },
+      ],
+    });
+
+    OpenAI.mockImplementation(() => ({
+      images: {
+        generate: mockGenerate,
+      },
+    }));
+
+    await imageGenTool.func({ prompt: 'test prompt' });
+
+    expect(mockGenerate).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'gpt-image-1',
+      }),
+      expect.any(Object),
+    );
+  });
+
+  it('should use "gpt-image-1.5" when IMAGE_GEN_OAI_MODEL is set to "gpt-image-1.5"', async () => {
+    process.env.IMAGE_GEN_OAI_MODEL = 'gpt-image-1.5';
+
+    const mockGenerate = jest.fn().mockResolvedValue({
+      data: [
+        {
+          b64_json: 'base64-encoded-image-data',
+        },
+      ],
+    });
+
+    OpenAI.mockImplementation(() => ({
+      images: {
+        generate: mockGenerate,
+      },
+    }));
+
+    const [imageGenTool] = createOpenAIImageTools({
+      isAgent: true,
+      override: false,
+      req: { user: { id: 'test-user' } },
+    });
+
+    await imageGenTool.func({ prompt: 'test prompt' });
+
+    expect(mockGenerate).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'gpt-image-1.5',
+      }),
+      expect.any(Object),
+    );
+  });
+
+  it('should use custom model name from IMAGE_GEN_OAI_MODEL environment variable', async () => {
+    process.env.IMAGE_GEN_OAI_MODEL = 'custom-image-model';
+
+    const mockGenerate = jest.fn().mockResolvedValue({
+      data: [
+        {
+          b64_json: 'base64-encoded-image-data',
+        },
+      ],
+    });
+
+    OpenAI.mockImplementation(() => ({
+      images: {
+        generate: mockGenerate,
+      },
+    }));
+
+    const [imageGenTool] = createOpenAIImageTools({
+      isAgent: true,
+      override: false,
+      req: { user: { id: 'test-user' } },
+    });
+
+    await imageGenTool.func({ prompt: 'test prompt' });
+
+    expect(mockGenerate).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'custom-image-model',
+      }),
+      expect.any(Object),
+    );
+  });
+});