feat: optimize rag

app/api/langchain-tools/myfiles_browser.ts

@@ -0,0 +1,78 @@
+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 { getServerSideConfig } from "@/app/config/server";
+import { SupabaseVectorStore } from "@langchain/community/vectorstores/supabase";
+import { createClient } from "@supabase/supabase-js";
+import { z } from "zod";
+import { StructuredTool } from "@langchain/core/tools";
+
+export class MyFilesBrowser extends StructuredTool {
+  static lc_name() {
+    return "MyFilesBrowser";
+  }
+
+  get lc_namespace() {
+    return [...super.lc_namespace, "myfilesbrowser"];
+  }
+
+  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;
+  }
+
+  schema = z.object({
+    queries: z.array(z.string()).describe("A query list."),
+  });
+
+  /** @ignore */
+  async _call({ queries }: z.infer<typeof this.schema>) {
+    const serverConfig = getServerSideConfig();
+    if (!serverConfig.isEnableRAG)
+      throw new Error("env ENABLE_RAG not configured");
+
+    const privateKey = process.env.SUPABASE_PRIVATE_KEY;
+    if (!privateKey) throw new Error(`Expected env var SUPABASE_PRIVATE_KEY`);
+
+    const url = process.env.SUPABASE_URL;
+    if (!url) throw new Error(`Expected env var SUPABASE_URL`);
+    const client = createClient(url, privateKey);
+    const vectorStore = new SupabaseVectorStore(this.embeddings, {
+      client,
+      tableName: "documents",
+      queryName: "match_documents",
+    });
+
+    let context;
+    const returnCunt = serverConfig.ragReturnCount
+      ? parseInt(serverConfig.ragReturnCount, 10)
+      : 4;
+    console.log("[myfiles_browser]", { queries, returnCunt });
+    let documents: any[] = [];
+    for (var i = 0; i < queries.length; i++) {
+      let results = await vectorStore.similaritySearch(queries[i], returnCunt, {
+        sessionId: this.sessionId,
+      });
+      results.forEach((item) => documents.push(item));
+    }
+    context = formatDocumentsAsString(documents);
+    console.log("[myfiles_browser]", { context });
+    return context;
+  }
+
+  name = "myfiles_browser";
+
+  description = `queries to a search over the file(s) uploaded in the current conversation and displays the results.`;
+}

app/api/langchain-tools/nodejs_tools.ts

@@ -10,7 +10,7 @@ import { WolframAlphaTool } from "@/app/api/langchain-tools/wolframalpha";
 import { BilibiliVideoInfoTool } from "./bilibili_vid_info";
 import { BilibiliVideoSearchTool } from "./bilibili_vid_search";
 import { BilibiliMusicRecognitionTool } from "./bilibili_music_recognition";
-import { RAGSearch } from "./rag_search";
+import { MyFilesBrowser } from "./myfiles_browser";
 import { BilibiliVideoConclusionTool } from "./bilibili_vid_conclusion";
 
 export class NodeJSTool {
@@ -59,7 +59,7 @@ export class NodeJSTool {
     const bilibiliVideoSearchTool = new BilibiliVideoSearchTool();
     const bilibiliVideoConclusionTool = new BilibiliVideoConclusionTool();
     const bilibiliMusicRecognitionTool = new BilibiliMusicRecognitionTool();
-    let tools = [
+    let tools: any = [
       calculatorTool,
       webBrowserTool,
       dallEAPITool,
@@ -73,7 +73,9 @@ export class NodeJSTool {
       bilibiliVideoConclusionTool,
     ];
     if (!!process.env.ENABLE_RAG) {
-      tools.push(new RAGSearch(this.sessionId, this.model, this.ragEmbeddings));
+      tools.push(
+        new MyFilesBrowser(this.sessionId, this.model, this.ragEmbeddings),
+      );
     }
     return tools;
   }

app/api/langchain-tools/rag_search.ts

@@ -1,79 +0,0 @@
-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";
-import { getServerSideConfig } from "@/app/config/server";
-import { QdrantVectorStore } from "@langchain/community/vectorstores/qdrant";
-
-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 serverConfig = getServerSideConfig();
-    if (!serverConfig.isEnableRAG)
-      throw new Error("env ENABLE_RAG not configured");
-    // const pinecone = new Pinecone();
-    // const pineconeIndex = pinecone.Index(serverConfig.pineconeIndex!);
-    // const vectorStore = await PineconeStore.fromExistingIndex(this.embeddings, {
-    //   pineconeIndex,
-    // });
-    const vectorStore = await QdrantVectorStore.fromExistingCollection(
-      this.embeddings,
-      {
-        url: process.env.QDRANT_URL,
-        apiKey: process.env.QDRANT_API_KEY,
-        collectionName: this.sessionId,
-      },
-    );
-
-    let context;
-    const returnCunt = serverConfig.ragReturnCount
-      ? parseInt(serverConfig.ragReturnCount, 10)
-      : 4;
-    console.log("[rag-search]", { inputs, returnCunt });
-    // const results = await vectorStore.similaritySearch(inputs, returnCunt, {
-    //   sessionId: this.sessionId,
-    // });
-    const results = await vectorStore.similaritySearch(inputs, returnCunt);
-    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.`;
-}

app/api/langchain/rag/store/route.ts

@@ -20,7 +20,10 @@ import { FileInfo } from "@/app/client/platforms/utils";
 import mime from "mime";
 import LocalFileStorage from "@/app/utils/local_file_storage";
 import S3FileStorage from "@/app/utils/s3_file_storage";
-import { QdrantVectorStore } from "@langchain/community/vectorstores/qdrant";
+import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
+import { SupabaseVectorStore } from "@langchain/community/vectorstores/supabase";
+import { createClient } from "@supabase/supabase-js";
+import { Embeddings } from "langchain/dist/embeddings/base";
 
 interface RequestBody {
   sessionId: string;
@@ -67,6 +70,11 @@ async function handle(req: NextRequest) {
   if (req.method === "OPTIONS") {
     return NextResponse.json({ body: "OK" }, { status: 200 });
   }
+  const privateKey = process.env.SUPABASE_PRIVATE_KEY;
+  if (!privateKey) throw new Error(`Expected env var SUPABASE_PRIVATE_KEY`);
+  const url = process.env.SUPABASE_URL;
+  if (!url) throw new Error(`Expected env var SUPABASE_URL`);
+
   try {
     const authResult = auth(req, ModelProvider.GPT);
     if (authResult.error) {
@@ -81,18 +89,25 @@ async function handle(req: NextRequest) {
     const apiKey = getOpenAIApiKey(token);
     const baseUrl = getOpenAIBaseUrl(reqBody.baseUrl);
     const serverConfig = getServerSideConfig();
-    // const pinecone = new Pinecone();
-    // const pineconeIndex = pinecone.Index(serverConfig.pineconeIndex!);
-    const embeddings = new OpenAIEmbeddings(
-      {
-        modelName: serverConfig.ragEmbeddingModel,
-        openAIApiKey: apiKey,
-      },
-      { basePath: baseUrl },
-    );
+    let embeddings: Embeddings;
+    if (process.env.OLLAMA_BASE_URL) {
+      embeddings = new OllamaEmbeddings({
+        model: serverConfig.ragEmbeddingModel,
+        baseUrl: process.env.OLLAMA_BASE_URL,
+      });
+    } else {
+      embeddings = new OpenAIEmbeddings(
+        {
+          modelName: serverConfig.ragEmbeddingModel,
+          openAIApiKey: apiKey,
+        },
+        { basePath: baseUrl },
+      );
+    }
     // https://js.langchain.com/docs/integrations/vectorstores/pinecone
     // https://js.langchain.com/docs/integrations/vectorstores/qdrant
     // process files
+    let partial = "";
     for (let i = 0; i < reqBody.fileInfos.length; i++) {
       const fileInfo = reqBody.fileInfos[i];
       const contentType = mime.getType(fileInfo.fileName);
@@ -134,26 +149,25 @@ async function handle(req: NextRequest) {
         chunkOverlap: chunkOverlap,
       });
       const splits = await textSplitter.splitDocuments(docs);
-      const vectorStore = await QdrantVectorStore.fromDocuments(
+      const client = createClient(url, privateKey);
+      const vectorStore = await SupabaseVectorStore.fromDocuments(
         splits,
         embeddings,
         {
-          url: process.env.QDRANT_URL,
-          apiKey: process.env.QDRANT_API_KEY,
-          collectionName: reqBody.sessionId,
+          client,
+          tableName: "documents",
+          queryName: "match_documents",
         },
       );
-      // await PineconeStore.fromDocuments(splits, embeddings, {
-      //   pineconeIndex,
-      //   maxConcurrency: 5,
-      // });
-      // const vectorStore = await PineconeStore.fromExistingIndex(embeddings, {
-      //   pineconeIndex,
-      // });
+      partial = splits
+        .slice(0, 2)
+        .map((v) => v.pageContent)
+        .join("\n");
     }
     return NextResponse.json(
       {
         sessionId: reqBody.sessionId,
+        partial: partial,
       },
       {
         status: 200,

app/api/langchain/tool/agent/nodejs/route.ts

@@ -4,6 +4,8 @@ import { auth } from "@/app/api/auth";
 import { NodeJSTool } from "@/app/api/langchain-tools/nodejs_tools";
 import { ModelProvider } from "@/app/constant";
 import { OpenAI, OpenAIEmbeddings } from "@langchain/openai";
+import { Embeddings } from "langchain/dist/embeddings/base";
+import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
 
 async function handle(req: NextRequest) {
   if (req.method === "OPTIONS") {
@@ -44,13 +46,22 @@ async function handle(req: NextRequest) {
       },
       { basePath: baseUrl },
     );
-    const ragEmbeddings = new OpenAIEmbeddings(
-      {
-        modelName: process.env.RAG_EMBEDDING_MODEL ?? "text-embedding-3-large",
-        openAIApiKey: apiKey,
-      },
-      { basePath: baseUrl },
-    );
+    let ragEmbeddings: Embeddings;
+    if (process.env.OLLAMA_BASE_URL) {
+      ragEmbeddings = new OllamaEmbeddings({
+        model: process.env.RAG_EMBEDDING_MODEL,
+        baseUrl: process.env.OLLAMA_BASE_URL,
+      });
+    } else {
+      ragEmbeddings = new OpenAIEmbeddings(
+        {
+          modelName:
+            process.env.RAG_EMBEDDING_MODEL ?? "text-embedding-3-large",
+          openAIApiKey: apiKey,
+        },
+        { basePath: baseUrl },
+      );
+    }
 
     var dalleCallback = async (data: string) => {
       var response = new ResponseBody();