From 4044891f8509ec0aff2688e50a0bd7574978cc8e Mon Sep 17 00:00:00 2001
From: Hk-Gosuto <yominogosuto@gmail.com>
Date: Thu, 21 Dec 2023 12:36:57 +0800
Subject: [PATCH] feat: support custom dalle model

---
 README.md                                        |  1 +
 app/api/langchain-tools/dalle_image_generator.ts | 15 ++++++++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index ec8ba2bd9..bdb68f616 100644
--- a/README.md
+++ b/README.md
@@ -54,6 +54,7 @@
     - DALL-E 3
       - DALL-E 3 插件需要配置对象存储服务，请参考 [对象存储服务配置指南](./docs/s3-oss.md) 配置
       - 如无需图像转存则可以配置  `DALLE_NO_IMAGE_STORAGE=1` ，此时将直接将 DALL-E 服务返回的临时 URL 用于图像显示，注意：该链接具有时效性
+      - 默认使用 `dall-e-3` 模型，如果想使用 `dall-e-2` ，可以配置环境变量 `DALLE_MODEL=dall-e-2`
     - StableDiffusion
       - 本插件目前为测试版本，后续可能会有较大的变更，请谨慎使用
       - 使用本插件需要一定的专业知识，Stable Diffusion 本身的相关问题不在本项目的解答范围内，如果您确定要使用本插件请参考 [Stable Diffusion 插件配置指南](./docs/stable-diffusion-plugin-cn.md) 文档进行配置
diff --git a/app/api/langchain-tools/dalle_image_generator.ts b/app/api/langchain-tools/dalle_image_generator.ts
index a244a82bf..fe4d14eda 100644
--- a/app/api/langchain-tools/dalle_image_generator.ts
+++ b/app/api/langchain-tools/dalle_image_generator.ts
@@ -7,6 +7,7 @@ export class DallEAPIWrapper extends StructuredTool {
   n = 1;
   apiKey: string;
   baseURL?: string;
+  model: string;
 
   noStorage: boolean;
 
@@ -26,6 +27,7 @@ export class DallEAPIWrapper extends StructuredTool {
     this.callback = callback;
 
     this.noStorage = !!process.env.DALLE_NO_IMAGE_STORAGE;
+    this.model = process.env.DALLE_MODEL ?? "dall-e-3";
   }
 
   async saveImageFromUrl(url: string) {
@@ -59,18 +61,25 @@ export class DallEAPIWrapper extends StructuredTool {
           Authorization: `Bearer ${this.apiKey}`,
         },
         body: JSON.stringify({
-          model: "dall-e-3",
+          model: this.model,
           prompt: prompt,
           n: this.n,
           size: size,
         }),
       };
+      console.log(requestOptions);
       const response = await fetch(apiUrl, requestOptions);
       const json = await response.json();
-      console.log("[DALL-E]", json);
-      imageUrl = json.data[0].url;
+      try {
+        console.log("[DALL-E]", json);
+        imageUrl = json.data[0].url;
+      } catch (e) {
+        if (this.callback != null) await this.callback(JSON.stringify(json));
+        throw e;
+      }
     } catch (e) {
       console.error("[DALL-E]", e);
+      return (e as Error).message;
     }
     if (!imageUrl) return "No image was generated";
     try {