Merge main into sweep/add-unit-tests

Add HTMLReader, sample app and HTML file

.changeset/good-badgers-train.md

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

apps/simple/CHANGELOG.md

@@ -1,5 +1,14 @@
 # simple
 
+## 0.0.29
+
+### Patch Changes
+
+- Updated dependencies [6c55b2d]
+- Updated dependencies [8aa8c65]
+- Updated dependencies [6c55b2d]
+  - llamaindex@0.0.31
+
 ## 0.0.28
 
 ### Patch Changes

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.28",
+  "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,13 @@
 # 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

packages/core/package.json

@@ -1,33 +1,35 @@
 {
   "name": "llamaindex",
-  "version": "0.0.30",
+  "version": "0.0.31",
+  "license": "MIT",
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.6.2",
+    "@anthropic-ai/sdk": "^0.8.1",
     "@notionhq/client": "^2.2.13",
     "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.13",
     "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"
   },
   "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,10 +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",
     "dev": "tsup src/index.ts --format esm,cjs --dts --watch"
   }
-}
+}

packages/core/src/ChatEngine.ts

@@ -1,6 +1,7 @@
 import { v4 as uuidv4 } from "uuid";
 import { Event } from "./callbacks/CallbackManager";
 import { ChatHistory } from "./ChatHistory";
+import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
 import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
 import { NodeWithScore, TextNode } from "./Node";
 import {
@@ -178,14 +179,24 @@ export interface ContextGenerator {
 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> {
@@ -201,16 +212,16 @@ export class DefaultContextGenerator implements ContextGenerator {
       parentEvent,
     );
 
+    const nodes = this.applyNodePostprocessors(sourceNodesWithScore);
+
     return {
       message: {
         content: this.contextSystemPrompt({
-          context: sourceNodesWithScore
-            .map((r) => (r.node as TextNode).text)
-            .join("\n\n"),
+          context: nodes.map((r) => (r.node as TextNode).text).join("\n\n"),
         }),
         role: "system",
       },
-      nodes: sourceNodesWithScore,
+      nodes,
     };
   }
 }
@@ -230,6 +241,7 @@ export class ContextChatEngine implements ChatEngine {
     chatModel?: LLM;
     chatHistory?: ChatMessage[];
     contextSystemPrompt?: ContextSystemPrompt;
+    nodePostprocessors?: BaseNodePostprocessor[];
   }) {
     this.chatModel =
       init.chatModel ?? new OpenAI({ model: "gpt-3.5-turbo-16k" });

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

@@ -18,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/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

@@ -5,6 +5,7 @@ import { CompleteFileSystem, walk } from "../storage/FileSystem";
 import { BaseReader } from "./base";
 import { PapaCSVReader } from "./CSVReader";
 import { DocxReader } from "./DocxReader";
+import { HTMLReader } from "./HTMLReader";
 import { MarkdownReader } from "./MarkdownReader";
 import { PDFReader } from "./PDFReader";
 
@@ -27,6 +28,8 @@ const FILE_EXT_TO_READER: Record<string, BaseReader> = {
   csv: new PapaCSVReader(),
   md: new MarkdownReader(),
   docx: new DocxReader(),
+  htm: new HTMLReader(),
+  html: new HTMLReader(),
 };
 
 export type SimpleDirectoryReaderLoadDataProps = {

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

@@ -32,13 +32,37 @@ describe("similarity", () => {
   });
 
   test("calculates euclidean similarity", () => {
-    const queryEmbedding = [1, 0];
-    const docEmbedding1 = [0, 1]; // farther from query, distance 1.414
-    const docEmbedding2 = [1, 1]; // closer to query distance 1
-    expect(
-      similarity(queryEmbedding, docEmbedding1, SimilarityType.EUCLIDEAN),
-    ).toBeLessThan(
-      similarity(queryEmbedding, docEmbedding2, SimilarityType.EUCLIDEAN),
-    );
+      const queryEmbedding = [1, 0];
+      const docEmbedding1 = [0, 1]; // farther from query, distance 1.414
+      const docEmbedding2 = [1, 1]; // closer to query distance 1
+      expect(
+        similarity(queryEmbedding, docEmbedding1, SimilarityType.EUCLIDEAN),
+      ).toBeLessThan(
+        similarity(queryEmbedding, docEmbedding2, SimilarityType.EUCLIDEAN),
+      );
+    });
+  
+    test("calculates cosine similarity with zero vectors", () => {
+      const embedding1 = [0, 0];
+      const embedding2 = [0, 0];
+      expect(similarity(embedding1, embedding2, SimilarityType.DEFAULT)).toEqual(
+        NaN,
+      );
+    });
+  
+    test("calculates dot product with zero vectors", () => {
+      const embedding1 = [0, 0];
+      const embedding2 = [0, 0];
+      expect(
+        similarity(embedding1, embedding2, SimilarityType.DOT_PRODUCT),
+      ).toEqual(0);
+    });
+  
+    test("calculates euclidean similarity with zero vectors", () => {
+      const embedding1 = [0, 0];
+      const embedding2 = [0, 0];
+      expect(
+        similarity(embedding1, embedding2, SimilarityType.EUCLIDEAN),
+      ).toEqual(0);
+    });
   });
-});

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

@@ -0,0 +1,34 @@
+import { Node, ObjectType, MetadataMode } from '../Node';
+import { v4 as uuidv4 } from 'uuid';
+
+describe('Node class', () => {
+  let node: Node;
+
+  beforeEach(() => {
+    node = new Node();
+  });
+
+  test('getType method', () => {
+    expect(node.getType()).toBe(ObjectType.TEXT);
+  });
+
+  test('getContent method', () => {
+    const content = 'Test content';
+    node.setContent(content);
+    expect(node.getContent(MetadataMode.ALL)).toBe(content);
+  });
+
+  test('getMetadataStr method', () => {
+    const metadata = { key: 'value' };
+    node.metadata = metadata;
+    expect(node.getMetadataStr(MetadataMode.ALL)).toBe(JSON.stringify(metadata));
+  });
+
+  test('setContent method', () => {
+    const content = 'Test content';
+    node.setContent(content);
+    expect(node.getContent(MetadataMode.ALL)).toBe(content);
+  });
+
+  // Mock LLM completions and test other methods...
+});

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

@@ -0,0 +1,34 @@
+import { RouterQueryEngine } from '../RouterQueryEngine';
+import { LLM } from '../LLM'; // Assuming LLM is a class or module that can be mocked
+
+jest.mock('../LLM'); // Mock the LLM completions
+
+describe('RouterQueryEngine', () => {
+  let routerQueryEngine: RouterQueryEngine;
+  let mockLLM: jest.Mocked<LLM>;
+
+  beforeEach(() => {
+    mockLLM = new LLM() as jest.Mocked<LLM>;
+    routerQueryEngine = new RouterQueryEngine(mockLLM);
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  test('method1', () => {
+    // Arrange
+    const input = 'some input';
+    const expectedOutput = 'some output';
+    mockLLM.completion.mockReturnValue(expectedOutput);
+
+    // Act
+    const output = routerQueryEngine.method1(input);
+
+    // Assert
+    expect(output).toBe(expectedOutput);
+    expect(mockLLM.completion).toHaveBeenCalledWith(input);
+  });
+
+  // Repeat the test block for other methods in the RouterQueryEngine class
+});

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

@@ -0,0 +1,21 @@
+import { Selector } from '../Selector';
+import { LLM } from '../LLM'; // Assuming LLM is a class or function that can be mocked
+import { jest } from '@jest/globals';
+
+describe('Selector', () => {
+  let mockLLM: jest.Mocked<LLM>;
+
+  beforeEach(() => {
+    mockLLM = {
+      complete: jest.fn(),
+      // Add other methods to be mocked here
+    };
+  });
+
+  test('method1', () => {
+    const selector = new Selector(mockLLM);
+    // Call method1 and assert the output
+  });
+
+  // Add more tests for other methods here
+});

packages/core/tsconfig.json

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