changeset

Add observer/callback feature to SimpleDirectoryReader

.changeset/big-hairs-eat.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Streaming improvements including Anthropic (thanks @kkang2097)

.changeset/breezy-days-admire.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Portkey integration (Thank you @noble-varghese)

.changeset/good-badgers-train.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Add HTMLReader (thanks @mtutty)

.changeset/honest-students-care.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Add export for PromptHelper (thanks @zigamall)

.changeset/shiny-lions-clap.md

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

.changeset/six-hats-compare.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Publish ESM module again

.changeset/twelve-otters-wash.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Pinecone demo (thanks @Einsenhorn)

.gitignore

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

README.md

@@ -84,6 +84,26 @@ Check out our NextJS playground at https://llama-playground.vercel.app/. The sou
 
 - [SimplePrompt](/packages/core/src/Prompt.ts): A simple standardized function call definition that takes in inputs and formats them in a template literal. SimplePrompts can be specialized using currying and combined using other SimplePrompt functions.
 
+## Note: NextJS:
+
+If you're using NextJS App Router, you'll need to use the NodeJS runtime (default) and add the follow config to your next.config.js to have it use imports/exports in the same way Node does.
+
+```js
+export const runtime = "nodejs" // default
+```
+
+```js
+// next.config.js
+/** @type {import('next').NextConfig} */
+const nextConfig = {
+  experimental: {
+    serverComponentsExternalPackages: ["pdf-parse"], // Puts pdf-parse in actual NodeJS mode with NextJS App Router
+  },
+};
+
+module.exports = nextConfig;
+```
+
 ## Supported LLMs:
 
 - OpenAI GPT-3.5-turbo and GPT-4

apps/docs/docs/environments.md

@@ -0,0 +1,29 @@
+---
+sidebar_position: 5
+---
+
+# Environments
+
+LlamaIndex currently officially supports NodeJS 18 and NodeJS 20.
+
+## NextJS App Router
+
+If you're using NextJS App Router route handlers/serverless functions, you'll need to use the NodeJS mode:
+
+```js
+export const runtime = "nodejs" // default
+```
+
+and you'll need to add an exception for pdf-parse in your next.config.js
+
+```js
+// next.config.js
+/** @type {import('next').NextConfig} */
+const nextConfig = {
+  experimental: {
+    serverComponentsExternalPackages: ["pdf-parse"], // Puts pdf-parse in actual NodeJS mode with NextJS App Router
+  },
+};
+
+module.exports = nextConfig;
+```

apps/simple/CHANGELOG.md

@@ -1,5 +1,25 @@
 # simple
 
+## 0.0.29
+
+### Patch Changes
+
+- Updated dependencies [6c55b2d]
+- Updated dependencies [8aa8c65]
+- Updated dependencies [6c55b2d]
+  - llamaindex@0.0.31
+
+## 0.0.28
+
+### Patch Changes
+
+- Updated dependencies [139abad]
+- Updated dependencies [139abad]
+- Updated dependencies [eb0e994]
+- Updated dependencies [eb0e994]
+- Updated dependencies [139abad]
+  - llamaindex@0.0.30
+
 ## 0.0.27
 
 ### Patch Changes

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);

apps/simple/html.ts

@@ -0,0 +1,21 @@
+import { HTMLReader, VectorStoreIndex } from "llamaindex";
+
+async function main() {
+  // Load page
+  const reader = new HTMLReader();
+  const documents = await reader.loadData("data/18-1_Changelog.html");
+
+  // Split text and create embeddings. Store them in a VectorStoreIndex
+  const index = await VectorStoreIndex.fromDocuments(documents);
+
+  // Query the index
+  const queryEngine = index.asQueryEngine();
+  const response = await queryEngine.query(
+    "What were the notable changes in 18.1?",
+  );
+
+  // Output response
+  console.log(response.toString());
+}
+
+main().catch(console.error);

apps/simple/package.json

@@ -1,15 +1,16 @@
 {
-  "version": "0.0.27",
+  "version": "0.0.29",
   "private": true,
   "name": "simple",
   "dependencies": {
-    "@notionhq/client": "^2.2.12",
-    "@pinecone-database/pinecone": "^1.0.1",
-    "commander": "^11.0.0",
+    "@notionhq/client": "^2.2.13",
+    "@pinecone-database/pinecone": "^1.1.2",
+    "commander": "^11.1.0",
     "llamaindex": "workspace:*"
   },
   "devDependencies": {
-    "@types/node": "^18.17.12"
+    "@types/node": "^18.18.6",
+    "ts-node": "^10.9.1"
   },
   "scripts": {
     "lint": "eslint ."

apps/simple/vectorIndexCustomize.ts

@@ -3,6 +3,7 @@ import {
   OpenAI,
   RetrieverQueryEngine,
   serviceContextFromDefaults,
+  SimilarityPostprocessor,
   VectorStoreIndex,
 } from "llamaindex";
 import essay from "./essay";
@@ -21,8 +22,16 @@ async function main() {
 
   const retriever = index.asRetriever();
   retriever.similarityTopK = 5;
+  const nodePostprocessor = new SimilarityPostprocessor({
+    similarityCutoff: 0.7,
+  });
   // TODO: cannot pass responseSynthesizer into retriever query engine
-  const queryEngine = new RetrieverQueryEngine(retriever);
+  const queryEngine = new RetrieverQueryEngine(
+    retriever,
+    undefined,
+    undefined,
+    [nodePostprocessor],
+  );
 
   const response = await queryEngine.query(
     "What did the author do growing up?",

examples/readonly.ts

@@ -0,0 +1,37 @@
+import { execSync } from "child_process";
+import {
+  PDFReader,
+  serviceContextFromDefaults,
+  storageContextFromDefaults,
+  VectorStoreIndex,
+} from "llamaindex";
+
+const STORAGE_DIR = "./cache";
+
+async function main() {
+  // write the index to disk
+  const serviceContext = serviceContextFromDefaults({});
+  const storageContext = await storageContextFromDefaults({
+    persistDir: `${STORAGE_DIR}`,
+  });
+  const reader = new PDFReader();
+  const documents = await reader.loadData("data/brk-2022.pdf");
+  await VectorStoreIndex.fromDocuments(documents, {
+    storageContext,
+    serviceContext,
+  });
+  console.log("wrote index to disk - now trying to read it");
+  // make index dir read only
+  execSync(`chmod -R 555 ${STORAGE_DIR}`);
+  // reopen index
+  const readOnlyStorageContext = await storageContextFromDefaults({
+    persistDir: `${STORAGE_DIR}`,
+  });
+  await VectorStoreIndex.init({
+    storageContext: readOnlyStorageContext,
+    serviceContext,
+  });
+  console.log("read only index successfully opened");
+}
+
+main().catch(console.error);

package.json

@@ -11,16 +11,16 @@
     "publish-snapshot": "turbo run build lint test && changeset version --snapshot && changeset publish"
   },
   "devDependencies": {
-    "@turbo/gen": "^1.10.15",
-    "@types/jest": "^29.5.5",
-    "eslint": "^7.32.0",
+    "@turbo/gen": "^1.10.16",
+    "@types/jest": "^29.5.6",
+    "eslint": "^8.52.0",
     "eslint-config-custom": "workspace:*",
     "husky": "^8.0.3",
     "jest": "^29.7.0",
     "prettier": "^3.0.3",
     "prettier-plugin-organize-imports": "^3.2.3",
     "ts-jest": "^29.1.1",
-    "turbo": "^1.10.15"
+    "turbo": "^1.10.16"
   },
   "packageManager": "pnpm@7.15.0",
   "dependencies": {
@@ -28,7 +28,8 @@
   },
   "pnpm": {
     "overrides": {
-      "trim": "1.0.1"
+      "trim": "1.0.1",
+      "@babel/traverse": "7.23.2"
     }
   }
 }

packages/core/CHANGELOG.md

@@ -1,5 +1,23 @@
 # llamaindex
 
+## 0.0.31
+
+### Patch Changes
+
+- 6c55b2d: Give HistoryChatEngine pluggable options (thanks @marcusschiesser)
+- 8aa8c65: Add SimilarityPostProcessor (thanks @TomPenguin)
+- 6c55b2d: Added LLMMetadata (thanks @marcusschiesser)
+
+## 0.0.30
+
+### Patch Changes
+
+- 139abad: Streaming improvements including Anthropic (thanks @kkang2097)
+- 139abad: Portkey integration (Thank you @noble-varghese)
+- eb0e994: Add export for PromptHelper (thanks @zigamall)
+- eb0e994: Publish ESM module again
+- 139abad: Pinecone demo (thanks @Einsenhorn)
+
 ## 0.0.29
 
 ### Patch Changes

packages/core/package.json

@@ -1,33 +1,35 @@
 {
   "name": "llamaindex",
-  "version": "0.0.29",
+  "version": "0.0.31",
+  "license": "MIT",
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.6.2",
+    "@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",
-    "mongodb": "^6.1.0",
+    "mongodb": "^6.2.0",
     "notion-md-crawler": "^0.0.2",
-    "openai": "^4.11.1",
+    "openai": "^4.14.0",
     "papaparse": "^5.4.1",
     "pdf-parse": "^1.1.1",
-    "portkey-ai": "^0.1.11",
+    "portkey-ai": "^0.1.13",
     "rake-modified": "^1.0.8",
-    "replicate": "^0.20.0",
-    "tiktoken": "^1.0.10",
+    "replicate": "^0.20.1",
+    "string-strip-html": "^13.4.3",
     "uuid": "^9.0.1",
     "wink-nlp": "^1.14.3"
   },
   "devDependencies": {
-    "@types/lodash": "^4.14.199",
-    "@types/node": "^18.18.4",
-    "@types/papaparse": "^5.3.9",
-    "@types/pdf-parse": "^1.1.2",
-    "@types/uuid": "^9.0.5",
+    "@types/lodash": "^4.14.200",
+    "@types/node": "^18.18.7",
+    "@types/papaparse": "^5.3.10",
+    "@types/pdf-parse": "^1.1.3",
+    "@types/uuid": "^9.0.6",
     "node-stdlib-browser": "^1.2.0",
     "tsup": "^7.2.0",
-    "typescript": "^4.9.5"
+    "typescript": "^5.2.2"
   },
   "engines": {
     "node": ">=18.0.0"
@@ -35,9 +37,11 @@
   "types": "./dist/index.d.ts",
   "main": "./dist/index.js",
   "module": "./dist/index.mjs",
+  "repository": "run-llama/LlamaIndexTS",
   "scripts": {
     "lint": "eslint .",
     "test": "jest",
-    "build": "tsup src/index.ts --format esm,cjs --dts"
+    "build": "tsup src/index.ts --format esm,cjs --dts",
+    "dev": "tsup src/index.ts --format esm,cjs --dts --watch"
   }
 }

packages/core/src/ChatEngine.ts

@@ -1,8 +1,9 @@
 import { v4 as uuidv4 } from "uuid";
 import { Event } from "./callbacks/CallbackManager";
-import { ChatHistory, SimpleChatHistory } from "./ChatHistory";
+import { ChatHistory } from "./ChatHistory";
+import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
 import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
-import { TextNode } from "./Node";
+import { NodeWithScore, TextNode } from "./Node";
 import {
   CondenseQuestionPrompt,
   ContextSystemPrompt,
@@ -166,29 +167,89 @@ export class CondenseQuestionChatEngine implements ChatEngine {
   }
 }
 
+export interface Context {
+  message: ChatMessage;
+  nodes: NodeWithScore[];
+}
+
+export interface ContextGenerator {
+  generate(message: string, parentEvent?: Event): Promise<Context>;
+}
+
+export class DefaultContextGenerator implements ContextGenerator {
+  retriever: BaseRetriever;
+  contextSystemPrompt: ContextSystemPrompt;
+  nodePostprocessors: BaseNodePostprocessor[];
+
+  constructor(init: {
+    retriever: BaseRetriever;
+    contextSystemPrompt?: ContextSystemPrompt;
+    nodePostprocessors?: BaseNodePostprocessor[];
+  }) {
+    this.retriever = init.retriever;
+    this.contextSystemPrompt =
+      init?.contextSystemPrompt ?? defaultContextSystemPrompt;
+    this.nodePostprocessors = init.nodePostprocessors || [];
+  }
+
+  private applyNodePostprocessors(nodes: NodeWithScore[]) {
+    return this.nodePostprocessors.reduce(
+      (nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
+      nodes,
+    );
+  }
+
+  async generate(message: string, parentEvent?: Event): Promise<Context> {
+    if (!parentEvent) {
+      parentEvent = {
+        id: uuidv4(),
+        type: "wrapper",
+        tags: ["final"],
+      };
+    }
+    const sourceNodesWithScore = await this.retriever.retrieve(
+      message,
+      parentEvent,
+    );
+
+    const nodes = this.applyNodePostprocessors(sourceNodesWithScore);
+
+    return {
+      message: {
+        content: this.contextSystemPrompt({
+          context: nodes.map((r) => (r.node as TextNode).text).join("\n\n"),
+        }),
+        role: "system",
+      },
+      nodes,
+    };
+  }
+}
+
 /**
  * ContextChatEngine uses the Index to get the appropriate context for each query.
  * The context is stored in the system prompt, and the chat history is preserved,
  * ideally allowing the appropriate context to be surfaced for each query.
  */
 export class ContextChatEngine implements ChatEngine {
-  retriever: BaseRetriever;
   chatModel: LLM;
   chatHistory: ChatMessage[];
-  contextSystemPrompt: ContextSystemPrompt;
+  contextGenerator: ContextGenerator;
 
   constructor(init: {
     retriever: BaseRetriever;
     chatModel?: LLM;
     chatHistory?: ChatMessage[];
     contextSystemPrompt?: ContextSystemPrompt;
+    nodePostprocessors?: BaseNodePostprocessor[];
   }) {
-    this.retriever = init.retriever;
     this.chatModel =
       init.chatModel ?? new OpenAI({ model: "gpt-3.5-turbo-16k" });
     this.chatHistory = init?.chatHistory ?? [];
-    this.contextSystemPrompt =
-      init?.contextSystemPrompt ?? defaultContextSystemPrompt;
+    this.contextGenerator = new DefaultContextGenerator({
+      retriever: init.retriever,
+      contextSystemPrompt: init?.contextSystemPrompt,
+    });
   }
 
   async chat<
@@ -211,24 +272,12 @@ export class ContextChatEngine implements ChatEngine {
       type: "wrapper",
       tags: ["final"],
     };
-    const sourceNodesWithScore = await this.retriever.retrieve(
-      message,
-      parentEvent,
-    );
-
-    const systemMessage: ChatMessage = {
-      content: this.contextSystemPrompt({
-        context: sourceNodesWithScore
-          .map((r) => (r.node as TextNode).text)
-          .join("\n\n"),
-      }),
-      role: "system",
-    };
+    const context = await this.contextGenerator.generate(message, parentEvent);
 
     chatHistory.push({ content: message, role: "user" });
 
     const response = await this.chatModel.chat(
-      [systemMessage, ...chatHistory],
+      [context.message, ...chatHistory],
       parentEvent,
     );
     chatHistory.push(response.message);
@@ -237,7 +286,7 @@ export class ContextChatEngine implements ChatEngine {
 
     return new Response(
       response.message.content,
-      sourceNodesWithScore.map((r) => r.node),
+      context.nodes.map((r) => r.node),
     ) as R;
   }
 
@@ -252,24 +301,12 @@ export class ContextChatEngine implements ChatEngine {
       type: "wrapper",
       tags: ["final"],
     };
-    const sourceNodesWithScore = await this.retriever.retrieve(
-      message,
-      parentEvent,
-    );
-
-    const systemMessage: ChatMessage = {
-      content: this.contextSystemPrompt({
-        context: sourceNodesWithScore
-          .map((r) => (r.node as TextNode).text)
-          .join("\n\n"),
-      }),
-      role: "system",
-    };
+    const context = await this.contextGenerator.generate(message, parentEvent);
 
     chatHistory.push({ content: message, role: "user" });
 
     const response_stream = await this.chatModel.chat(
-      [systemMessage, ...chatHistory],
+      [context.message, ...chatHistory],
       parentEvent,
       true,
     );
@@ -279,7 +316,7 @@ export class ContextChatEngine implements ChatEngine {
       yield part;
     }
 
-    chatHistory.push({ content: accumulator, role: "system" });
+    chatHistory.push({ content: accumulator, role: "assistant" });
 
     this.chatHistory = chatHistory;
 
@@ -292,43 +329,56 @@ export class ContextChatEngine implements ChatEngine {
 }
 
 /**
- * HistoryChatEngine is a ChatEngine that uses a ChatHistory to keep track of the chat history. This is an example with the same behavior as SimpleChatEngine
- * TODO: generally use the ChatHistory instead of ChatMessage[] - breaking change
+ * HistoryChatEngine is a ChatEngine that uses a `ChatHistory` object
+ * to keeps track of chat's message history.
+ * A `ChatHistory` object is passed as a parameter for each call to the `chat` method,
+ * so the state of the chat engine is preserved between calls.
+ * Optionally, a `ContextGenerator` can be used to generate an additional context for each call to `chat`.
  */
-export class HistoryChatEngine implements ChatEngine {
-  chatHistory: ChatHistory;
+export class HistoryChatEngine {
   llm: LLM;
+  contextGenerator?: ContextGenerator;
 
   constructor(init?: Partial<HistoryChatEngine>) {
-    this.chatHistory = init?.chatHistory ?? new SimpleChatHistory();
     this.llm = init?.llm ?? new OpenAI();
+    this.contextGenerator = init?.contextGenerator;
   }
 
   async chat<
     T extends boolean | undefined = undefined,
     R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
-  >(
-    message: string,
-    chatHistory?: ChatMessage[] | undefined,
-    streaming?: T,
-  ): Promise<R> {
+  >(message: string, chatHistory: ChatHistory, streaming?: T): Promise<R> {
     //Streaming option
     if (streaming) {
       return this.streamChat(message, chatHistory) as R;
     }
-    this.chatHistory.addMessage({ content: message, role: "user" });
-    const response = await this.llm.chat(this.chatHistory.requestMessages);
-    this.chatHistory.addMessage(response.message);
+    const context = await this.contextGenerator?.generate(message);
+    chatHistory.addMessage({
+      content: message,
+      role: "user",
+    });
+    const response = await this.llm.chat(
+      await chatHistory.requestMessages(
+        context ? [context.message] : undefined,
+      ),
+    );
+    chatHistory.addMessage(response.message);
     return new Response(response.message.content) as R;
   }
 
   protected async *streamChat(
     message: string,
-    chatHistory?: ChatMessage[] | undefined,
+    chatHistory: ChatHistory,
   ): AsyncGenerator<string, void, unknown> {
-    this.chatHistory.addMessage({ content: message, role: "user" });
+    const context = await this.contextGenerator?.generate(message);
+    chatHistory.addMessage({
+      content: message,
+      role: "user",
+    });
     const response_stream = await this.llm.chat(
-      this.chatHistory.requestMessages,
+      await chatHistory.requestMessages(
+        context ? [context.message] : undefined,
+      ),
       undefined,
       true,
     );
@@ -338,11 +388,10 @@ export class HistoryChatEngine implements ChatEngine {
       accumulator += part;
       yield part;
     }
-    this.chatHistory.addMessage({ content: accumulator, role: "user" });
+    chatHistory.addMessage({
+      content: accumulator,
+      role: "assistant",
+    });
     return;
   }
-
-  reset() {
-    this.chatHistory.reset();
-  }
 }

packages/core/src/ChatHistory.ts

@@ -1,4 +1,4 @@
-import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
+import { ChatMessage, LLM, MessageType, OpenAI } from "./llm/LLM";
 import {
   defaultSummaryPrompt,
   messagesToHistoryStr,
@@ -14,106 +14,187 @@ export interface ChatHistory {
    * Adds a message to the chat history.
    * @param message
    */
-  addMessage(message: ChatMessage): Promise<void>;
+  addMessage(message: ChatMessage): void;
 
   /**
    * Returns the messages that should be used as input to the LLM.
    */
-  requestMessages: ChatMessage[];
+  requestMessages(transientMessages?: ChatMessage[]): Promise<ChatMessage[]>;
 
   /**
    * Resets the chat history so that it's empty.
    */
   reset(): void;
+
+  /**
+   * Returns the new messages since the last call to this function (or since calling the constructor)
+   */
+  newMessages(): ChatMessage[];
 }
 
 export class SimpleChatHistory implements ChatHistory {
   messages: ChatMessage[];
+  private messagesBefore: number;
 
   constructor(init?: Partial<SimpleChatHistory>) {
     this.messages = init?.messages ?? [];
+    this.messagesBefore = this.messages.length;
   }
-  async addMessage(message: ChatMessage) {
+
+  addMessage(message: ChatMessage) {
     this.messages.push(message);
   }
 
-  get requestMessages() {
-    return this.messages;
+  async requestMessages(transientMessages?: ChatMessage[]) {
+    return [...(transientMessages ?? []), ...this.messages];
   }
 
   reset() {
     this.messages = [];
   }
+
+  newMessages() {
+    const newMessages = this.messages.slice(this.messagesBefore);
+    this.messagesBefore = this.messages.length;
+    return newMessages;
+  }
 }
 
 export class SummaryChatHistory implements ChatHistory {
-  messagesToSummarize: number;
+  tokensToSummarize: number;
   messages: ChatMessage[];
   summaryPrompt: SummaryPrompt;
   llm: LLM;
+  private messagesBefore: number;
 
   constructor(init?: Partial<SummaryChatHistory>) {
-    this.messagesToSummarize = init?.messagesToSummarize ?? 5;
     this.messages = init?.messages ?? [];
+    this.messagesBefore = this.messages.length;
     this.summaryPrompt = init?.summaryPrompt ?? defaultSummaryPrompt;
     this.llm = init?.llm ?? new OpenAI();
-  }
-
-  private async summarize() {
-    // get all messages after the last summary message (including)
-    const chatHistoryStr = messagesToHistoryStr(
-      this.messages.slice(this.getLastSummaryIndex()),
-    );
-
-    const response = await this.llm.complete(
-      this.summaryPrompt({ context: chatHistoryStr }),
-    );
-
-    this.messages.push({ content: response.message.content, role: "memory" });
-  }
-
-  async addMessage(message: ChatMessage) {
-    const lastSummaryIndex = this.getLastSummaryIndex();
-    // if there are more than or equal `messagesToSummarize` messages since the last summary, call summarize
-    if (
-      lastSummaryIndex !== -1 &&
-      this.messages.length - lastSummaryIndex - 1 >= this.messagesToSummarize
-    ) {
-      // TODO: define what are better conditions, e.g. depending on the context length of the LLM?
-      // for now we just summarize each `messagesToSummarize` messages
-      await this.summarize();
+    if (!this.llm.metadata.maxTokens) {
+      throw new Error(
+        "LLM maxTokens is not set. Needed so the summarizer ensures the context window size of the LLM.",
+      );
     }
+    this.tokensToSummarize =
+      this.llm.metadata.contextWindow - this.llm.metadata.maxTokens;
+  }
+
+  private async summarize(): Promise<ChatMessage> {
+    // get the conversation messages to create summary
+    const messagesToSummarize = this.calcConversationMessages();
+
+    let promptMessages;
+    do {
+      promptMessages = [
+        {
+          content: this.summaryPrompt({
+            context: messagesToHistoryStr(messagesToSummarize),
+          }),
+          role: "user" as MessageType,
+        },
+      ];
+      // remove oldest message until the chat history is short enough for the context window
+      messagesToSummarize.shift();
+    } while (this.llm.tokens(promptMessages) > this.tokensToSummarize);
+
+    const response = await this.llm.chat(promptMessages);
+    return { content: response.message.content, role: "memory" };
+  }
+
+  addMessage(message: ChatMessage) {
     this.messages.push(message);
   }
 
   // Find last summary message
-  private getLastSummaryIndex() {
-    return this.messages
-      .slice()
-      .reverse()
-      .findIndex((message) => message.role === "memory");
+  private getLastSummaryIndex(): number | null {
+    const reversedMessages = this.messages.slice().reverse();
+    const index = reversedMessages.findIndex(
+      (message) => message.role === "memory",
+    );
+    if (index === -1) {
+      return null;
+    }
+    return this.messages.length - 1 - index;
   }
 
-  get requestMessages() {
-    const lastSummaryIndex = this.getLastSummaryIndex();
+  private get systemMessages() {
     // get array of all system messages
-    const systemMessages = this.messages.filter(
-      (message) => message.role === "system",
-    );
-    // convert summary message so it can be send to the LLM
-    const summaryMessage: ChatMessage = {
-      content: `This is a summary of conversation so far: ${this.messages[lastSummaryIndex].content}`,
-      role: "system",
-    };
-    // return system messages, last summary and all messages after the last summary message
+    return this.messages.filter((message) => message.role === "system");
+  }
+
+  private get nonSystemMessages() {
+    // get array of all non-system messages
+    return this.messages.filter((message) => message.role !== "system");
+  }
+
+  /**
+   * Calculates the messages that describe the conversation so far.
+   * If there's no memory, all non-system messages are used.
+   * If there's a memory, uses all messages after the last summary message.
+   */
+  private calcConversationMessages(transformSummary?: boolean): ChatMessage[] {
+    const lastSummaryIndex = this.getLastSummaryIndex();
+    if (!lastSummaryIndex) {
+      // there's no memory, so just use all non-system messages
+      return this.nonSystemMessages;
+    } else {
+      // there's a memory, so use all messages after the last summary message
+      // and convert summary message so it can be send to the LLM
+      const summaryMessage: ChatMessage = transformSummary
+        ? {
+            content: `Summary of the conversation so far: ${this.messages[lastSummaryIndex].content}`,
+            role: "system",
+          }
+        : this.messages[lastSummaryIndex];
+      return [summaryMessage, ...this.messages.slice(lastSummaryIndex + 1)];
+    }
+  }
+
+  private calcCurrentRequestMessages(transientMessages?: ChatMessage[]) {
+    // TODO: check order: currently, we're sending:
+    // system messages first, then transient messages and then the messages that describe the conversation so far
     return [
-      ...systemMessages,
-      summaryMessage,
-      ...this.messages.slice(lastSummaryIndex + 1),
+      ...this.systemMessages,
+      ...(transientMessages ? transientMessages : []),
+      ...this.calcConversationMessages(true),
     ];
   }
 
+  async requestMessages(transientMessages?: ChatMessage[]) {
+    const requestMessages = this.calcCurrentRequestMessages(transientMessages);
+
+    // get tokens of current request messages and the transient messages
+    const tokens = this.llm.tokens(requestMessages);
+    if (tokens > this.tokensToSummarize) {
+      // if there are too many tokens for the next request, call summarize
+      const memoryMessage = await this.summarize();
+      const lastMessage = this.messages.at(-1);
+      if (lastMessage && lastMessage.role === "user") {
+        // if last message is a user message, ensure that it's sent after the new memory message
+        this.messages.pop();
+        this.messages.push(memoryMessage);
+        this.messages.push(lastMessage);
+      } else {
+        // otherwise just add the memory message
+        this.messages.push(memoryMessage);
+      }
+      // TODO: we still might have too many tokens
+      // e.g. too large system messages or transient messages
+      // how should we deal with that?
+      return this.calcCurrentRequestMessages(transientMessages);
+    }
+    return requestMessages;
+  }
+
   reset() {
     this.messages = [];
   }
+
+  newMessages() {
+    const newMessages = this.messages.slice(this.messagesBefore);
+    this.messagesBefore = this.messages.length;
+    return newMessages;
+  }
 }

packages/core/src/GlobalsHelper.ts

@@ -1,9 +1,12 @@
-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";
 
+export enum Tokenizers {
+  CL100K_BASE = "cl100k_base",
+}
+
 /**
  * Helper class singleton
  */
@@ -14,35 +17,40 @@ 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);      
       },
     };
   }
 
-  tokenizer() {
+  tokenizer(encoding?: string) {
+    if (encoding && encoding !== Tokenizers.CL100K_BASE) {
+      throw new Error(`Tokenizer encoding ${encoding} not yet supported`);
+    }
     if (!this.defaultTokenizer) {
       this.initDefaultTokenizer();
     }
-
+  
     return this.defaultTokenizer!.encode.bind(this.defaultTokenizer);
   }
-
-  tokenizerDecoder() {
+  
+  tokenizerDecoder(encoding?: string) {
+    if (encoding && encoding !== Tokenizers.CL100K_BASE) {
+      throw new Error(`Tokenizer encoding ${encoding} not yet supported`);
+    }
     if (!this.defaultTokenizer) {
       this.initDefaultTokenizer();
     }
-
+  
     return this.defaultTokenizer!.decode.bind(this.defaultTokenizer);
   }
 

packages/core/src/QueryEngine.ts

@@ -1,5 +1,6 @@
 import { v4 as uuidv4 } from "uuid";
 import { Event } from "./callbacks/CallbackManager";
+import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
 import { NodeWithScore, TextNode } from "./Node";
 import {
   BaseQuestionGenerator,
@@ -30,12 +31,14 @@ export interface BaseQueryEngine {
 export class RetrieverQueryEngine implements BaseQueryEngine {
   retriever: BaseRetriever;
   responseSynthesizer: ResponseSynthesizer;
+  nodePostprocessors: BaseNodePostprocessor[];
   preFilters?: unknown;
 
   constructor(
     retriever: BaseRetriever,
     responseSynthesizer?: ResponseSynthesizer,
     preFilters?: unknown,
+    nodePostprocessors?: BaseNodePostprocessor[],
   ) {
     this.retriever = retriever;
     const serviceContext: ServiceContext | undefined =
@@ -43,6 +46,24 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
     this.responseSynthesizer =
       responseSynthesizer || new ResponseSynthesizer({ serviceContext });
     this.preFilters = preFilters;
+    this.nodePostprocessors = nodePostprocessors || [];
+  }
+
+  private applyNodePostprocessors(nodes: NodeWithScore[]) {
+    return this.nodePostprocessors.reduce(
+      (nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
+      nodes,
+    );
+  }
+
+  private async retrieve(query: string, parentEvent: Event) {
+    const nodes = await this.retriever.retrieve(
+      query,
+      parentEvent,
+      this.preFilters,
+    );
+
+    return this.applyNodePostprocessors(nodes);
   }
 
   async query(query: string, parentEvent?: Event) {
@@ -51,11 +72,7 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
       type: "wrapper",
       tags: ["final"],
     };
-    const nodes = await this.retriever.retrieve(
-      query,
-      _parentEvent,
-      this.preFilters,
-    );
+    const nodes = await this.retrieve(query, _parentEvent);
     return this.responseSynthesizer.synthesize(query, nodes, _parentEvent);
   }
 }

packages/core/src/index.ts

@@ -1,5 +1,6 @@
 export * from "./callbacks/CallbackManager";
 export * from "./ChatEngine";
+export * from "./ChatHistory";
 export * from "./constants";
 export * from "./Embedding";
 export * from "./GlobalsHelper";
@@ -17,6 +18,7 @@ export * from "./readers/CSVReader";
 export * from "./readers/MarkdownReader";
 export * from "./readers/NotionReader";
 export * from "./readers/PDFReader";
+export * from "./readers/HTMLReader";
 export * from "./readers/SimpleDirectoryReader";
 export * from "./Response";
 export * from "./ResponseSynthesizer";

packages/core/src/indices/BaseNodePostprocessor.ts

@@ -0,0 +1,20 @@
+import { NodeWithScore } from "../Node";
+
+export interface BaseNodePostprocessor {
+  postprocessNodes: (nodes: NodeWithScore[]) => NodeWithScore[];
+}
+
+export class SimilarityPostprocessor implements BaseNodePostprocessor {
+  similarityCutoff?: number;
+
+  constructor(options?: { similarityCutoff?: number }) {
+    this.similarityCutoff = options?.similarityCutoff;
+  }
+
+  postprocessNodes(nodes: NodeWithScore[]) {
+    if (this.similarityCutoff === undefined) return nodes;
+
+    const cutoff = this.similarityCutoff || 0;
+    return nodes.filter((node) => node.score && node.score >= cutoff);
+  }
+}

packages/core/src/indices/index.ts

@@ -1,4 +1,5 @@
 export * from "./BaseIndex";
+export * from "./BaseNodePostprocessor";
 export * from "./keyword";
 export * from "./summary";
 export * from "./vectorStore";

packages/core/src/indices/keyword/KeywordTableIndex.ts

@@ -15,6 +15,7 @@ import {
   IndexStructType,
   KeywordTable,
 } from "../BaseIndex";
+import { BaseNodePostprocessor } from "../BaseNodePostprocessor";
 import {
   KeywordTableLLMRetriever,
   KeywordTableRAKERetriever,
@@ -129,11 +130,15 @@ export class KeywordTableIndex extends BaseIndex<KeywordTable> {
   asQueryEngine(options?: {
     retriever?: BaseRetriever;
     responseSynthesizer?: ResponseSynthesizer;
+    preFilters?: unknown;
+    nodePostprocessors?: BaseNodePostprocessor[];
   }): BaseQueryEngine {
     const { retriever, responseSynthesizer } = options ?? {};
     return new RetrieverQueryEngine(
       retriever ?? this.asRetriever(),
       responseSynthesizer,
+      options?.preFilters,
+      options?.nodePostprocessors,
     );
   }
 

packages/core/src/indices/summary/SummaryIndex.ts

@@ -10,17 +10,18 @@ import {
   ServiceContext,
   serviceContextFromDefaults,
 } from "../../ServiceContext";
+import { BaseDocumentStore, RefDocInfo } from "../../storage/docStore/types";
 import {
   StorageContext,
   storageContextFromDefaults,
 } from "../../storage/StorageContext";
-import { BaseDocumentStore, RefDocInfo } from "../../storage/docStore/types";
 import {
   BaseIndex,
   BaseIndexInit,
   IndexList,
   IndexStructType,
 } from "../BaseIndex";
+import { BaseNodePostprocessor } from "../BaseNodePostprocessor";
 import {
   SummaryIndexLLMRetriever,
   SummaryIndexRetriever,
@@ -155,6 +156,8 @@ export class SummaryIndex extends BaseIndex<IndexList> {
   asQueryEngine(options?: {
     retriever?: BaseRetriever;
     responseSynthesizer?: ResponseSynthesizer;
+    preFilters?: unknown;
+    nodePostprocessors?: BaseNodePostprocessor[];
   }): BaseQueryEngine {
     let { retriever, responseSynthesizer } = options ?? {};
 
@@ -170,7 +173,12 @@ export class SummaryIndex extends BaseIndex<IndexList> {
       });
     }
 
-    return new RetrieverQueryEngine(retriever, responseSynthesizer);
+    return new RetrieverQueryEngine(
+      retriever,
+      responseSynthesizer,
+      options?.preFilters,
+      options?.nodePostprocessors,
+    );
   }
 
   static async buildIndexFromNodes(

packages/core/src/indices/vectorStore/VectorStoreIndex.ts

@@ -18,6 +18,7 @@ import {
   IndexDict,
   IndexStructType,
 } from "../BaseIndex";
+import { BaseNodePostprocessor } from "../BaseNodePostprocessor";
 import { VectorIndexRetriever } from "./VectorIndexRetriever";
 
 export interface VectorIndexOptions {
@@ -87,24 +88,23 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
       );
     }
 
-    if (!indexStruct && !options.nodes) {
+    if (options.nodes) {
+      // If nodes are passed in, then we need to update the index
+      indexStruct = await VectorStoreIndex.buildIndexFromNodes(
+        options.nodes,
+        serviceContext,
+        vectorStore,
+        docStore,
+        indexStruct,
+      );
+
+      await indexStore.addIndexStruct(indexStruct);
+    } else if (!indexStruct) {
       throw new Error(
         "Cannot initialize VectorStoreIndex without nodes or indexStruct",
       );
     }
 
-    const nodes = options.nodes ?? [];
-
-    indexStruct = await VectorStoreIndex.buildIndexFromNodes(
-      nodes,
-      serviceContext,
-      vectorStore,
-      docStore,
-      indexStruct,
-    );
-
-    await indexStore.addIndexStruct(indexStruct);
-
     return new VectorStoreIndex({
       storageContext,
       serviceContext,
@@ -247,11 +247,15 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
   asQueryEngine(options?: {
     retriever?: BaseRetriever;
     responseSynthesizer?: ResponseSynthesizer;
+    preFilters?: unknown;
+    nodePostprocessors?: BaseNodePostprocessor[];
   }): BaseQueryEngine {
     const { retriever, responseSynthesizer } = options ?? {};
     return new RetrieverQueryEngine(
       retriever ?? this.asRetriever(),
       responseSynthesizer,
+      options?.preFilters,
+      options?.nodePostprocessors,
     );
   }
 

packages/core/src/llm/LLM.ts

@@ -9,6 +9,7 @@ import {
 } from "../callbacks/CallbackManager";
 
 import { LLMOptions } from "portkey-ai";
+import { globalsHelper, Tokenizers } from "../GlobalsHelper";
 import {
   AnthropicSession,
   ANTHROPIC_AI_PROMPT,
@@ -48,10 +49,20 @@ export interface ChatResponse {
 // NOTE in case we need CompletionResponse to diverge from ChatResponse in the future
 export type CompletionResponse = ChatResponse;
 
+export interface LLMMetadata {
+  model: string;
+  temperature: number;
+  topP: number;
+  maxTokens?: number;
+  contextWindow: number;
+  tokenizer: Tokenizers | undefined;
+}
+
 /**
  * Unified language model interface
  */
 export interface LLM {
+  metadata: LLMMetadata;
   // Whether a LLM has streaming support
   hasStreaming: boolean;
   /**
@@ -81,6 +92,11 @@ export interface LLM {
     parentEvent?: Event,
     streaming?: T,
   ): Promise<R>;
+
+  /**
+   * Calculates the number of tokens needed for the given chat messages
+   */
+  tokens(messages: ChatMessage[]): number;
 }
 
 export const GPT4_MODELS = {
@@ -183,6 +199,32 @@ export class OpenAI implements LLM {
     this.callbackManager = init?.callbackManager;
   }
 
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_OPENAI_MODELS[this.model].contextWindow,
+      tokenizer: Tokenizers.CL100K_BASE,
+    };
+  }
+
+  tokens(messages: ChatMessage[]): number {
+    // for latest OpenAI models, see https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+    const tokenizer = globalsHelper.tokenizer(this.metadata.tokenizer);
+    const tokensPerMessage = 3;
+    let numTokens = 0;
+    for (const message of messages) {
+      numTokens += tokensPerMessage;
+      for (const value of Object.values(message)) {
+        numTokens += tokenizer(value).length;
+      }
+    }
+    numTokens += 3; // every reply is primed with <|im_start|>assistant<|im_sep|>
+    return numTokens;
+  }
+
   mapMessageType(
     messageType: MessageType,
   ): "user" | "assistant" | "system" | "function" {
@@ -393,6 +435,21 @@ export class LlamaDeuce implements LLM {
     this.hasStreaming = init?.hasStreaming ?? false;
   }
 
+  tokens(messages: ChatMessage[]): number {
+    throw new Error("Method not implemented.");
+  }
+
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow,
+      tokenizer: undefined,
+    };
+  }
+
   mapMessagesToPrompt(messages: ChatMessage[]) {
     if (this.chatStrategy === DeuceChatStrategy.A16Z) {
       return this.mapMessagesToPromptA16Z(messages);
@@ -545,6 +602,12 @@ If a question does not make any sense, or is not factually coherent, explain why
   }
 }
 
+export const ALL_AVAILABLE_ANTHROPIC_MODELS = {
+  // both models have 100k context window, see https://docs.anthropic.com/claude/reference/selecting-a-model
+  "claude-2": { contextWindow: 100000 },
+  "claude-instant-1": { contextWindow: 100000 },
+};
+
 /**
  * Anthropic LLM implementation
  */
@@ -553,7 +616,7 @@ export class Anthropic implements LLM {
   hasStreaming: boolean = true;
 
   // Per completion Anthropic params
-  model: string;
+  model: keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS;
   temperature: number;
   topP: number;
   maxTokens?: number;
@@ -585,6 +648,21 @@ export class Anthropic implements LLM {
 
     this.callbackManager = init?.callbackManager;
   }
+  
+  tokens(messages: ChatMessage[]): number {
+    throw new Error("Method not implemented.");
+  }
+
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_ANTHROPIC_MODELS[this.model].contextWindow,
+      tokenizer: undefined,
+    };
+  }
 
   mapMessagesToPrompt(messages: ChatMessage[]) {
     return (
@@ -707,6 +785,14 @@ export class Portkey implements LLM {
     this.callbackManager = init?.callbackManager;
   }
 
+  tokens(messages: ChatMessage[]): number {
+    throw new Error("Method not implemented.");
+  }
+
+  get metadata(): LLMMetadata {
+    throw new Error("metadata not implemented for Portkey");
+  }
+
   async chat<
     T extends boolean | undefined = undefined,
     R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,

packages/core/src/readers/HTMLReader.ts

@@ -0,0 +1,77 @@
+import { Document } from "../Node";
+import { DEFAULT_FS } from "../storage/constants";
+import { GenericFileSystem } from "../storage/FileSystem";
+import { BaseReader } from "./base";
+
+/**
+ * Extract the significant text from an arbitrary HTML document.
+ * The contents of any head, script, style, and xml tags are removed completely.
+ * The URLs for a[href] tags are extracted, along with the inner text of the tag.
+ * All other tags are removed, and the inner text is kept intact.
+ * Html entities (e.g., &) are not decoded.
+ */
+export class HTMLReader implements BaseReader {
+  /**
+   * Public method for this reader.
+   * Required by BaseReader interface.
+   * @param file Path/name of the file to be loaded.
+   * @param fs fs wrapper interface for getting the file content.
+   * @returns Promise<Document[]> A Promise object, eventually yielding zero or one Document parsed from the HTML content of the specified file.
+   */
+  async loadData(
+    file: string,
+    fs: GenericFileSystem = DEFAULT_FS,
+  ): Promise<Document[]> {
+    const dataBuffer = await fs.readFile(file, "utf-8");
+    const htmlOptions = this.getOptions();
+    const content = await this.parseContent(dataBuffer, htmlOptions);
+    return [new Document({ text: content, id_: file })];
+  }
+
+  /**
+   * Wrapper for string-strip-html usage.
+   * @param html Raw HTML content to be parsed.
+   * @param options An object of options for the underlying library
+   * @see getOptions
+   * @returns The HTML content, stripped of unwanted tags and attributes
+   */
+  async parseContent(html: string, options: any = {}): Promise<string> {
+    const { stripHtml } = await import("string-strip-html"); // ESM only
+    return stripHtml(html).result;
+  }
+
+  /**
+   * Wrapper for our configuration options passed to string-strip-html library
+   * @see https://codsen.com/os/string-strip-html/examples
+   * @returns An object of options for the underlying library
+   */
+  getOptions() {
+    return {
+      skipHtmlDecoding: true,
+      stripTogetherWithTheirContents: [
+        "script", // default
+        "style", // default
+        "xml", // default
+        "head", // <-- custom-added
+      ],
+      // Keep the URLs for embedded links
+      // cb: (tag: any, deleteFrom: number, deleteTo: number, insert: string, rangesArr: any, proposedReturn: string) => {
+      //   let temp;
+      //   if (
+      //     tag.name === "a" &&
+      //     tag.attributes &&
+      //     tag.attributes.some((attr: any) => {
+      //       if (attr.name === "href") {
+      //         temp = attr.value;
+      //         return true;
+      //       }
+      //     })
+      //   ) {
+      //     rangesArr.push([deleteFrom, deleteTo, `${temp} ${insert || ""}`]);
+      //   } else {
+      //     rangesArr.push(proposedReturn);
+      //   }
+      // },
+    };
+  }
+}

packages/core/src/readers/SimpleDirectoryReader.ts

@@ -1,12 +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
@@ -21,12 +34,14 @@ 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(),
   md: new MarkdownReader(),
   docx: new DocxReader(),
+  htm: new HTMLReader(),
+  html: new HTMLReader(),
 };
 
 export type SimpleDirectoryReaderLoadDataProps = {
@@ -37,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) {
@@ -58,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/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,6 +3,7 @@
     "esModuleInterop": true,
     "forceConsistentCasingInFileNames": true,
     "isolatedModules": true,
+    "module": "esnext",
     "moduleResolution": "node",
     "preserveWatchOutput": true,
     "skipLibCheck": true,