feat(node-parser): support async function (#1682)

.changeset/selfish-kids-own.md

@@ -0,0 +1,7 @@
+---
+"@llamaindex/core": patch
+"llamaindex": patch
+"@llamaindex/core-tests": patch
+---
+
+feat(node-parser): support async function

packages/core/src/node-parser/base.ts

@@ -7,21 +7,27 @@ import {
   TextNode,
   TransformComponent,
 } from "../schema";
+import { isPromise } from "../utils";
 
-export abstract class NodeParser extends TransformComponent<BaseNode[]> {
+export abstract class NodeParser<
+  Result extends TextNode[] | Promise<TextNode[]> =
+    | TextNode[]
+    | Promise<TextNode[]>,
+> extends TransformComponent<Result> {
   includeMetadata: boolean = true;
   includePrevNextRel: boolean = true;
 
   constructor() {
-    super((nodes: BaseNode[]): BaseNode[] => {
+    super((nodes: BaseNode[]): Result => {
+      // alex: should we fix `as` type?
       return this.getNodesFromDocuments(nodes as TextNode[]);
     });
   }
 
   protected postProcessParsedNodes(
-    nodes: TextNode[],
+    nodes: Awaited<Result>,
     parentDocMap: Map<string, TextNode>,
-  ): TextNode[] {
+  ): Awaited<Result> {
     nodes.forEach((node, i) => {
       const parentDoc = parentDocMap.get(node.sourceNode?.nodeId || "");
 
@@ -73,9 +79,9 @@ export abstract class NodeParser extends TransformComponent<BaseNode[]> {
   protected abstract parseNodes(
     documents: TextNode[],
     showProgress?: boolean,
-  ): TextNode[];
+  ): Result;
 
-  public getNodesFromDocuments(documents: TextNode[]): TextNode[] {
+  public getNodesFromDocuments(documents: TextNode[]): Result {
     const docsId: Map<string, TextNode> = new Map(
       documents.map((doc) => [doc.id_, doc]),
     );
@@ -85,20 +91,36 @@ export abstract class NodeParser extends TransformComponent<BaseNode[]> {
       documents,
     });
 
-    const nodes = this.postProcessParsedNodes(
-      this.parseNodes(documents),
-      docsId,
-    );
+    const parsedNodes = this.parseNodes(documents);
+    if (isPromise(parsedNodes)) {
+      return parsedNodes.then((parsedNodes) => {
+        const nodes = this.postProcessParsedNodes(
+          parsedNodes as Awaited<Result>,
+          docsId,
+        );
 
-    callbackManager.dispatchEvent("node-parsing-end", {
-      nodes,
-    });
+        callbackManager.dispatchEvent("node-parsing-end", {
+          nodes,
+        });
 
-    return nodes;
+        return nodes;
+      }) as Result;
+    } else {
+      const nodes = this.postProcessParsedNodes(
+        parsedNodes as Awaited<Result>,
+        docsId,
+      );
+
+      callbackManager.dispatchEvent("node-parsing-end", {
+        nodes,
+      });
+
+      return nodes;
+    }
   }
 }
 
-export abstract class TextSplitter extends NodeParser {
+export abstract class TextSplitter extends NodeParser<TextNode[]> {
   abstract splitText(text: string): string[];
 
   public splitTexts(texts: string[]): string[] {

packages/core/src/node-parser/markdown.ts

@@ -6,7 +6,7 @@ import {
 } from "../schema";
 import { NodeParser } from "./base";
 
-export class MarkdownNodeParser extends NodeParser {
+export class MarkdownNodeParser extends NodeParser<TextNode[]> {
   override parseNodes(nodes: TextNode[], showProgress?: boolean): TextNode[] {
     return nodes.reduce<TextNode[]>((allNodes, node) => {
       const markdownNodes = this.getNodesFromNode(node);

packages/core/src/node-parser/sentence-window.ts

@@ -9,7 +9,7 @@ import {
 import { NodeParser } from "./base";
 import { splitBySentenceTokenizer, type TextSplitterFn } from "./utils";
 
-export class SentenceWindowNodeParser extends NodeParser {
+export class SentenceWindowNodeParser extends NodeParser<TextNode[]> {
   static DEFAULT_WINDOW_SIZE = 3;
   static DEFAULT_WINDOW_METADATA_KEY = "window";
   static DEFAULT_ORIGINAL_TEXT_METADATA_KEY = "originalText";

packages/core/src/utils/index.ts

@@ -1,5 +1,9 @@
 import type { JSONValue } from "../global";
 
+export const isPromise = <T>(obj: unknown): obj is Promise<T> => {
+  return obj != null && typeof obj === "object" && "then" in obj;
+};
+
 export const isAsyncIterable = (
   obj: unknown,
 ): obj is AsyncIterable<unknown> => {

packages/core/tests/node-parser/node-parser.test.ts

@@ -0,0 +1,24 @@
+import { NodeParser } from "@llamaindex/core/node-parser";
+import { TextNode } from "@llamaindex/core/schema";
+import { describe, expect, test } from "vitest";
+
+describe("NodeParser", () => {
+  test("node parser should allow async parse function", async () => {
+    class MyNodeParser extends NodeParser<Promise<TextNode[]>> {
+      protected async parseNodes(documents: TextNode[]): Promise<TextNode[]> {
+        await new Promise((resolve) => setTimeout(resolve, 1000));
+        return documents;
+      }
+    }
+
+    const nodeParser = new MyNodeParser();
+    const nodes = [
+      new TextNode({
+        text: "Hello, world!",
+      }),
+    ];
+    const result = nodeParser(nodes);
+    expect(result).toBeInstanceOf(Promise);
+    await expect(result).resolves.toEqual(nodes);
+  });
+});

packages/llamaindex/src/indices/keyword/index.ts

@@ -296,7 +296,7 @@ export class KeywordTableIndex extends BaseIndex<KeywordTable> {
       await docStore.setDocumentHash(doc.id_, doc.hash);
     }
 
-    const nodes = Settings.nodeParser.getNodesFromDocuments(documents);
+    const nodes = await Settings.nodeParser.getNodesFromDocuments(documents);
     const index = await KeywordTableIndex.init({
       nodes,
       storageContext,

packages/llamaindex/src/indices/summary/index.ts

@@ -145,7 +145,7 @@ export class SummaryIndex extends BaseIndex<IndexList> {
       await docStore.setDocumentHash(doc.id_, doc.hash);
     }
 
-    const nodes = Settings.nodeParser.getNodesFromDocuments(documents);
+    const nodes = await Settings.nodeParser.getNodesFromDocuments(documents);
 
     const index = await SummaryIndex.init({
       nodes,