alpha version

app/api/langchain-tools/nodejs_tools.ts

@@ -7,16 +7,15 @@ import { StableDiffusionNodeWrapper } from "@/app/api/langchain-tools/stable_dif
 import { Calculator } from "langchain/tools/calculator";
 import { WebBrowser } from "langchain/tools/webbrowser";
 import { WolframAlphaTool } from "@/app/api/langchain-tools/wolframalpha";
+import { RAGSearch } from "./rag_search";
 
 export class NodeJSTool {
   private apiKey: string | undefined;
-
   private baseUrl: string;
-
   private model: BaseLanguageModel;
-
   private embeddings: Embeddings;
-
+  private sessionId: string;
+  private ragEmbeddings: Embeddings;
   private callback?: (data: string) => Promise<void>;
 
   constructor(
@@ -24,12 +23,16 @@ export class NodeJSTool {
     baseUrl: string,
     model: BaseLanguageModel,
     embeddings: Embeddings,
+    sessionId: string,
+    ragEmbeddings: Embeddings,
     callback?: (data: string) => Promise<void>,
   ) {
     this.apiKey = apiKey;
     this.baseUrl = baseUrl;
     this.model = model;
     this.embeddings = embeddings;
+    this.sessionId = sessionId;
+    this.ragEmbeddings = ragEmbeddings;
     this.callback = callback;
   }
 
@@ -57,6 +60,9 @@ export class NodeJSTool {
       wolframAlphaTool,
       pdfBrowserTool,
     ];
+    if (!!process.env.NEXT_PUBLIC_ENABLE_RAG) {
+      tools.push(new RAGSearch(this.sessionId, this.model, this.ragEmbeddings));
+    }
     return tools;
   }
 }

app/api/langchain-tools/rag_search.ts

@@ -0,0 +1,64 @@
+import { Tool } from "@langchain/core/tools";
+import { CallbackManagerForToolRun } from "@langchain/core/callbacks/manager";
+import { BaseLanguageModel } from "langchain/dist/base_language";
+import { formatDocumentsAsString } from "langchain/util/document";
+import { Embeddings } from "langchain/dist/embeddings/base.js";
+import { RunnableSequence } from "@langchain/core/runnables";
+import { StringOutputParser } from "@langchain/core/output_parsers";
+import { Pinecone } from "@pinecone-database/pinecone";
+import { PineconeStore } from "@langchain/pinecone";
+
+export class RAGSearch extends Tool {
+  static lc_name() {
+    return "RAGSearch";
+  }
+
+  get lc_namespace() {
+    return [...super.lc_namespace, "ragsearch"];
+  }
+
+  private sessionId: string;
+  private model: BaseLanguageModel;
+  private embeddings: Embeddings;
+
+  constructor(
+    sessionId: string,
+    model: BaseLanguageModel,
+    embeddings: Embeddings,
+  ) {
+    super();
+    this.sessionId = sessionId;
+    this.model = model;
+    this.embeddings = embeddings;
+  }
+
+  /** @ignore */
+  async _call(inputs: string, runManager?: CallbackManagerForToolRun) {
+    const pinecone = new Pinecone();
+    const pineconeIndex = pinecone.Index(process.env.PINECONE_INDEX!);
+    const vectorStore = await PineconeStore.fromExistingIndex(this.embeddings, {
+      pineconeIndex,
+    });
+
+    let context;
+    const returnCunt = process.env.RAG_RETURN_COUNT
+      ? parseInt(process.env.RAG_RETURN_COUNT, 10)
+      : 4;
+    const results = await vectorStore.similaritySearch(inputs, returnCunt, {
+      sessionId: this.sessionId,
+    });
+    context = formatDocumentsAsString(results);
+    console.log("[rag-search]", context);
+    return context;
+    // const input = `Text:${context}\n\nQuestion:${inputs}\n\nI need you to answer the question based on the text.`;
+
+    // console.log("[rag-search]", input);
+
+    // const chain = RunnableSequence.from([this.model, new StringOutputParser()]);
+    // return chain.invoke(input, runManager?.getChild());
+  }
+
+  name = "rag-search";
+
+  description = `It is used to query documents entered by the user.The input content is the keywords extracted from the user's question, and multiple keywords are separated by spaces and passed in.`;
+}