feat: Add DALL-E reverse proxy settings and handle errors in image ge…

…neration (danny-avila#1173) * feat: Add DALL-E reverse proxy settings and handle errors in image generation * fix(ci): avoid importing extra utilities
aitok-ai · Nov 13, 2023 · 94fe949 · 94fe949
1 parent 464ff91
commit 94fe949
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 22 deletions.
diff --git a/app/clients/tools/DALL-E.js b/app/clients/tools/DALL-E.js
@@ -1,19 +1,30 @@
 // From https://platform.openai.com/docs/api-reference/images/create
 // To use this tool, you must pass in a configured OpenAIApi object.
 const fs = require('fs');
+const path = require('path');
 const OpenAI = require('openai');
 // const { genAzureEndpoint } = require('../../../utils/genAzureEndpoints');
 const { Tool } = require('langchain/tools');
+const { HttpsProxyAgent } = require('https-proxy-agent');
 const saveImageFromUrl = require('./saveImageFromUrl');
-const path = require('path');
+const extractBaseURL = require('../../../utils/extractBaseURL');
+const { DALLE_REVERSE_PROXY, PROXY } = process.env;
 
 class OpenAICreateImage extends Tool {
   constructor(fields = {}) {
     super();
 
     let apiKey = fields.DALLE_API_KEY || this.getApiKey();
+    const config = { apiKey };
+    if (DALLE_REVERSE_PROXY) {
+      config.baseURL = extractBaseURL(DALLE_REVERSE_PROXY);
+    }
+
+    if (PROXY) {
+      config.httpAgent = new HttpsProxyAgent(PROXY);
+    }
+
     // let azureKey = fields.AZURE_API_KEY || process.env.AZURE_API_KEY;
-    let config = { apiKey };
 
     // if (azureKey) {
     //   apiKey = azureKey;

diff --git a/app/clients/tools/structured/DALLE3.js b/app/clients/tools/structured/DALLE3.js
@@ -5,14 +5,24 @@ const path = require('path');
 const { z } = require('zod');
 const OpenAI = require('openai');
 const { Tool } = require('langchain/tools');
+const { HttpsProxyAgent } = require('https-proxy-agent');
 const saveImageFromUrl = require('../saveImageFromUrl');
-const { DALLE3_SYSTEM_PROMPT } = process.env;
+const extractBaseURL = require('../../../../utils/extractBaseURL');
+const { DALLE3_SYSTEM_PROMPT, DALLE_REVERSE_PROXY, PROXY } = process.env;
 class DALLE3 extends Tool {
   constructor(fields = {}) {
     super();
 
     let apiKey = fields.DALLE_API_KEY || this.getApiKey();
-    let config = { apiKey };
+    const config = { apiKey };
+    if (DALLE_REVERSE_PROXY) {
+      config.baseURL = extractBaseURL(DALLE_REVERSE_PROXY);
+    }
+
+    if (PROXY) {
+      config.httpAgent = new HttpsProxyAgent(PROXY);
+    }
+
     this.openai = new OpenAI(config);
     this.name = 'dalle';
     this.description = `Use DALLE to create images from text descriptions.
@@ -84,19 +94,30 @@ class DALLE3 extends Tool {
     if (!prompt) {
       throw new Error('Missing required field: prompt');
     }
-    const resp = await this.openai.images.generate({
-      model: 'dall-e-3',
-      quality,
-      style,
-      size,
-      prompt: this.replaceUnwantedChars(prompt),
-      n: 1,
-    });
+
+    let resp;
+    try {
+      resp = await this.openai.images.generate({
+        model: 'dall-e-3',
+        quality,
+        style,
+        size,
+        prompt: this.replaceUnwantedChars(prompt),
+        n: 1,
+      });
+    } catch (error) {
+      return `Something went wrong when trying to generate the image. The DALL-E API may unavailable:
+Error Message: ${error.message}`;
+    }
+
+    if (!resp) {
+      return 'Something went wrong when trying to generate the image. The DALL-E API may unavailable';
+    }
 
     const theImageUrl = resp.data[0].url;
 
     if (!theImageUrl) {
-      throw new Error('No image URL returned from OpenAI API.');
+      return 'No image URL returned from OpenAI API. There may be a problem with the API or your configuration.';
     }
 
     const regex = /img-[\w\d]+.png/;

diff --git a/app/clients/tools/structured/specs/DALLE3.spec.js b/app/clients/tools/structured/specs/DALLE3.spec.js
@@ -134,15 +134,6 @@ describe('DALLE3', () => {
     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',