From 4551abdce86bb0a960bb9464cb926338ea0f2910 Mon Sep 17 00:00:00 2001
From: Hk-Gosuto <yominogosuto@gmail.com>
Date: Fri, 29 Dec 2023 18:34:38 +0800
Subject: [PATCH] feat: dall-e plugin upgrade

---
 .../langchain-tools/dalle_image_generator.ts  | 37 +++++++++++++++++--
 1 file changed, 33 insertions(+), 4 deletions(-)

diff --git a/app/api/langchain-tools/dalle_image_generator.ts b/app/api/langchain-tools/dalle_image_generator.ts
index fe4d14eda..1ce20234f 100644
--- a/app/api/langchain-tools/dalle_image_generator.ts
+++ b/app/api/langchain-tools/dalle_image_generator.ts
@@ -41,20 +41,32 @@ export class DallEAPIWrapper extends StructuredTool {
     prompt: z
       .string()
       .describe(
-        'input must be a english prompt. you can set `quality: "hd"` for enhanced detail.',
+        "A text description of the desired image(s). input must be a english prompt.",
       ),
     size: z
       .enum(["1024x1024", "1024x1792", "1792x1024"])
       .default("1024x1024")
       .describe("images size"),
+    quality: z
+      .enum(["standard", "hd"])
+      .default("standard")
+      .describe(
+        "hd increases image detail and clarity at the cost of doubled consumption, and should not be used unless specified by the user.",
+      ),
+    style: z
+      .enum(["vivid", "natural"])
+      .default("vivid")
+      .describe(
+        "vivid leads to the creation of more intense and dramatic images, while Natural results in images that look more realistic and less exaggerated.",
+      ),
   });
 
   /** @ignore */
-  async _call({ prompt, size }: z.infer<typeof this.schema>) {
+  async _call({ prompt, size, quality, style }: z.infer<typeof this.schema>) {
     let imageUrl;
     const apiUrl = `${this.baseURL}/images/generations`;
     try {
-      const requestOptions = {
+      let requestOptions = {
         method: "POST",
         headers: {
           "Content-Type": "application/json",
@@ -65,8 +77,25 @@ export class DallEAPIWrapper extends StructuredTool {
           prompt: prompt,
           n: this.n,
           size: size,
+          quality: quality,
+          style: style,
         }),
       };
+      if (this.model != "dall-e-3") {
+        requestOptions = {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            Authorization: `Bearer ${this.apiKey}`,
+          },
+          body: JSON.stringify({
+            model: this.model,
+            prompt: prompt,
+            n: this.n,
+            size: size,
+          }),
+        };
+      }
       console.log(requestOptions);
       const response = await fetch(apiUrl, requestOptions);
       const json = await response.json();
@@ -98,5 +127,5 @@ export class DallEAPIWrapper extends StructuredTool {
     }
   }
 
-  description = `openai's dall-e 3 image generator.`;
+  description = `openai's dall-e image generator.`;
 }