From 6a0791f0a2312ffb4e4707a3cf8ee4e14bd06f72 Mon Sep 17 00:00:00 2001
From: Hk-Gosuto <yominogosuto@gmail.com>
Date: Thu, 14 Sep 2023 23:16:18 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0dall-e=E6=8F=92?=
 =?UTF-8?q?=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore                                    |  4 +-
 .../langchain-tools/dalle_image_generator.ts  | 65 +++++++++++++++++++
 app/api/langchain/tool/agent/route.ts         |  3 +
 app/plugins/cn.ts                             |  9 +++
 app/plugins/en.ts                             | 10 +++
 package.json                                  |  2 +
 yarn.lock                                     | 23 ++++++-
 7 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 app/api/langchain-tools/dalle_image_generator.ts

diff --git a/.gitignore b/.gitignore
index b00b0e325..a769884d0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -43,4 +43,6 @@ dev
 .env
 
 *.key
-*.key.pub
\ No newline at end of file
+*.key.pub
+
+/public/uploads
\ No newline at end of file
diff --git a/app/api/langchain-tools/dalle_image_generator.ts b/app/api/langchain-tools/dalle_image_generator.ts
new file mode 100644
index 000000000..c8b37a980
--- /dev/null
+++ b/app/api/langchain-tools/dalle_image_generator.ts
@@ -0,0 +1,65 @@
+import * as fs from "fs";
+import * as path from "path";
+import axios from "axios";
+
+import { Tool } from "langchain/tools";
+import OpenAI from "openai";
+
+export class DallEAPIWrapper extends Tool {
+  name = "dalle_image_generator";
+  n = 1;
+  size: "256x256" | "512x512" | "1024x1024" | null = "1024x1024";
+  apiKey: string;
+  baseURL?: string;
+
+  constructor(apiKey?: string | undefined, baseURL?: string | undefined) {
+    super();
+    if (!apiKey) {
+      throw new Error("OpenAI API key not set.");
+    }
+    this.apiKey = apiKey;
+    this.baseURL = baseURL;
+  }
+
+  async saveImageFromUrl(url: string) {
+    const response = await axios.get(url, { responseType: "arraybuffer" });
+
+    const uploadsDir = "public/uploads";
+    console.log("[fileUpload]", { uploadsDir });
+    if (!fs.existsSync(uploadsDir)) {
+      fs.mkdirSync(uploadsDir, { recursive: true });
+    }
+
+    const filename = `${Date.now()}.png`;
+    const filePath = path.join(uploadsDir, filename);
+
+    fs.writeFileSync(filePath, Buffer.from(response.data, "binary"));
+
+    return `uploads/${filename}`;
+  }
+
+  /** @ignore */
+  async _call(prompt: string) {
+    const openai = new OpenAI({
+      apiKey: this.apiKey,
+      baseURL: this.baseURL,
+    });
+    const response = await openai.images.generate({
+      prompt: prompt,
+      n: this.n,
+      size: this.size,
+    });
+
+    let image_url = response.data[0].url;
+    console.log(image_url);
+    if (!image_url) return "No image was generated";
+    let filePath = await this.saveImageFromUrl(image_url);
+    console.log(filePath);
+    return filePath;
+  }
+
+  description = `openai's dall-e image generator.
+    input must be a english prompt.
+    output will be the image link url.
+    use markdown to display images.`;
+}
diff --git a/app/api/langchain/tool/agent/route.ts b/app/api/langchain/tool/agent/route.ts
index 5e8f9a665..5a31a2ea8 100644
--- a/app/api/langchain/tool/agent/route.ts
+++ b/app/api/langchain/tool/agent/route.ts
@@ -18,6 +18,7 @@ import { DuckDuckGo } from "@/app/api/langchain-tools/duckduckgo_search";
 import { WebBrowser } from "langchain/tools/webbrowser";
 import { Calculator } from "langchain/tools/calculator";
 import { DynamicTool, Tool } from "langchain/tools";
+import { DallEAPIWrapper } from "@/app/api/langchain-tools/dalle_image_generator";
 
 const serverConfig = getServerSideConfig();
 
@@ -214,9 +215,11 @@ async function handle(req: NextRequest) {
     ];
     const webBrowserTool = new WebBrowser({ model, embeddings });
     const calculatorTool = new Calculator();
+    const dallEAPITool = new DallEAPIWrapper(apiKey, baseUrl);
     if (useTools.includes("web-search")) tools.push(searchTool);
     if (useTools.includes(webBrowserTool.name)) tools.push(webBrowserTool);
     if (useTools.includes(calculatorTool.name)) tools.push(calculatorTool);
+    if (useTools.includes(dallEAPITool.name)) tools.push(dallEAPITool);
 
     useTools.forEach((toolName) => {
       if (toolName) {
diff --git a/app/plugins/cn.ts b/app/plugins/cn.ts
index a18ad1229..4df7e7144 100644
--- a/app/plugins/cn.ts
+++ b/app/plugins/cn.ts
@@ -38,4 +38,13 @@ export const CN_PLUGINS: BuiltinPlugin[] = [
     createdAt: 1694235989000,
     enable: false,
   },
+  {
+    name: "DALL·E",
+    toolName: "dalle_image_generator",
+    lang: "cn",
+    description: "DALL·E 可以根据自然语言的描述创建逼真的图像和艺术。",
+    builtin: true,
+    createdAt: 1694703673000,
+    enable: false,
+  },
 ];
diff --git a/app/plugins/en.ts b/app/plugins/en.ts
index c7508cd72..852262e83 100644
--- a/app/plugins/en.ts
+++ b/app/plugins/en.ts
@@ -40,4 +40,14 @@ export const EN_PLUGINS: BuiltinPlugin[] = [
     createdAt: 1694235989000,
     enable: false,
   },
+  {
+    name: "DALL·E",
+    toolName: "dalle_image_generator",
+    lang: "en",
+    description:
+      "DALL·E 2 is an AI system that can create realistic images and art from a description in natural language.",
+    builtin: true,
+    createdAt: 1694703673000,
+    enable: false,
+  },
 ];
diff --git a/package.json b/package.json
index a9cbf2ee1..20088853d 100644
--- a/package.json
+++ b/package.json
@@ -24,6 +24,7 @@
     "cheerio": "^1.0.0-rc.12",
     "duck-duck-scrape": "^2.2.4",
     "emoji-picker-react": "^4.5.1",
+    "encoding": "^0.1.13",
     "fuse.js": "^6.6.2",
     "html-entities": "^2.4.0",
     "html-to-image": "^1.11.11",
@@ -34,6 +35,7 @@
     "nanoid": "^4.0.2",
     "next": "^13.4.9",
     "node-fetch": "^3.3.1",
+    "openai": "^4.6.0",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-markdown": "^8.0.7",
diff --git a/yarn.lock b/yarn.lock
index e58eed5e3..c3419b64b 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -3005,6 +3005,13 @@ emoji-regex@^9.2.2:
   resolved "https://registry.yarnpkg.com/emoji-regex/-/emoji-regex-9.2.2.tgz#840c8803b0d8047f4ff0cf963176b32d4ef3ed72"
   integrity sha512-L18DaJsXSUk2+42pv8mLs5jJT2hqFkFE4j21wOmgbUqsZ2hL72NsUU785g9RXgo3s0ZNgVl42TiHp3ZtOv/Vyg==
 
+encoding@^0.1.13:
+  version "0.1.13"
+  resolved "https://registry.yarnpkg.com/encoding/-/encoding-0.1.13.tgz#56574afdd791f54a8e9b2785c0582a2d26210fa9"
+  integrity sha512-ETBauow1T35Y/WZMkio9jiM0Z5xjHHmJ4XmjZOq1l/dXz3lr2sRn87nJy20RupqSh1F2m3HHPSp8ShIPQJrJ3A==
+  dependencies:
+    iconv-lite "^0.6.2"
+
 enhanced-resolve@^5.12.0:
   version "5.12.0"
   resolved "https://registry.yarnpkg.com/enhanced-resolve/-/enhanced-resolve-5.12.0.tgz#300e1c90228f5b570c4d35babf263f6da7155634"
@@ -3944,7 +3951,7 @@ husky@^8.0.0:
   resolved "https://registry.yarnpkg.com/husky/-/husky-8.0.3.tgz#4936d7212e46d1dea28fef29bb3a108872cd9184"
   integrity sha512-+dQSyqPh4x1hlO1swXBiNb2HzTDN1I2IGLQx1GrBuiqFJfoMrnZWwVmatvSiO+Iz8fBUnf+lekwNo4c2LlXItg==
 
-iconv-lite@0.6, iconv-lite@^0.6.3:
+iconv-lite@0.6, iconv-lite@^0.6.2, iconv-lite@^0.6.3:
   version "0.6.3"
   resolved "https://registry.yarnpkg.com/iconv-lite/-/iconv-lite-0.6.3.tgz#a52f80bf38da1952eb5c681790719871a1a72501"
   integrity sha512-4fCk79wshMdzMp2rH06qWrJE4iolqLhCUH+OiuIgU++RB0+94NlDL81atO7GX55uUKueo0txHNtvEyI6D7WdMw==
@@ -5347,6 +5354,20 @@ openai@^3.3.0:
     axios "^0.26.0"
     form-data "^4.0.0"
 
+openai@^4.6.0:
+  version "4.6.0"
+  resolved "https://registry.yarnpkg.com/openai/-/openai-4.6.0.tgz#0335111cf71c608b68dbdca19f8ccbd9cfb0cd97"
+  integrity sha512-LuONkTgoe4D172raQCv+eEK5OdLGnY/M4JrUz/pxRGevZwqDqy3xhBbCeWX8QLCbFcnITYsu/VBJXZJ0rDAMpA==
+  dependencies:
+    "@types/node" "^18.11.18"
+    "@types/node-fetch" "^2.6.4"
+    abort-controller "^3.0.0"
+    agentkeepalive "^4.2.1"
+    digest-fetch "^1.3.0"
+    form-data-encoder "1.7.2"
+    formdata-node "^4.3.2"
+    node-fetch "^2.6.7"
+
 openapi-types@^12.1.3:
   version "12.1.3"
   resolved "https://registry.yarnpkg.com/openapi-types/-/openapi-types-12.1.3.tgz#471995eb26c4b97b7bd356aacf7b91b73e777dd3"