changeset

Add observer/callback feature to SimpleDirectoryReader

.changeset/shiny-lions-clap.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Add observer/filter to the SimpleDirectoryReader (thanks @mtutty)

.gitignore

@@ -3,6 +3,7 @@
 # dependencies
 node_modules
 .pnp
+.pnpm-store
 .pnp.js
 
 # testing

apps/simple/directory.ts

@@ -0,0 +1,24 @@
+import { SimpleDirectoryReader } from "llamaindex";
+
+function callback(
+  category: string,
+  name: string,
+  status: any,
+  message?: string,
+): boolean {
+  console.log(category, name, status, message);
+  if (name.endsWith(".pdf")) {
+    console.log("I DON'T WANT PDF FILES!");
+    return false;
+  }
+  return true;
+}
+
+async function main() {
+  // Load page
+  const reader = new SimpleDirectoryReader(callback);
+  const params = { directoryPath: "./data" };
+  await reader.loadData(params);
+}
+
+main().catch(console.error);

packages/core/package.json

@@ -5,6 +5,7 @@
   "dependencies": {
     "@anthropic-ai/sdk": "^0.8.1",
     "@notionhq/client": "^2.2.13",
+    "js-tiktoken": "^1.0.7",
     "lodash": "^4.17.21",
     "mammoth": "^1.6.0",
     "md-utils-ts": "^2.0.0",
@@ -17,7 +18,6 @@
     "rake-modified": "^1.0.8",
     "replicate": "^0.20.1",
     "string-strip-html": "^13.4.3",
-    "tiktoken": "^1.0.10",
     "uuid": "^9.0.1",
     "wink-nlp": "^1.14.3"
   },

packages/core/src/GlobalsHelper.ts

@@ -1,5 +1,4 @@
-import cl100k_base from "tiktoken/encoders/cl100k_base.json";
-import { Tiktoken } from "tiktoken/lite";
+import { encodingForModel, TiktokenModel } from "js-tiktoken";
 
 import { v4 as uuidv4 } from "uuid";
 import { Event, EventTag, EventType } from "./callbacks/CallbackManager";
@@ -18,18 +17,17 @@ class GlobalsHelper {
   } | null = null;
 
   private initDefaultTokenizer() {
-    const encoding = new Tiktoken(
-      cl100k_base.bpe_ranks,
-      cl100k_base.special_tokens,
-      cl100k_base.pat_str,
-    );
+    const encoding = encodingForModel("text-embedding-ada-002"); // cl100k_base
 
     this.defaultTokenizer = {
       encode: (text: string) => {
-        return encoding.encode(text);
+        return new Uint32Array(encoding.encode(text));
       },
       decode: (tokens: Uint32Array) => {
-        return new TextDecoder().decode(encoding.decode(tokens));
+        const numberArray = Array.from(tokens);
+        const text = encoding.decode(numberArray);
+        const uint8Array = new TextEncoder().encode(text);
+        return new TextDecoder().decode(uint8Array);      
       },
     };
   }
@@ -41,10 +39,10 @@ class GlobalsHelper {
     if (!this.defaultTokenizer) {
       this.initDefaultTokenizer();
     }
-
+  
     return this.defaultTokenizer!.encode.bind(this.defaultTokenizer);
   }
-
+  
   tokenizerDecoder(encoding?: string) {
     if (encoding && encoding !== Tokenizers.CL100K_BASE) {
       throw new Error(`Tokenizer encoding ${encoding} not yet supported`);
@@ -52,7 +50,7 @@ class GlobalsHelper {
     if (!this.defaultTokenizer) {
       this.initDefaultTokenizer();
     }
-
+  
     return this.defaultTokenizer!.decode.bind(this.defaultTokenizer);
   }
 

packages/core/src/readers/SimpleDirectoryReader.ts

@@ -1,13 +1,25 @@
 import _ from "lodash";
 import { Document } from "../Node";
-import { DEFAULT_FS } from "../storage/constants";
 import { CompleteFileSystem, walk } from "../storage/FileSystem";
-import { BaseReader } from "./base";
+import { DEFAULT_FS } from "../storage/constants";
 import { PapaCSVReader } from "./CSVReader";
 import { DocxReader } from "./DocxReader";
 import { HTMLReader } from "./HTMLReader";
 import { MarkdownReader } from "./MarkdownReader";
 import { PDFReader } from "./PDFReader";
+import { BaseReader } from "./base";
+
+type ReaderCallback = (
+  category: "file" | "directory",
+  name: string,
+  status: ReaderStatus,
+  message?: string,
+) => boolean;
+enum ReaderStatus {
+  STARTED = 0,
+  COMPLETE,
+  ERROR,
+}
 
 /**
  * Read a .txt file
@@ -22,7 +34,7 @@ export class TextFileReader implements BaseReader {
   }
 }
 
-const FILE_EXT_TO_READER: Record<string, BaseReader> = {
+export const FILE_EXT_TO_READER: Record<string, BaseReader> = {
   txt: new TextFileReader(),
   pdf: new PDFReader(),
   csv: new PapaCSVReader(),
@@ -40,20 +52,37 @@ export type SimpleDirectoryReaderLoadDataProps = {
 };
 
 /**
- * Read all of the documents in a directory. Currently supports PDF and TXT files.
+ * Read all of the documents in a directory.
+ * By default, supports the list of file types
+ * in the FILE_EXIT_TO_READER map.
  */
 export class SimpleDirectoryReader implements BaseReader {
+  constructor(private observer?: ReaderCallback) {}
+
   async loadData({
     directoryPath,
     fs = DEFAULT_FS as CompleteFileSystem,
     defaultReader = new TextFileReader(),
     fileExtToReader = FILE_EXT_TO_READER,
   }: SimpleDirectoryReaderLoadDataProps): Promise<Document[]> {
+    // Observer can decide to skip the directory
+    if (
+      !this.doObserverCheck("directory", directoryPath, ReaderStatus.STARTED)
+    ) {
+      return [];
+    }
+
     let docs: Document[] = [];
     for await (const filePath of walk(fs, directoryPath)) {
       try {
         const fileExt = _.last(filePath.split(".")) || "";
 
+        // Observer can decide to skip each file
+        if (!this.doObserverCheck("file", filePath, ReaderStatus.STARTED)) {
+          // Skip this file
+          continue;
+        }
+
         let reader = null;
 
         if (fileExt in fileExtToReader) {
@@ -61,16 +90,52 @@ export class SimpleDirectoryReader implements BaseReader {
         } else if (!_.isNil(defaultReader)) {
           reader = defaultReader;
         } else {
-          console.warn(`No reader for file extension of ${filePath}`);
+          const msg = `No reader for file extension of ${filePath}`;
+          console.warn(msg);
+
+          // In an error condition, observer's false cancels the whole process.
+          if (
+            !this.doObserverCheck("file", filePath, ReaderStatus.ERROR, msg)
+          ) {
+            return [];
+          }
+
           continue;
         }
 
         const fileDocs = await reader.loadData(filePath, fs);
-        docs.push(...fileDocs);
+
+        // Observer can still cancel addition of the resulting docs from this file
+        if (this.doObserverCheck("file", filePath, ReaderStatus.COMPLETE)) {
+          docs.push(...fileDocs);
+        }
       } catch (e) {
-        console.error(`Error reading file ${filePath}: ${e}`);
+        const msg = `Error reading file ${filePath}: ${e}`;
+        console.error(msg);
+
+        // In an error condition, observer's false cancels the whole process.
+        if (!this.doObserverCheck("file", filePath, ReaderStatus.ERROR, msg)) {
+          return [];
+        }
       }
     }
+
+    // After successful import of all files, directory completion
+    // is only a notification for observer, cannot be cancelled.
+    this.doObserverCheck("directory", directoryPath, ReaderStatus.COMPLETE);
+
     return docs;
   }
+
+  private doObserverCheck(
+    category: "file" | "directory",
+    name: string,
+    status: ReaderStatus,
+    message?: string,
+  ): boolean {
+    if (this.observer) {
+      return this.observer(category, name, status, message);
+    }
+    return true;
+  }
 }

packages/core/src/tests/SimpleMongoReader.test.ts

@@ -1,48 +0,0 @@
-import { MongoClient } from "mongodb";
-import { SimpleMongoReader } from "../readers/SimpleMongoReader";
-
-describe("SimpleMongoReader", () => {
-  let mockClient: MongoClient;
-  let reader: SimpleMongoReader;
-
-  beforeEach(() => {
-    mockClient = new MongoClient("");
-    reader = new SimpleMongoReader(mockClient);
-  });
-
-  it("should construct correctly", () => {
-    expect(reader).toBeInstanceOf(SimpleMongoReader);
-  });
-
-  it("should load data correctly", async () => {
-    const mockDbName = "testDb";
-    const mockCollectionName = "testCollection";
-    const mockMaxDocs = 10;
-    const mockQueryDict = { key: "value" };
-    const mockQueryOptions = { sort: { key: 1 } };
-    const mockProjection = { key: 1 };
-
-    const mockCursor = {
-      toArray: jest.fn().mockResolvedValue([{ key: "value" }]),
-      limit: jest.fn().mockReturnThis(),
-      project: jest.fn().mockReturnThis(),
-    };
-
-    mockClient.db = jest.fn().mockReturnValue({
-      collection: jest.fn().mockReturnValue({
-        find: jest.fn().mockReturnValue(mockCursor),
-      }),
-    });
-
-    const result = await reader.loadData(
-      mockDbName,
-      mockCollectionName,
-      mockMaxDocs,
-      mockQueryDict,
-      mockQueryOptions,
-      mockProjection
-    );
-
-    expect(result).toEqual([{ text: '{"key":"value"}' }]);
-  });
-});

packages/core/src/tests/TextSplitter.test.ts

@@ -73,6 +73,7 @@ describe("SentenceSplitter", () => {
     let splits = sentenceSplitter.splitText(
       "This is a sentence. This is another sentence. 1.0",
     );
+
     expect(splits).toEqual([
       "This is a sentence.",
       "This is another sentence.",

packages/core/tsconfig.json

@@ -3,7 +3,7 @@
     "esModuleInterop": true,
     "forceConsistentCasingInFileNames": true,
     "isolatedModules": true,
-    "module": "ES6",
+    "module": "esnext",
     "moduleResolution": "node",
     "preserveWatchOutput": true,
     "skipLibCheck": true,