feat: support local storage

app/api/langchain-tools/dalle_image_generator_node.ts

@@ -0,0 +1,22 @@
+import { getServerSideConfig } from "@/app/config/server";
+import { DallEAPIWrapper } from "./dalle_image_generator";
+import S3FileStorage from "@/app/utils/s3_file_storage";
+import LocalFileStorage from "@/app/utils/local_file_storage";
+
+export class DallEAPINodeWrapper extends DallEAPIWrapper {
+  async saveImageFromUrl(url: string) {
+    const response = await fetch(url);
+    const content = await response.arrayBuffer();
+    const buffer = Buffer.from(content);
+
+    var filePath = "";
+    const serverConfig = getServerSideConfig();
+    var fileName = `${Date.now()}.png`;
+    if (serverConfig.isStoreFileToLocal) {
+      filePath = await LocalFileStorage.put(fileName, buffer);
+    } else {
+      filePath = await S3FileStorage.put(fileName, buffer);
+    }
+    return filePath;
+  }
+}

app/api/langchain-tools/edge_tools.ts

@@ -4,13 +4,8 @@ import { StableDiffusionWrapper } from "@/app/api/langchain-tools/stable_diffusi
 import { BaseLanguageModel } from "langchain/dist/base_language";
 import { Calculator } from "langchain/tools/calculator";
 import { WebBrowser } from "langchain/tools/webbrowser";
-import { BaiduSearch } from "@/app/api/langchain-tools/baidu_search";
-import { DuckDuckGo } from "@/app/api/langchain-tools/duckduckgo_search";
-import { GoogleSearch } from "@/app/api/langchain-tools/google_search";
-import { Tool, DynamicTool } from "langchain/tools";
-import * as langchainTools from "langchain/tools";
 import { Embeddings } from "langchain/dist/embeddings/base.js";
-import { WolframAlphaTool } from "./wolframalpha";
+import { WolframAlphaTool } from "@/app/api/langchain-tools/wolframalpha";
 
 export class EdgeTool {
   private apiKey: string | undefined;
@@ -52,7 +47,6 @@ export class EdgeTool {
     const arxivAPITool = new ArxivAPIWrapper();
     const wolframAlphaTool = new WolframAlphaTool();
     let tools = [
-      // searchTool,
       calculatorTool,
       webBrowserTool,
       dallEAPITool,

app/api/langchain-tools/nodejs_tools.ts

@@ -1,7 +1,12 @@
 import { BaseLanguageModel } from "langchain/dist/base_language";
 import { PDFBrowser } from "@/app/api/langchain-tools/pdf_browser";
-
 import { Embeddings } from "langchain/dist/embeddings/base.js";
+import { ArxivAPIWrapper } from "@/app/api/langchain-tools/arxiv";
+import { DallEAPINodeWrapper } from "@/app/api/langchain-tools/dalle_image_generator_node";
+import { StableDiffusionNodeWrapper } from "@/app/api/langchain-tools/stable_diffusion_image_generator_node";
+import { Calculator } from "langchain/tools/calculator";
+import { WebBrowser } from "langchain/tools/webbrowser";
+import { WolframAlphaTool } from "@/app/api/langchain-tools/wolframalpha";
 
 export class NodeJSTool {
   private apiKey: string | undefined;
@@ -29,7 +34,29 @@ export class NodeJSTool {
   }
 
   async getCustomTools(): Promise<any[]> {
+    const webBrowserTool = new WebBrowser({
+      model: this.model,
+      embeddings: this.embeddings,
+    });
+    const calculatorTool = new Calculator();
+    const dallEAPITool = new DallEAPINodeWrapper(
+      this.apiKey,
+      this.baseUrl,
+      this.callback,
+    );
+    const stableDiffusionTool = new StableDiffusionNodeWrapper();
+    const arxivAPITool = new ArxivAPIWrapper();
+    const wolframAlphaTool = new WolframAlphaTool();
     const pdfBrowserTool = new PDFBrowser(this.model, this.embeddings);
-    return [pdfBrowserTool];
+    let tools = [
+      calculatorTool,
+      webBrowserTool,
+      dallEAPITool,
+      stableDiffusionTool,
+      arxivAPITool,
+      wolframAlphaTool,
+      pdfBrowserTool,
+    ];
+    return tools;
   }
 }

app/api/langchain-tools/stable_diffusion_image_generator.ts

@@ -8,6 +8,13 @@ export class StableDiffusionWrapper extends Tool {
     super();
   }
 
+  async saveImage(imageBase64: string) {
+    const buffer = Buffer.from(imageBase64, "base64");
+    var fileName = `${Date.now()}.png`;
+    const filePath = await S3FileStorage.put(fileName, buffer);
+    return filePath;
+  }
+
   /** @ignore */
   async _call(prompt: string) {
     let url = process.env.STABLE_DIFFUSION_API_URL;
@@ -40,8 +47,7 @@ export class StableDiffusionWrapper extends Tool {
     const json = await response.json();
     let imageBase64 = json.images[0];
     if (!imageBase64) return "No image was generated";
-    const buffer = Buffer.from(imageBase64, "base64");
-    const filePath = await S3FileStorage.put(`${Date.now()}.png`, buffer);
+    const filePath = await this.saveImage(imageBase64);
     console.log(`[${this.name}]`, filePath);
     return filePath;
   }

app/api/langchain-tools/stable_diffusion_image_generator_node.ts

@@ -0,0 +1,19 @@
+import S3FileStorage from "@/app/utils/s3_file_storage";
+import { StableDiffusionWrapper } from "./stable_diffusion_image_generator";
+import { getServerSideConfig } from "@/app/config/server";
+import LocalFileStorage from "@/app/utils/local_file_storage";
+
+export class StableDiffusionNodeWrapper extends StableDiffusionWrapper {
+  async saveImage(imageBase64: string) {
+    var filePath = "";
+    var fileName = `${Date.now()}.png`;
+    const buffer = Buffer.from(imageBase64, "base64");
+    const serverConfig = getServerSideConfig();
+    if (serverConfig.isStoreFileToLocal) {
+      filePath = await LocalFileStorage.put(fileName, buffer);
+    } else {
+      filePath = await S3FileStorage.put(fileName, buffer);
+    }
+    return filePath;
+  }
+}