feat: allow config of DALL-E-3 System Prompt via env 🎨 (danny-avila#1150

)
ywkim · Nov 7, 2023 · 43d7a75 · 43d7a75
1 parent 4f3b667
commit 43d7a75
Show file tree

Hide file tree

Showing 4 changed files with 214 additions and 7 deletions.
diff --git a/.env.example b/.env.example
@@ -132,6 +132,12 @@ DEBUG_OPENAI=false # Set to true to enable debug mode for the OpenAI endpoint
 # This may be the case for LocalAI with some models. To do so, uncomment the following:
 # OPENAI_FORCE_PROMPT=true
 
+# (Advanced) For customization of the DALL-E-3 System prompt, 
+# uncomment the following, and provide your own prompt:
+# See official prompt for reference:
+# https://github.com/spdustin/ChatGPT-AutoExpert/blob/main/_system-prompts/dall-e.md
+# DALLE3_SYSTEM_PROMPT="Your System Prompt here"
+
 ##########################
 # OpenRouter (overrides OpenAI and Plugins Endpoints): 
 ##########################

diff --git a/api/app/clients/tools/DALL-E.js b/api/app/clients/tools/DALL-E.js
@@ -58,7 +58,7 @@ Guidelines:
   replaceUnwantedChars(inputString) {
     return inputString
       .replace(/\r\n|\r|\n/g, ' ')
-      .replace('"', '')
+      .replace(/"/g, '')
       .trim();
   }
 

diff --git a/api/app/clients/tools/structured/DALLE3.js b/api/app/clients/tools/structured/DALLE3.js
@@ -1,12 +1,12 @@
 // From https://platform.openai.com/docs/guides/images/usage?context=node
 // To use this tool, you must pass in a configured OpenAIApi object.
 const fs = require('fs');
+const path = require('path');
 const { z } = require('zod');
 const OpenAI = require('openai');
 const { Tool } = require('langchain/tools');
 const saveImageFromUrl = require('../saveImageFromUrl');
-const path = require('path');
-
+const { DALLE3_SYSTEM_PROMPT } = process.env;
 class DALLE3 extends Tool {
   constructor(fields = {}) {
     super();
@@ -17,11 +17,13 @@ class DALLE3 extends Tool {
     this.name = 'dalle';
     this.description = `Use DALLE to create images from text descriptions.
     - It requires prompts to be in English, detailed, and to specify image type and human features for diversity.
-    - Only one image is produced per call, without repeating or listing descriptions outside the "prompts" field.
+    - Create only one image, without repeating or listing descriptions outside the "prompts" field.
     - Maintains the original intent of the description, with parameters for image style, quality, and size to tailor the output.`;
-    this.description_for_model = `// Whenever a description of an image is given, generate prompts (following these rules), and use dalle to create the image. If the user does not ask for a specific number of images, default to creating 2 prompts to send to dalle that are written to be as diverse as possible. All prompts sent to dalle must abide by the following policies:
+    this.description_for_model =
+      DALLE3_SYSTEM_PROMPT ??
+      `// Whenever a description of an image is given, generate prompts (following these rules), and use dalle to create the image. If the user does not ask for a specific number of images, default to creating 2 prompts to send to dalle that are written to be as diverse as possible. All prompts sent to dalle must abide by the following policies:
     // 1. Prompts must be in English. Translate to English if needed.
-    // 2. Only one image can be created per function call.
+    // 2. One image per function call. Create only 1 image per request unless explicitly told to generate more than 1 image.
     // 3. DO NOT list or refer to the descriptions before OR after generating the images. They should ONLY ever be written out ONCE, in the \`"prompts"\` field of the request. You do not need to ask for permission to generate, just do it!
     // 4. Always mention the image type (photo, oil painting, watercolor painting, illustration, cartoon, drawing, vector, render, etc.) at the beginning of the caption. Unless the captions suggests otherwise, make one of the images a photo.
     // 5. Diversify depictions of ALL images with people to always include always DESCENT and GENDER for EACH person using direct terms. Adjust only human descriptions.
@@ -65,7 +67,7 @@ class DALLE3 extends Tool {
   replaceUnwantedChars(inputString) {
     return inputString
       .replace(/\r\n|\r|\n/g, ' ')
-      .replace('"', '')
+      .replace(/"/g, '')
       .trim();
   }
 

diff --git a/api/app/clients/tools/structured/specs/DALLE3.spec.js b/api/app/clients/tools/structured/specs/DALLE3.spec.js
@@ -0,0 +1,199 @@
+const fs = require('fs');
+const path = require('path');
+const OpenAI = require('openai');
+const DALLE3 = require('../DALLE3');
+const saveImageFromUrl = require('../../saveImageFromUrl');
+
+jest.mock('openai');
+
+const generate = jest.fn();
+OpenAI.mockImplementation(() => ({
+  images: {
+    generate,
+  },
+}));
+
+jest.mock('fs', () => {
+  return {
+    existsSync: jest.fn(),
+    mkdirSync: jest.fn(),
+  };
+});
+
+jest.mock('../../saveImageFromUrl', () => {
+  return jest.fn();
+});
+
+jest.mock('path', () => {
+  return {
+    resolve: jest.fn(),
+    join: jest.fn(),
+    relative: jest.fn(),
+  };
+});
+
+describe('DALLE3', () => {
+  let originalEnv;
+  let dalle; // Keep this declaration if you need to use dalle in other tests
+  const mockApiKey = 'mock_api_key';
+
+  beforeAll(() => {
+    // Save the original process.env
+    originalEnv = { ...process.env };
+  });
+
+  beforeEach(() => {
+    // Reset the process.env before each test
+    jest.resetModules();
+    process.env = { ...originalEnv, DALLE_API_KEY: mockApiKey };
+    // Instantiate DALLE3 for tests that do not depend on DALLE3_SYSTEM_PROMPT
+    dalle = new DALLE3();
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+    // Restore the original process.env after each test
+    process.env = originalEnv;
+  });
+
+  it('should throw an error if DALLE_API_KEY is missing', () => {
+    delete process.env.DALLE_API_KEY;
+    expect(() => new DALLE3()).toThrow('Missing DALLE_API_KEY environment variable.');
+  });
+
+  it('should replace unwanted characters in input string', () => {
+    const input = 'This is a test\nstring with "quotes" and new lines.';
+    const expectedOutput = 'This is a test string with quotes and new lines.';
+    expect(dalle.replaceUnwantedChars(input)).toBe(expectedOutput);
+  });
+
+  it('should generate markdown image URL correctly', () => {
+    const imageName = 'test.png';
+    path.join.mockReturnValue('images/test.png');
+    path.relative.mockReturnValue('images/test.png');
+    const markdownImage = dalle.getMarkdownImageUrl(imageName);
+    expect(markdownImage).toBe('![generated image](/images/test.png)');
+  });
+
+  it('should call OpenAI API with correct parameters', async () => {
+    const mockData = {
+      prompt: 'A test prompt',
+      quality: 'standard',
+      size: '1024x1024',
+      style: 'vivid',
+    };
+
+    const mockResponse = {
+      data: [
+        {
+          url: 'http://example.com/img-test.png',
+        },
+      ],
+    };
+
+    generate.mockResolvedValue(mockResponse);
+    saveImageFromUrl.mockResolvedValue(true);
+    fs.existsSync.mockReturnValue(true);
+    path.resolve.mockReturnValue('/fakepath/images');
+    path.join.mockReturnValue('/fakepath/images/img-test.png');
+    path.relative.mockReturnValue('images/img-test.png');
+
+    const result = await dalle._call(mockData);
+
+    expect(generate).toHaveBeenCalledWith({
+      model: 'dall-e-3',
+      quality: mockData.quality,
+      style: mockData.style,
+      size: mockData.size,
+      prompt: mockData.prompt,
+      n: 1,
+    });
+    expect(result).toContain('![generated image]');
+  });
+
+  it('should use the system prompt if provided', () => {
+    process.env.DALLE3_SYSTEM_PROMPT = 'System prompt for testing';
+    jest.resetModules(); // This will ensure the module is fresh and will read the new env var
+    const DALLE3 = require('../DALLE3'); // Re-require after setting the env var
+    const dalleWithSystemPrompt = new DALLE3();
+    expect(dalleWithSystemPrompt.description_for_model).toBe('System prompt for testing');
+  });
+
+  it('should not use the system prompt if not provided', async () => {
+    delete process.env.DALLE3_SYSTEM_PROMPT;
+    const dalleWithoutSystemPrompt = new DALLE3();
+    expect(dalleWithoutSystemPrompt.description_for_model).not.toBe('System prompt for testing');
+  });
+
+  it('should throw an error if prompt is missing', async () => {
+    const mockData = {
+      quality: 'standard',
+      size: '1024x1024',
+      style: 'vivid',
+    };
+    await expect(dalle._call(mockData)).rejects.toThrow('Missing required field: prompt');
+  });
+
+  it('should throw an error if no image URL is returned from OpenAI API', async () => {
+    const mockData = {
+      prompt: 'A test prompt',
+    };
+    // Simulate a response with an object that has a `url` property set to `undefined`
+    generate.mockResolvedValue({ data: [{ url: undefined }] });
+    await expect(dalle._call(mockData)).rejects.toThrow('No image URL returned from OpenAI API.');
+  });
+
+  it('should log to console if no image name is found in the URL', async () => {
+    const mockData = {
+      prompt: 'A test prompt',
+    };
+    const mockResponse = {
+      data: [
+        {
+          url: 'http://example.com/invalid-url',
+        },
+      ],
+    };
+    console.log = jest.fn(); // Mock console.log
+    generate.mockResolvedValue(mockResponse);
+    await dalle._call(mockData);
+    expect(console.log).toHaveBeenCalledWith('No image name found in the string.');
+  });
+
+  it('should create the directory if it does not exist', async () => {
+    const mockData = {
+      prompt: 'A test prompt',
+    };
+    const mockResponse = {
+      data: [
+        {
+          url: 'http://example.com/img-test.png',
+        },
+      ],
+    };
+    generate.mockResolvedValue(mockResponse);
+    fs.existsSync.mockReturnValue(false); // Simulate directory does not exist
+    await dalle._call(mockData);
+    expect(fs.mkdirSync).toHaveBeenCalledWith(expect.any(String), { recursive: true });
+  });
+
+  it('should log an error and return the image URL if there is an error saving the image', async () => {
+    const mockData = {
+      prompt: 'A test prompt',
+    };
+    const mockResponse = {
+      data: [
+        {
+          url: 'http://example.com/img-test.png',
+        },
+      ],
+    };
+    const error = new Error('Error while saving the image');
+    generate.mockResolvedValue(mockResponse);
+    saveImageFromUrl.mockRejectedValue(error);
+    console.error = jest.fn(); // Mock console.error
+    const result = await dalle._call(mockData);
+    expect(console.error).toHaveBeenCalledWith('Error while saving the image:', error);
+    expect(result).toBe(mockResponse.data[0].url);
+  });
+});