mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-27 13:48:51 +01:00
🏞️ feat: Modifiable OpenAI Image Gen Model Environment Variable (#11082)
This commit is contained in:
parent
6ffb176056
commit
b7ea340769
3 changed files with 167 additions and 2 deletions
|
|
@ -78,6 +78,8 @@ function createOpenAIImageTools(fields = {}) {
|
|||
let apiKey = fields.IMAGE_GEN_OAI_API_KEY ?? getApiKey();
|
||||
const closureConfig = { apiKey };
|
||||
|
||||
const imageModel = process.env.IMAGE_GEN_OAI_MODEL || 'gpt-image-1';
|
||||
|
||||
let baseURL = 'https://api.openai.com/v1/';
|
||||
if (!override && process.env.IMAGE_GEN_OAI_BASEURL) {
|
||||
baseURL = extractBaseURL(process.env.IMAGE_GEN_OAI_BASEURL);
|
||||
|
|
@ -157,7 +159,7 @@ function createOpenAIImageTools(fields = {}) {
|
|||
|
||||
resp = await openai.images.generate(
|
||||
{
|
||||
model: 'gpt-image-1',
|
||||
model: imageModel,
|
||||
prompt: replaceUnwantedChars(prompt),
|
||||
n: Math.min(Math.max(1, n), 10),
|
||||
background,
|
||||
|
|
@ -239,7 +241,7 @@ Error Message: ${error.message}`);
|
|||
}
|
||||
|
||||
const formData = new FormData();
|
||||
formData.append('model', 'gpt-image-1');
|
||||
formData.append('model', imageModel);
|
||||
formData.append('prompt', replaceUnwantedChars(prompt));
|
||||
// TODO: `mask` support
|
||||
// TODO: more than 1 image support
|
||||
|
|
|
|||
162
api/test/app/clients/tools/structured/OpenAIImageTools.test.js
Normal file
162
api/test/app/clients/tools/structured/OpenAIImageTools.test.js
Normal file
|
|
@ -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),
|
||||
);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue