feat: add arxiv plugin (#55)

app/api/langchain-tools/arxiv.ts

@@ -0,0 +1,77 @@
+import { StructuredTool } from "langchain/tools";
+import { z } from "zod";
+
+export class ArxivAPIWrapper extends StructuredTool {
+  get lc_namespace() {
+    return [...super.lc_namespace, "test"];
+  }
+
+  name = "arxiv";
+  description = "Run Arxiv search and get the article information.";
+
+  SORT_BY = {
+    RELEVANCE: "relevance",
+    LAST_UPDATED_DATE: "lastUpdatedDate",
+    SUBMITTED_DATE: "submittedDate",
+  };
+
+  SORT_ORDER = {
+    ASCENDING: "ascending",
+    DESCENDING: "descending",
+  };
+
+  schema = z.object({
+    searchQuery: z
+      .string()
+      .describe("same as the search_query parameter rules of the arxiv API."),
+    sortBy: z
+      .string()
+      .describe('can be "relevance", "lastUpdatedDate", "submittedDate".'),
+    sortOrder: z
+      .string()
+      .describe('can be either "ascending" or "descending".'),
+    start: z
+      .number()
+      .default(0)
+      .describe("the index of the first returned result."),
+    maxResults: z
+      .number()
+      .default(10)
+      .describe("the number of results returned by the query."),
+  });
+
+  async _call({
+    searchQuery,
+    sortBy,
+    sortOrder,
+    start,
+    maxResults,
+  }: z.infer<typeof this.schema>) {
+    if (sortBy && !Object.values(this.SORT_BY).includes(sortBy)) {
+      throw new Error(
+        `unsupported sort by option. should be one of: ${Object.values(
+          this.SORT_BY,
+        ).join(" ")}`,
+      );
+    }
+    if (sortOrder && !Object.values(this.SORT_ORDER).includes(sortOrder)) {
+      throw new Error(
+        `unsupported sort order option. should be one of: ${Object.values(
+          this.SORT_ORDER,
+        ).join(" ")}`,
+      );
+    }
+    try {
+      let url = `http://export.arxiv.org/api/query?search_query=${searchQuery}&start=${start}&max_results=${maxResults}${
+        sortBy ? `&sortBy=${sortBy}` : ""
+      }${sortOrder ? `&sortOrder=${sortOrder}` : ""}`;
+      console.error("[arxiv]", url);
+      const response = await fetch(url);
+      const data = await response.text();
+      return data;
+    } catch (e) {
+      console.error("[arxiv]", e);
+    }
+    return "not found";
+  }
+}

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

@@ -22,6 +22,7 @@ import { DallEAPIWrapper } from "@/app/api/langchain-tools/dalle_image_generator
 import { BaiduSearch } from "@/app/api/langchain-tools/baidu_search";
 import { GoogleSearch } from "@/app/api/langchain-tools/google_search";
 import { StableDiffusionWrapper } from "@/app/api/langchain-tools/stable_diffusion_image_generator";
+import { ArxivAPIWrapper } from "@/app/api/langchain-tools/arxiv";
 
 const serverConfig = getServerSideConfig();
 
@@ -95,8 +96,8 @@ async function handle(req: NextRequest) {
 
     const handler = BaseCallbackHandler.fromMethods({
       async handleLLMNewToken(token: string) {
+        // console.log("[Token]", token);
         if (token) {
-          console.log("[Token]", token);
           var response = new ResponseBody();
           response.message = token;
           await writer.ready;
@@ -151,8 +152,7 @@ async function handle(req: NextRequest) {
           if (!reqBody.returnIntermediateSteps) return;
           var response = new ResponseBody();
           response.isToolMessage = true;
-          let toolInput = <ToolInput>(<unknown>action.toolInput);
-          response.message = toolInput.input;
+          response.message = JSON.stringify(action.toolInput);
           response.toolName = action.tool;
           await writer.ready;
           await writer.write(
@@ -230,12 +230,14 @@ async function handle(req: NextRequest) {
     const calculatorTool = new Calculator();
     const dallEAPITool = new DallEAPIWrapper(apiKey, baseUrl);
     const stableDiffusionTool = new StableDiffusionWrapper();
+    const arxivAPITool = new ArxivAPIWrapper();
     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);
     if (useTools.includes(stableDiffusionTool.name))
       tools.push(stableDiffusionTool);
+    if (useTools.includes(arxivAPITool.name)) tools.push(arxivAPITool);
 
     useTools.forEach((toolName) => {
       if (toolName) {

app/plugins/cn.ts

@@ -58,4 +58,13 @@ export const CN_PLUGINS: BuiltinPlugin[] = [
     createdAt: 1688899480510,
     enable: false,
   },
+  {
+    name: "Arxiv",
+    toolName: "arxiv",
+    lang: "cn",
+    description: "使用 Arxiv 接口搜索并获取文章信息。",
+    builtin: true,
+    createdAt: 1699265115000,
+    enable: false,
+  },
 ];

app/plugins/en.ts

@@ -60,4 +60,13 @@ export const EN_PLUGINS: BuiltinPlugin[] = [
     createdAt: 1688899480510,
     enable: false,
   },
+  {
+    name: "Arxiv",
+    toolName: "arxiv",
+    lang: "en",
+    description: "Arxiv search and get the article information.",
+    builtin: true,
+    createdAt: 1699265115000,
+    enable: false,
+  },
 ];