e2e install pack and run the package

Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/smart-eels-design.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: update `VectorIndexRetriever` constructor parameters' type.

.github/workflows/e2e.yml

@@ -50,7 +50,7 @@ jobs:
         run: pnpm run build
         working-directory: ./packages/create-llama
       - name: Run Playwright tests
-        run: pnpm exec playwright test
+        run: pnpm run e2e
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
         working-directory: ./packages/create-llama

.github/workflows/test.yml

@@ -32,10 +32,35 @@ jobs:
       - name: Install dependencies
         run: pnpm install
       - name: Build
-        run: pnpm run build
-        working-directory: ./packages/core
+        run: pnpm run build --filter llamaindex
       - name: Run Type Check
         run: pnpm run type-check
       - name: Run Circular Dependency Check
         run: pnpm run circular-check
         working-directory: ./packages/core
+  typecheck-examples:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: pnpm/action-setup@v2
+      - name: Setup Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version-file: ".nvmrc"
+          cache: "pnpm"
+      - name: Install dependencies
+        run: pnpm install
+      - name: Build
+        run: pnpm run build --filter llamaindex
+      - name: Copy examples
+        run: rsync -rv --exclude=node_modules ./examples ${{ runner.temp }}
+      - name: Pack
+        run: pnpm pack --pack-destination ${{ runner.temp }}
+        working-directory: packages/core
+      - name: Install llamaindex
+        run: npm add ${{ runner.temp }}/*.tgz
+        working-directory: ${{ runner.temp }}/examples
+      - name: Run Type Check
+        run: npx tsc --project ./tsconfig.json
+        working-directory: ${{ runner.temp }}/examples

.vscode/settings.json

@@ -8,5 +8,11 @@
   "jest.rootPath": "./packages/core",
   "[python]": {
     "editor.defaultFormatter": "ms-python.black-formatter"
+  },
+  "[jsonc]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "[json]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
   }
 }

README.md

@@ -70,7 +70,7 @@ main();
 Then you can run it using
 
 ```bash
-pnpx ts-node example.ts
+pnpm dlx ts-node example.ts
 ```
 
 ## Playground
@@ -105,6 +105,9 @@ export const runtime = "nodejs"; // default
 // next.config.js
 /** @type {import('next').NextConfig} */
 const nextConfig = {
+  experimental: {
+    serverComponentsExternalPackages: ["pdf2json"],
+  },
   webpack: (config) => {
     config.resolve.alias = {
       ...config.resolve.alias,

apps/docs/CHANGELOG.md

@@ -1,5 +1,11 @@
 # docs
 
+## 0.0.2
+
+### Patch Changes
+
+- 0f64084: docs: update API references
+
 ## 0.0.1
 
 ### Patch Changes

apps/docs/docs/examples/agent.mdx

@@ -0,0 +1,85 @@
+# Agents
+
+A built-in agent that can take decisions and reasoning based on the tools provided to it.
+
+## OpenAI Agent
+
+```ts
+import { FunctionTool, OpenAIAgent } from "llamaindex";
+
+// Define a function to sum two numbers
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+  return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+  return a / b;
+}
+
+// Define the parameters of the sum function as a JSON schema
+const sumJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The first number",
+    },
+    b: {
+      type: "number",
+      description: "The second number",
+    },
+  },
+  required: ["a", "b"],
+};
+
+// Define the parameters of the divide function as a JSON schema
+const divideJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The dividend to divide",
+    },
+    b: {
+      type: "number",
+      description: "The divisor to divide by",
+    },
+  },
+  required: ["a", "b"],
+};
+
+async function main() {
+  // Create a function tool from the sum function
+  const sumFunctionTool = new FunctionTool(sumNumbers, {
+    name: "sumNumbers",
+    description: "Use this function to sum two numbers",
+    parameters: sumJSON,
+  });
+
+  // Create a function tool from the divide function
+  const divideFunctionTool = new FunctionTool(divideNumbers, {
+    name: "divideNumbers",
+    description: "Use this function to divide two numbers"
+    parameters: divideJSON,
+  });
+
+  // Create an OpenAIAgent with the function tools
+  const agent = new OpenAIAgent({
+    tools: [sumFunctionTool, divideFunctionTool],
+    verbose: true,
+  });
+
+  // Chat with the agent
+  const response = await agent.chat({
+    message: "How much is 5 + 5? then divide by 2",
+  });
+
+  // Print the response
+  console.log(String(response));
+}
+
+main().then(() => {
+  console.log("Done");
+});
+```

apps/docs/docs/getting_started/concepts.md

@@ -35,7 +35,7 @@ LlamaIndex.TS help you prepare the knowledge base with a suite of data connector
 [**Data Loaders**](../modules/data_loader.md):
 A data connector (i.e. `Reader`) ingest data from different data sources and data formats into a simple `Document` representation (text and simple metadata).
 
-[**Documents / Nodes**](../modules/documents_and_nodes.md): A `Document` is a generic container around any data source - for instance, a PDF, an API output, or retrieved data from a database. A `Node` is the atomic unit of data in LlamaIndex and represents a "chunk" of a source `Document`. It's a rich representation that includes metadata and relationships (to other nodes) to enable accurate and expressive retrieval operations.
+[**Documents / Nodes**](../modules/documents_and_nodes/index.md): A `Document` is a generic container around any data source - for instance, a PDF, an API output, or retrieved data from a database. A `Node` is the atomic unit of data in LlamaIndex and represents a "chunk" of a source `Document`. It's a rich representation that includes metadata and relationships (to other nodes) to enable accurate and expressive retrieval operations.
 
 [**Data Indexes**](../modules/data_index.md):
 Once you've ingested your data, LlamaIndex helps you index data into a format that's easy to retrieve.
@@ -69,7 +69,7 @@ A response synthesizer generates a response from an LLM, using a user query and
 
 #### Pipelines
 
-[**Query Engines**](../modules/query_engine.md):
+[**Query Engines**](../modules/query_engines):
 A query engine is an end-to-end pipeline that allow you to ask question over your data.
 It takes in a natural language query, and returns a response, along with reference context retrieved and passed to the LLM.
 

apps/docs/docs/getting_started/installation.mdx

@@ -58,6 +58,6 @@ Our examples use OpenAI by default. You'll need to set up your Open AI key like
 export OPENAI_API_KEY="sk-......" # Replace with your key from https://platform.openai.com/account/api-keys
 ```
 
-If you want to have it automatically loaded every time, add it to your .zshrc/.bashrc.
+If you want to have it automatically loaded every time, add it to your `.zshrc/.bashrc`.
 
 WARNING: do not check in your OpenAI key into version control.

apps/docs/docs/getting_started/starter.md

@@ -36,9 +36,9 @@ async function main() {
 
   // Query the index
   const queryEngine = index.asQueryEngine();
-  const response = await queryEngine.query(
-    "What did the author do in college?",
-  );
+  const response = await queryEngine.query({
+    query: "What did the author do in college?",
+  });
 
   // Output response
   console.log(response.toString());

apps/docs/docs/introduction.md

@@ -37,9 +37,9 @@ For more complex applications, our lower-level APIs allow advanced users to cust
 
 `npm install llamaindex`
 
-Our documentation includes [Installation Instructions](./installation.mdx) and a [Starter Tutorial](./starter.md) to build your first application.
+Our documentation includes [Installation Instructions](./getting_started/installation.mdx) and a [Starter Tutorial](./getting_started/starter.md) to build your first application.
 
-Once you're up and running, [High-Level Concepts](./getting_started/concepts.md) has an overview of LlamaIndex's modular architecture. For more hands-on practical examples, look through our [End-to-End Tutorials](./end_to_end.md).
+Once you're up and running, [High-Level Concepts](./getting_started/concepts.md) has an overview of LlamaIndex's modular architecture. For more hands-on practical examples, look through our Examples section on the sidebar.
 
 ## 🗺️ Ecosystem
 

apps/docs/docs/modules/agent/_category_.yml

@@ -0,0 +1 @@
+label: "Agents"

apps/docs/docs/modules/agent/index.md

@@ -0,0 +1,14 @@
+# Agents
+
+An “agent” is an automated reasoning and decision engine. It takes in a user input/query and can make internal decisions for executing that query in order to return the correct result. The key agent components can include, but are not limited to:
+
+- Breaking down a complex question into smaller ones
+- Choosing an external Tool to use + coming up with parameters for calling the Tool
+- Planning out a set of tasks
+- Storing previously completed tasks in a memory module
+
+## Getting Started
+
+LlamaIndex.TS comes with a few built-in agents, but you can also create your own. The built-in agents include:
+
+- [OpenAI Agent](./openai.mdx)

apps/docs/docs/modules/agent/openai.mdx

@@ -0,0 +1,183 @@
+# OpenAI Agent
+
+OpenAI API that supports function calling, it’s never been easier to build your own agent!
+
+In this notebook tutorial, we showcase how to write your own OpenAI agent
+
+## Setup
+
+First, you need to install the `llamaindex` package. You can do this by running the following command in your terminal:
+
+```bash
+pnpm i llamaindex
+```
+
+Then we can define a function to sum two numbers and another function to divide two numbers.
+
+```ts
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+  return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+  return a / b;
+}
+```
+
+## Create a function tool
+
+Now we can create a function tool from the sum function and another function tool from the divide function.
+
+For the parameters of the sum function, we can define a JSON schema.
+
+### JSON Schema
+
+```ts
+const sumJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The first number",
+    },
+    b: {
+      type: "number",
+      description: "The second number",
+    },
+  },
+  required: ["a", "b"],
+};
+
+const divideJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The dividend a to divide",
+    },
+    b: {
+      type: "number",
+      description: "The divisor b to divide by",
+    },
+  },
+  required: ["a", "b"],
+};
+
+const sumFunctionTool = new FunctionTool(sumNumbers, {
+  name: "sumNumbers",
+  description: "Use this function to sum two numbers",
+  parameters: sumJSON,
+});
+
+const divideFunctionTool = new FunctionTool(divideNumbers, {
+  name: "divideNumbers",
+  description: "Use this function to divide two numbers",
+  parameters: divideJSON,
+});
+```
+
+## Create an OpenAIAgent
+
+Now we can create an OpenAIAgent with the function tools.
+
+```ts
+const worker = new OpenAIAgent({
+  tools: [sumFunctionTool, divideFunctionTool],
+  verbose: true,
+});
+```
+
+## Chat with the agent
+
+Now we can chat with the agent.
+
+```ts
+const response = await worker.chat({
+  message: "How much is 5 + 5? then divide by 2",
+});
+
+console.log(String(response));
+```
+
+## Full code
+
+```ts
+import { FunctionTool, OpenAIAgent } from "llamaindex";
+
+// Define a function to sum two numbers
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+  return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+  return a / b;
+}
+
+// Define the parameters of the sum function as a JSON schema
+const sumJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The first number",
+    },
+    b: {
+      type: "number",
+      description: "The second number",
+    },
+  },
+  required: ["a", "b"],
+};
+
+// Define the parameters of the divide function as a JSON schema
+const divideJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The argument a to divide",
+    },
+    b: {
+      type: "number",
+      description: "The argument b to divide",
+    },
+  },
+  required: ["a", "b"],
+};
+
+async function main() {
+  // Create a function tool from the sum function
+  const sumFunctionTool = new FunctionTool(sumNumbers, {
+    name: "sumNumbers",
+    description: "Use this function to sum two numbers",
+    parameters: sumJSON,
+  });
+
+  // Create a function tool from the divide function
+  const divideFunctionTool = new FunctionTool(divideNumbers, {
+    name: "divideNumbers",
+    description: "Use this function to divide two numbers",
+    parameters: divideJSON,
+  });
+
+  // Create an OpenAIAgent with the function tools
+  const agent = new OpenAIAgent({
+    tools: [sumFunctionTool, divideFunctionTool],
+    verbose: true,
+  });
+
+  // Chat with the agent
+  const response = await agent.chat({
+    message: "How much is 5 + 5? then divide by 2",
+  });
+
+  // Print the response
+  console.log(String(response));
+}
+
+main().then(() => {
+  console.log("Done");
+});
+```

apps/docs/docs/modules/chat_engine.md

@@ -25,5 +25,5 @@ for await (const chunk of stream) {
 
 ## Api References
 
-- [ContextChatEngine](../../api/classes/ContextChatEngine.md)
-- [CondenseQuestionChatEngine](../../api/classes/ContextChatEngine.md)
+- [ContextChatEngine](../api/classes/ContextChatEngine.md)
+- [CondenseQuestionChatEngine](../api/classes/ContextChatEngine.md)

apps/docs/docs/modules/data_index.md

@@ -19,5 +19,5 @@ const index = await VectorStoreIndex.fromDocuments([document]);
 
 ## API Reference
 
-- [SummaryIndex](../../api/classes/SummaryIndex.md)
-- [VectorStoreIndex](../../api/classes/VectorStoreIndex.md)
+- [SummaryIndex](../api/classes/SummaryIndex.md)
+- [VectorStoreIndex](../api/classes/VectorStoreIndex.md)

apps/docs/docs/modules/data_loader.md

@@ -14,4 +14,4 @@ documents = new SimpleDirectoryReader().loadData("./data");
 
 ## API Reference
 
-- [SimpleDirectoryReader](../../api/classes/SimpleDirectoryReader.md)
+- [SimpleDirectoryReader](../api/classes/SimpleDirectoryReader.md)

apps/docs/docs/modules/documents_and_nodes/index.md

@@ -14,5 +14,5 @@ document = new Document({ text: "text", metadata: { key: "val" } });
 
 ## API Reference
 
-- [Document](../../api/classes/Document.md)
-- [TextNode](../../api/classes/TextNode.md)
+- [Document](../api/classes/Document.md)
+- [TextNode](../api/classes/TextNode.md)

apps/docs/docs/modules/embedding.md

@@ -18,5 +18,5 @@ const serviceContext = serviceContextFromDefaults({ embedModel: openaiEmbeds });
 
 ## API Reference
 
-- [OpenAIEmbedding](../../api/classes/OpenAIEmbedding.md)
-- [ServiceContext](../../api/interfaces/ServiceContext.md)
+- [OpenAIEmbedding](../api/classes/OpenAIEmbedding.md)
+- [ServiceContext](../api/interfaces//ServiceContext.md)

apps/docs/docs/modules/ingestion_pipeline/transformations.md

@@ -4,7 +4,7 @@ A transformation is something that takes a list of nodes as an input, and return
 
 Currently, the following components are Transformation objects:
 
-- [SimpleNodeParser](../../api/classes/SimpleNodeParser.md)
+- [SimpleNodeParser](../api/classes/SimpleNodeParser.md)
 - [MetadataExtractor](../documents_and_nodes/metadata_extraction.md)
 - Embeddings
 

apps/docs/docs/modules/llm.md

@@ -16,7 +16,19 @@ const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0 });
 const serviceContext = serviceContextFromDefaults({ llm: openaiLLM });
 ```
 
+## Azure OpenAI
+
+To use Azure OpenAI, you only need to set a few environment variables.
+
+For example:
+
+```
+export AZURE_OPENAI_KEY="<YOUR KEY HERE>"
+export AZURE_OPENAI_ENDPOINT="<YOUR ENDPOINT, see https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart?tabs=command-line%2Cpython&pivots=rest-api>"
+export AZURE_OPENAI_DEPLOYMENT="gpt-4" # or some other deployment name
+```
+
 ## API Reference
 
-- [OpenAI](../../api/classes/OpenAI.md)
-- [ServiceContext](../../api/interfaces/ServiceContext.md)
+- [OpenAI](../api/classes/OpenAI.md)
+- [ServiceContext](../api/interfaces//ServiceContext.md)

apps/docs/docs/modules/node_parser.md

@@ -4,7 +4,7 @@ sidebar_position: 3
 
 # NodeParser
 
-The `NodeParser` in LlamaIndex is responbile for splitting `Document` objects into more manageable `Node` objects. When you call `.fromDocuments()`, the `NodeParser` from the `ServiceContext` is used to do this automatically for you. Alternatively, you can use it to split documents ahead of time.
+The `NodeParser` in LlamaIndex is responsible for splitting `Document` objects into more manageable `Node` objects. When you call `.fromDocuments()`, the `NodeParser` from the `ServiceContext` is used to do this automatically for you. Alternatively, you can use it to split documents ahead of time.
 
 ```typescript
 import { Document, SimpleNodeParser } from "llamaindex";
@@ -29,5 +29,5 @@ const textSplits = splitter.splitText("Hello World");
 
 ## API Reference
 
-- [SimpleNodeParser](../../api/classes/SimpleNodeParser.md)
-- [SentenceSplitter](../../api/classes/SentenceSplitter.md)
+- [SimpleNodeParser](../api/classes/SimpleNodeParser.md)
+- [SentenceSplitter](../api/classes/SentenceSplitter.md)

apps/docs/docs/modules/query_engines/_category_.yml

@@ -0,0 +1,2 @@
+label: "Query Engines"
+position: 2

apps/docs/docs/modules/query_engines/index.md

@@ -1,7 +1,3 @@
----
-sidebar_position: 4
----
-
 # QueryEngine
 
 A query engine wraps a `Retriever` and a `ResponseSynthesizer` into a pipeline, that will use the query string to fetech nodes and then send them to the LLM to generate a response.

apps/docs/docs/modules/query_engines/metadata_filtering.md

@@ -0,0 +1,152 @@
+# Metadata Filtering
+
+Metadata filtering is a way to filter the documents that are returned by a query based on the metadata associated with the documents. This is useful when you want to filter the documents based on some metadata that is not part of the document text.
+
+You can also check our multi-tenancy blog post to see how metadata filtering can be used in a multi-tenant environment. [https://blog.llamaindex.ai/building-multi-tenancy-rag-system-with-llamaindex-0d6ab4e0c44b] (the article uses the Python version of LlamaIndex, but the concepts are the same).
+
+## Setup
+
+Firstly if you haven't already, you need to install the `llamaindex` package:
+
+```bash
+pnpm i llamaindex
+```
+
+Then you can import the necessary modules from `llamaindex`:
+
+```ts
+import {
+  ChromaVectorStore,
+  Document,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
+
+const collectionName = "dog_colors";
+```
+
+## Creating documents with metadata
+
+You can create documents with metadata using the `Document` class:
+
+```ts
+const docs = [
+  new Document({
+    text: "The dog is brown",
+    metadata: {
+      color: "brown",
+      dogId: "1",
+    },
+  }),
+  new Document({
+    text: "The dog is red",
+    metadata: {
+      color: "red",
+      dogId: "2",
+    },
+  }),
+];
+```
+
+## Creating a ChromaDB vector store
+
+You can create a `ChromaVectorStore` to store the documents:
+
+```ts
+const chromaVS = new ChromaVectorStore({ collectionName });
+const serviceContext = await storageContextFromDefaults({
+  vectorStore: chromaVS,
+});
+
+const index = await VectorStoreIndex.fromDocuments(docs, {
+  storageContext: serviceContext,
+});
+```
+
+## Querying the index with metadata filtering
+
+Now you can query the index with metadata filtering using the `preFilters` option:
+
+```ts
+const queryEngine = index.asQueryEngine({
+  preFilters: {
+    filters: [
+      {
+        key: "dogId",
+        value: "2",
+        filterType: "ExactMatch",
+      },
+    ],
+  },
+});
+
+const response = await queryEngine.query({
+  query: "What is the color of the dog?",
+});
+
+console.log(response.toString());
+```
+
+## Full Code
+
+```ts
+import {
+  ChromaVectorStore,
+  Document,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
+
+const collectionName = "dog_colors";
+
+async function main() {
+  try {
+    const docs = [
+      new Document({
+        text: "The dog is brown",
+        metadata: {
+          color: "brown",
+          dogId: "1",
+        },
+      }),
+      new Document({
+        text: "The dog is red",
+        metadata: {
+          color: "red",
+          dogId: "2",
+        },
+      }),
+    ];
+
+    console.log("Creating ChromaDB vector store");
+    const chromaVS = new ChromaVectorStore({ collectionName });
+    const ctx = await storageContextFromDefaults({ vectorStore: chromaVS });
+
+    console.log("Embedding documents and adding to index");
+    const index = await VectorStoreIndex.fromDocuments(docs, {
+      storageContext: ctx,
+    });
+
+    console.log("Querying index");
+    const queryEngine = index.asQueryEngine({
+      preFilters: {
+        filters: [
+          {
+            key: "dogId",
+            value: "2",
+            filterType: "ExactMatch",
+          },
+        ],
+      },
+    });
+    const response = await queryEngine.query({
+      query: "What is the color of the dog?",
+    });
+    console.log(response.toString());
+  } catch (e) {
+    console.error(e);
+  }
+}
+
+main();
+```

apps/docs/docs/modules/query_engines/router_query_engine.md

@@ -0,0 +1,189 @@
+# Router Query Engine
+
+In this tutorial, we define a custom router query engine that selects one out of several candidate query engines to execute a query.
+
+## Setup
+
+First, we need to install import the necessary modules from `llamaindex`:
+
+```bash
+pnpm i lamaindex
+```
+
+```ts
+import {
+  OpenAI,
+  RouterQueryEngine,
+  SimpleDirectoryReader,
+  SimpleNodeParser,
+  SummaryIndex,
+  VectorStoreIndex,
+  serviceContextFromDefaults,
+} from "llamaindex";
+```
+
+## Loading Data
+
+Next, we need to load some data. We will use the `SimpleDirectoryReader` to load documents from a directory:
+
+```ts
+const documents = await new SimpleDirectoryReader().loadData({
+  directoryPath: "node_modules/llamaindex/examples",
+});
+```
+
+## Service Context
+
+Next, we need to define some basic rules and parse the documents into nodes. We will use the `SimpleNodeParser` to parse the documents into nodes and `ServiceContext` to define the rules (eg. LLM API key, chunk size, etc.):
+
+```ts
+const nodeParser = new SimpleNodeParser({
+  chunkSize: 1024,
+});
+
+const serviceContext = serviceContextFromDefaults({
+  nodeParser,
+  llm: new OpenAI(),
+});
+```
+
+## Creating Indices
+
+Next, we need to create some indices. We will create a `VectorStoreIndex` and a `SummaryIndex`:
+
+```ts
+const vectorIndex = await VectorStoreIndex.fromDocuments(documents, {
+  serviceContext,
+});
+
+const summaryIndex = await SummaryIndex.fromDocuments(documents, {
+  serviceContext,
+});
+```
+
+## Creating Query Engines
+
+Next, we need to create some query engines. We will create a `VectorStoreQueryEngine` and a `SummaryQueryEngine`:
+
+```ts
+const vectorQueryEngine = vectorIndex.asQueryEngine();
+const summaryQueryEngine = summaryIndex.asQueryEngine();
+```
+
+## Creating a Router Query Engine
+
+Next, we need to create a router query engine. We will use the `RouterQueryEngine` to create a router query engine:
+
+We're defining two query engines, one for summarization and one for retrieving specific context. The router query engine will select the most appropriate query engine based on the query.
+
+```ts
+const queryEngine = RouterQueryEngine.fromDefaults({
+  queryEngineTools: [
+    {
+      queryEngine: vectorQueryEngine,
+      description: "Useful for summarization questions related to Abramov",
+    },
+    {
+      queryEngine: summaryQueryEngine,
+      description: "Useful for retrieving specific context from Abramov",
+    },
+  ],
+  serviceContext,
+});
+```
+
+## Querying the Router Query Engine
+
+Finally, we can query the router query engine:
+
+```ts
+const summaryResponse = await queryEngine.query({
+  query: "Give me a summary about his past experiences?",
+});
+
+console.log({
+  answer: summaryResponse.response,
+  metadata: summaryResponse?.metadata?.selectorResult,
+});
+```
+
+## Full code
+
+```ts
+import {
+  OpenAI,
+  RouterQueryEngine,
+  SimpleDirectoryReader,
+  SimpleNodeParser,
+  SummaryIndex,
+  VectorStoreIndex,
+  serviceContextFromDefaults,
+} from "llamaindex";
+
+async function main() {
+  // Load documents from a directory
+  const documents = await new SimpleDirectoryReader().loadData({
+    directoryPath: "node_modules/llamaindex/examples",
+  });
+
+  // Parse the documents into nodes
+  const nodeParser = new SimpleNodeParser({
+    chunkSize: 1024,
+  });
+
+  // Create a service context
+  const serviceContext = serviceContextFromDefaults({
+    nodeParser,
+    llm: new OpenAI(),
+  });
+
+  // Create indices
+  const vectorIndex = await VectorStoreIndex.fromDocuments(documents, {
+    serviceContext,
+  });
+
+  const summaryIndex = await SummaryIndex.fromDocuments(documents, {
+    serviceContext,
+  });
+
+  // Create query engines
+  const vectorQueryEngine = vectorIndex.asQueryEngine();
+  const summaryQueryEngine = summaryIndex.asQueryEngine();
+
+  // Create a router query engine
+  const queryEngine = RouterQueryEngine.fromDefaults({
+    queryEngineTools: [
+      {
+        queryEngine: vectorQueryEngine,
+        description: "Useful for summarization questions related to Abramov",
+      },
+      {
+        queryEngine: summaryQueryEngine,
+        description: "Useful for retrieving specific context from Abramov",
+      },
+    ],
+    serviceContext,
+  });
+
+  // Query the router query engine
+  const summaryResponse = await queryEngine.query({
+    query: "Give me a summary about his past experiences?",
+  });
+
+  console.log({
+    answer: summaryResponse.response,
+    metadata: summaryResponse?.metadata?.selectorResult,
+  });
+
+  const specificResponse = await queryEngine.query({
+    query: "Tell me about abramov first job?",
+  });
+
+  console.log({
+    answer: specificResponse.response,
+    metadata: specificResponse.metadata.selectorResult,
+  });
+}
+
+main().then(() => console.log("Done"));
+```

apps/docs/docs/modules/response_synthesizer.md

@@ -57,8 +57,8 @@ for await (const chunk of stream) {
 
 ## API Reference
 
-- [ResponseSynthesizer](../../api/classes/ResponseSynthesizer.md)
-- [Refine](../../api/classes/Refine.md)
-- [CompactAndRefine](../../api/classes/CompactAndRefine.md)
-- [TreeSummarize](../../api/classes/TreeSummarize.md)
-- [SimpleResponseBuilder](../../api/classes/SimpleResponseBuilder.md)
+- [ResponseSynthesizer](../api/classes/ResponseSynthesizer.md)
+- [Refine](../api/classes/Refine.md)
+- [CompactAndRefine](../api/classes/CompactAndRefine.md)
+- [TreeSummarize](../api/classes/TreeSummarize.md)
+- [SimpleResponseBuilder](../api/classes/SimpleResponseBuilder.md)

apps/docs/docs/modules/retriever.md

@@ -16,6 +16,6 @@ const nodesWithScore = await retriever.retrieve("query string");
 
 ## API Reference
 
-- [SummaryIndexRetriever](../../api/classes/SummaryIndexRetriever.md)
-- [SummaryIndexLLMRetriever](../../api/classes/SummaryIndexLLMRetriever.md)
-- [VectorIndexRetriever](../../api/classes/VectorIndexRetriever.md)
+- [SummaryIndexRetriever](../api/classes/SummaryIndexRetriever.md)
+- [SummaryIndexLLMRetriever](../api/classes/SummaryIndexLLMRetriever.md)
+- [VectorIndexRetriever](../api/classes/VectorIndexRetriever.md)

apps/docs/docs/modules/storage.md

@@ -23,4 +23,4 @@ const index = await VectorStoreIndex.fromDocuments([document], {
 
 ## API Reference
 
-- [StorageContext](../../api/interfaces/StorageContext.md)
+- [StorageContext](../api/interfaces//StorageContext.md)

apps/docs/package.json

@@ -1,6 +1,6 @@
 {
   "name": "docs",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "private": true,
   "scripts": {
     "docusaurus": "docusaurus",

examples/agent/openai.ts

@@ -0,0 +1,76 @@
+import { FunctionTool, OpenAIAgent } from "llamaindex";
+
+// Define a function to sum two numbers
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+  return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+  return a / b;
+}
+
+// Define the parameters of the sum function as a JSON schema
+const sumJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The first number",
+    },
+    b: {
+      type: "number",
+      description: "The second number",
+    },
+  },
+  required: ["a", "b"],
+};
+
+const divideJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The dividend a to divide",
+    },
+    b: {
+      type: "number",
+      description: "The divisor b to divide by",
+    },
+  },
+  required: ["a", "b"],
+};
+
+async function main() {
+  // Create a function tool from the sum function
+  const functionTool = new FunctionTool(sumNumbers, {
+    name: "sumNumbers",
+    description: "Use this function to sum two numbers",
+    parameters: sumJSON,
+  });
+
+  // Create a function tool from the divide function
+  const functionTool2 = new FunctionTool(divideNumbers, {
+    name: "divideNumbers",
+    description: "Use this function to divide two numbers",
+    parameters: divideJSON,
+  });
+
+  // Create an OpenAIAgent with the function tools
+  const agent = new OpenAIAgent({
+    tools: [functionTool, functionTool2],
+    verbose: true,
+  });
+
+  // Chat with the agent
+  const response = await agent.chat({
+    message: "How much is 5 + 5? then divide by 2",
+  });
+
+  // Print the response
+  console.log(String(response));
+}
+
+main().then(() => {
+  console.log("Done");
+});

examples/chatEngine.ts

@@ -1,6 +1,4 @@
 import { stdin as input, stdout as output } from "node:process";
-// readline/promises is still experimental so not in @types/node yet
-// @ts-ignore
 import readline from "node:readline/promises";
 
 import {

examples/chromadb/README.md

@@ -6,7 +6,7 @@ Export your OpenAI API Key using `export OPEN_API_KEY=insert your api key here`
 
 If you haven't installed chromadb, run `pip install chromadb`. Start the server using `chroma run`.
 
-Now, open a new terminal window and inside `examples`, run `pnpx ts-node chromadb/test.ts`.
+Now, open a new terminal window and inside `examples`, run `pnpm dlx ts-node chromadb/test.ts`.
 
 Here's the output for the input query `Tell me about Godfrey Cheshire's rating of La Sapienza.`:
 

examples/gptllama.ts

@@ -1,6 +1,4 @@
 import { stdin as input, stdout as output } from "node:process";
-// readline/promises is still experimental so not in @types/node yet
-// @ts-ignore
 import readline from "node:readline/promises";
 
 import { ChatMessage, LlamaDeuce, OpenAI } from "llamaindex";

examples/package.json

@@ -14,7 +14,8 @@
   },
   "devDependencies": {
     "@types/node": "^18.19.10",
-    "ts-node": "^10.9.2"
+    "ts-node": "^10.9.2",
+    "typescript": "^5.3.3"
   },
   "scripts": {
     "lint": "eslint ."

examples/pipeline/ingestion.ts

@@ -19,7 +19,6 @@ async function main() {
   const pipeline = new IngestionPipeline({
     transformations: [
       new SimpleNodeParser({ chunkSize: 1024, chunkOverlap: 20 }),
-      // new TitleExtractor(llm),
       new OpenAIEmbedding(),
     ],
   });

examples/readers/load-assemblyai.ts

@@ -1,12 +1,7 @@
 import { program } from "commander";
-import {
-  AudioTranscriptReader,
-  TranscribeParams,
-  VectorStoreIndex,
-} from "llamaindex";
+import { TranscribeParams, VectorStoreIndex } from "llamaindex";
+import { AudioTranscriptReader } from "llamaindex/readers/AssemblyAIReader";
 import { stdin as input, stdout as output } from "node:process";
-// readline/promises is still experimental so not in @types/node yet
-// @ts-ignore
 import readline from "node:readline/promises";
 
 program

examples/readers/load-csv.ts

@@ -1,11 +1,11 @@
 import {
   CompactAndRefine,
   OpenAI,
-  PapaCSVReader,
   ResponseSynthesizer,
   serviceContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
+import { PapaCSVReader } from "llamaindex/readers/CSVReader";
 
 async function main() {
   // Load CSV

examples/readers/load-docx.ts

@@ -1,4 +1,5 @@
-import { DocxReader, VectorStoreIndex } from "llamaindex";
+import { VectorStoreIndex } from "llamaindex";
+import { DocxReader } from "llamaindex/readers/DocxReader";
 
 const FILE_PATH = "./data/stars.docx";
 const SAMPLE_QUERY = "Information about Zodiac";

examples/readers/load-html.ts

@@ -1,4 +1,5 @@
-import { HTMLReader, VectorStoreIndex } from "llamaindex";
+import { VectorStoreIndex } from "llamaindex";
+import { HTMLReader } from "llamaindex/readers/HTMLReader";
 
 async function main() {
   // Load page

examples/readers/load-md.ts

@@ -1,4 +1,5 @@
-import { MarkdownReader, VectorStoreIndex } from "llamaindex";
+import { VectorStoreIndex } from "llamaindex";
+import { MarkdownReader } from "llamaindex/readers/MarkdownReader";
 
 const FILE_PATH = "./data/planets.md";
 const SAMPLE_QUERY = "List all planets";

examples/readers/load-notion.ts

@@ -1,9 +1,8 @@
 import { Client } from "@notionhq/client";
 import { program } from "commander";
-import { NotionReader, VectorStoreIndex } from "llamaindex";
+import { VectorStoreIndex } from "llamaindex";
+import { NotionReader } from "llamaindex/readers/NotionReader";
 import { stdin as input, stdout as output } from "node:process";
-// readline/promises is still experimental so not in @types/node yet
-// @ts-ignore
 import readline from "node:readline/promises";
 
 program

examples/readers/load-pdf.ts

@@ -1,9 +1,13 @@
-import { PDFReader, VectorStoreIndex } from "llamaindex";
+import { VectorStoreIndex } from "llamaindex";
+import { PDFReader } from "llamaindex/readers/PDFReader";
+import { resolve } from "node:path";
 
 async function main() {
   // Load PDF
   const reader = new PDFReader();
-  const documents = await reader.loadData("data/brk-2022.pdf");
+  const documents = await reader.loadData(
+    resolve(__dirname, "../data/brk-2022.pdf"),
+  );
 
   // Split text and create embeddings. Store them in a VectorStoreIndex
   const index = await VectorStoreIndex.fromDocuments(documents);
@@ -11,7 +15,7 @@ async function main() {
   // Query the index
   const queryEngine = index.asQueryEngine();
   const response = await queryEngine.query({
-    query: "What mistakes did they make?",
+    query: "What mistakes did Warren E. Buffett make?",
   });
 
   // Output response

examples/routerQueryEngine.ts

@@ -0,0 +1,76 @@
+import {
+  OpenAI,
+  RouterQueryEngine,
+  SimpleDirectoryReader,
+  SimpleNodeParser,
+  SummaryIndex,
+  VectorStoreIndex,
+  serviceContextFromDefaults,
+} from "llamaindex";
+
+async function main() {
+  // Load documents from a directory
+  const documents = await new SimpleDirectoryReader().loadData({
+    directoryPath: "node_modules/llamaindex/examples",
+  });
+
+  // Parse the documents into nodes
+  const nodeParser = new SimpleNodeParser({
+    chunkSize: 1024,
+  });
+
+  // Create a service context
+  const serviceContext = serviceContextFromDefaults({
+    nodeParser,
+    llm: new OpenAI(),
+  });
+
+  // Create indices
+  const vectorIndex = await VectorStoreIndex.fromDocuments(documents, {
+    serviceContext,
+  });
+
+  const summaryIndex = await SummaryIndex.fromDocuments(documents, {
+    serviceContext,
+  });
+
+  // Create query engines
+  const vectorQueryEngine = vectorIndex.asQueryEngine();
+  const summaryQueryEngine = summaryIndex.asQueryEngine();
+
+  // Create a router query engine
+  const queryEngine = RouterQueryEngine.fromDefaults({
+    queryEngineTools: [
+      {
+        queryEngine: vectorQueryEngine,
+        description: "Useful for summarization questions related to Abramov",
+      },
+      {
+        queryEngine: summaryQueryEngine,
+        description: "Useful for retrieving specific context from Abramov",
+      },
+    ],
+    serviceContext,
+  });
+
+  // Query the router query engine
+  const summaryResponse = await queryEngine.query({
+    query: "Give me a summary about his past experiences?",
+  });
+
+  console.log({
+    answer: summaryResponse.response,
+    metadata: summaryResponse?.metadata?.selectorResult,
+  });
+
+  const specificResponse = await queryEngine.query({
+    query: "Tell me about abramov first job?",
+  });
+
+  console.log({
+    answer: specificResponse.response,
+    metadata: specificResponse.metadata.selectorResult,
+  });
+}
+
+main().then(() => console.log("Done"));

packages/core/CHANGELOG.md

@@ -1,5 +1,55 @@
 # llamaindex
 
+## 0.1.8
+
+### Patch Changes
+
+- d903da6: easier prompt customization for SimpleResponseBuilder
+- ab9d941: fix(cyclic): remove cyclic structures from transform hash
+- 177b446: chore: improve extractors prompt
+
+## 0.1.7
+
+### Patch Changes
+
+- d687c11: feat(router): add router query engine
+
+## 0.1.6
+
+### Patch Changes
+
+- cf44640: fix: `instanceof` issue
+
+  This will fix QueryEngine cannot run.
+
+- 7231ddb: feat: allow `SimpleDirectoryReader` to get a string
+
+## 0.1.5
+
+### Patch Changes
+
+- 8a9b78a: chore: split readers into different files
+
+## 0.1.4
+
+### Patch Changes
+
+- 88696e1: refactor: use `pdf2json` instead of `pdfjs-dist`
+
+  Please add `pdf2json` to `serverComponentsExternalPackages` if you have to parse pdf in runtime.
+
+  ```js
+  // next.config.js
+  /** @type {import('next').NextConfig} */
+  const nextConfig = {
+    experimental: {
+      serverComponentsExternalPackages: ["pdf2json"],
+    },
+  };
+
+  module.exports = nextConfig;
+  ```
+
 ## 0.1.3
 
 ### Patch Changes

packages/core/package.json

@@ -1,7 +1,7 @@
 {
   "name": "llamaindex",
   "private": true,
-  "version": "0.1.3",
+  "version": "0.1.8",
   "license": "MIT",
   "dependencies": {
     "@anthropic-ai/sdk": "^0.12.4",
@@ -23,7 +23,7 @@
     "openai": "^4.26.0",
     "papaparse": "^5.4.1",
     "pathe": "^1.1.2",
-    "pdfjs-dist": "4.0.269",
+    "pdf2json": "^3.0.5",
     "pg": "^8.11.3",
     "pgvector": "^0.1.7",
     "portkey-ai": "^0.1.16",
@@ -40,7 +40,7 @@
     "@types/node": "^18.19.10",
     "@types/papaparse": "^5.3.14",
     "@types/pg": "^8.11.0",
-    "bunchee": "^4.4.2",
+    "bunchee": "^4.4.3",
     "edit-json-file": "^1.8.0",
     "madge": "^6.1.0",
     "typescript": "^5.3.3"
@@ -62,6 +62,131 @@
       "import": "./dist/env.mjs",
       "edge-light": "./dist/env.edge-light.mjs",
       "require": "./dist/env.js"
+    },
+    "./ChatEngine": {
+      "types": "./dist/ChatEngine.d.mts",
+      "import": "./dist/ChatEngine.mjs",
+      "require": "./dist/ChatEngine.js"
+    },
+    "./ChatHistory": {
+      "types": "./dist/ChatHistory.d.mts",
+      "import": "./dist/ChatHistory.mjs",
+      "require": "./dist/ChatHistory.js"
+    },
+    "./constants": {
+      "types": "./dist/constants.d.mts",
+      "import": "./dist/constants.mjs",
+      "require": "./dist/constants.js"
+    },
+    "./GlobalsHelper": {
+      "types": "./dist/GlobalsHelper.d.mts",
+      "import": "./dist/GlobalsHelper.mjs",
+      "require": "./dist/GlobalsHelper.js"
+    },
+    "./Node": {
+      "types": "./dist/Node.d.mts",
+      "import": "./dist/Node.mjs",
+      "require": "./dist/Node.js"
+    },
+    "./OutputParser": {
+      "types": "./dist/OutputParser.d.mts",
+      "import": "./dist/OutputParser.mjs",
+      "require": "./dist/OutputParser.js"
+    },
+    "./Prompt": {
+      "types": "./dist/Prompt.d.mts",
+      "import": "./dist/Prompt.mjs",
+      "require": "./dist/Prompt.js"
+    },
+    "./PromptHelper": {
+      "types": "./dist/PromptHelper.d.mts",
+      "import": "./dist/PromptHelper.mjs",
+      "require": "./dist/PromptHelper.js"
+    },
+    "./QueryEngine": {
+      "types": "./dist/QueryEngine.d.mts",
+      "import": "./dist/QueryEngine.mjs",
+      "require": "./dist/QueryEngine.js"
+    },
+    "./QuestionGenerator": {
+      "types": "./dist/QuestionGenerator.d.mts",
+      "import": "./dist/QuestionGenerator.mjs",
+      "require": "./dist/QuestionGenerator.js"
+    },
+    "./Response": {
+      "types": "./dist/Response.d.mts",
+      "import": "./dist/Response.mjs",
+      "require": "./dist/Response.js"
+    },
+    "./Retriever": {
+      "types": "./dist/Retriever.d.mts",
+      "import": "./dist/Retriever.mjs",
+      "require": "./dist/Retriever.js"
+    },
+    "./ServiceContext": {
+      "types": "./dist/ServiceContext.d.mts",
+      "import": "./dist/ServiceContext.mjs",
+      "require": "./dist/ServiceContext.js"
+    },
+    "./TextSplitter": {
+      "types": "./dist/TextSplitter.d.mts",
+      "import": "./dist/TextSplitter.mjs",
+      "require": "./dist/TextSplitter.js"
+    },
+    "./Tool": {
+      "types": "./dist/Tool.d.mts",
+      "import": "./dist/Tool.mjs",
+      "require": "./dist/Tool.js"
+    },
+    "./readers/AssemblyAIReader": {
+      "types": "./dist/readers/AssemblyAIReader.d.mts",
+      "import": "./dist/readers/AssemblyAIReader.mjs",
+      "require": "./dist/readers/AssemblyAIReader.js"
+    },
+    "./readers/CSVReader": {
+      "types": "./dist/readers/CSVReader.d.mts",
+      "import": "./dist/readers/CSVReader.mjs",
+      "require": "./dist/readers/CSVReader.js"
+    },
+    "./readers/DocxReader": {
+      "types": "./dist/readers/DocxReader.d.mts",
+      "import": "./dist/readers/DocxReader.mjs",
+      "require": "./dist/readers/DocxReader.js"
+    },
+    "./readers/HTMLReader": {
+      "types": "./dist/readers/HTMLReader.d.mts",
+      "import": "./dist/readers/HTMLReader.mjs",
+      "require": "./dist/readers/HTMLReader.js"
+    },
+    "./readers/ImageReader": {
+      "types": "./dist/readers/ImageReader.d.mts",
+      "import": "./dist/readers/ImageReader.mjs",
+      "require": "./dist/readers/ImageReader.js"
+    },
+    "./readers/MarkdownReader": {
+      "types": "./dist/readers/MarkdownReader.d.mts",
+      "import": "./dist/readers/MarkdownReader.mjs",
+      "require": "./dist/readers/MarkdownReader.js"
+    },
+    "./readers/NotionReader": {
+      "types": "./dist/readers/NotionReader.d.mts",
+      "import": "./dist/readers/NotionReader.mjs",
+      "require": "./dist/readers/NotionReader.js"
+    },
+    "./readers/PDFReader": {
+      "types": "./dist/readers/PDFReader.d.mts",
+      "import": "./dist/readers/PDFReader.mjs",
+      "require": "./dist/readers/PDFReader.js"
+    },
+    "./readers/SimpleDirectoryReader": {
+      "types": "./dist/readers/SimpleDirectoryReader.d.mts",
+      "import": "./dist/readers/SimpleDirectoryReader.mjs",
+      "require": "./dist/readers/SimpleDirectoryReader.js"
+    },
+    "./readers/SimpleMongoReader": {
+      "types": "./dist/readers/SimpleMongoReader.d.mts",
+      "import": "./dist/readers/SimpleMongoReader.mjs",
+      "require": "./dist/readers/SimpleMongoReader.js"
     }
   },
   "files": [
@@ -75,12 +200,12 @@
   "scripts": {
     "lint": "eslint .",
     "test": "jest",
-    "build": "bunchee",
+    "build": "rm -rf ./dist && NODE_OPTIONS=\"--max-old-space-size=8192\" bunchee",
     "postbuild": "pnpm run copy && pnpm run modify-package-json",
     "copy": "cp -r package.json CHANGELOG.md ../../README.md ../../LICENSE examples src dist/",
     "modify-package-json": "node ./scripts/modify-package-json.mjs",
     "prepublish": "pnpm run modify-package-json && echo \"please cd ./dist and run pnpm publish\" && exit 1",
-    "dev": "bunchee -w",
+    "dev": "NODE_OPTIONS=\"--max-old-space-size=8192\" bunchee -w",
     "circular-check": "madge -c ./src/index.ts"
   }
 }

packages/core/src/OutputParser.ts

@@ -1,4 +1,5 @@
-import { BaseOutputParser, StructuredOutput, SubQuestion } from "./types";
+import { SubQuestion } from "./engines/query/types";
+import { BaseOutputParser, StructuredOutput } from "./types";
 
 /**
  * Error class for output parsing. Due to the nature of LLMs, anytime we use LLM
@@ -73,9 +74,6 @@ export class SubQuestionOutputParser
 {
   parse(output: string): StructuredOutput<SubQuestion[]> {
     const parsed = parseJsonMarkdown(output);
-
-    // TODO add zod validation
-
     return { rawOutput: output, parsedOutput: parsed };
   }
 

packages/core/src/Prompt.ts

@@ -1,5 +1,6 @@
+import { SubQuestion } from "./engines/query/types";
 import { ChatMessage } from "./llm/types";
-import { SubQuestion, ToolMetadata } from "./types";
+import { ToolMetadata } from "./types";
 
 /**
  * A SimplePrompt is a function that takes a dictionary of inputs and returns a string.
@@ -35,7 +36,10 @@ Answer:`;
 
 export type TextQaPrompt = typeof defaultTextQaPrompt;
 
-export const anthropicTextQaPrompt = ({ context = "", query = "" }) => {
+export const anthropicTextQaPrompt: TextQaPrompt = ({
+  context = "",
+  query = "",
+}) => {
   return `Context information:
 <context>
 ${context}
@@ -71,6 +75,16 @@ SUMMARY:"""
 
 export type SummaryPrompt = typeof defaultSummaryPrompt;
 
+export const anthropicSummaryPrompt: SummaryPrompt = ({ context = "" }) => {
+  return `Summarize the following text. Try to use only the information provided. Try to include as many key details as possible.
+<original-text>
+${context}
+</original-text>
+
+SUMMARY:
+`;
+};
+
 /*
 DEFAULT_REFINE_PROMPT_TMPL = (
     "The original query is as follows: {query_str}\n"

packages/core/src/QuestionGenerator.ts

@@ -4,15 +4,10 @@ import {
   buildToolsText,
   defaultSubQuestionPrompt,
 } from "./Prompt";
+import { BaseQuestionGenerator, SubQuestion } from "./engines/query/types";
 import { OpenAI } from "./llm/LLM";
 import { LLM } from "./llm/types";
-import {
-  BaseOutputParser,
-  BaseQuestionGenerator,
-  StructuredOutput,
-  SubQuestion,
-  ToolMetadata,
-} from "./types";
+import { BaseOutputParser, StructuredOutput, ToolMetadata } from "./types";
 
 /**
  * LLMQuestionGenerator uses the LLM to generate new questions for the LLM using tools and a user query.

packages/core/src/Response.ts

@@ -6,13 +6,14 @@ import { BaseNode } from "./Node";
 export class Response {
   response: string;
   sourceNodes?: BaseNode[];
+  metadata: Record<string, unknown> = {};
 
   constructor(response: string, sourceNodes?: BaseNode[]) {
     this.response = response;
     this.sourceNodes = sourceNodes || [];
   }
 
-  getFormattedSources() {
+  protected _getFormattedSources() {
     throw new Error("Not implemented yet");
   }
 

packages/core/src/agent/index.ts

@@ -0,0 +1,2 @@
+export * from "./openai/base";
+export * from "./openai/worker";

packages/core/src/agent/openai/base.ts

@@ -0,0 +1,55 @@
+import { CallbackManager } from "../../callbacks/CallbackManager";
+import { ChatMessage, OpenAI } from "../../llm";
+import { ObjectRetriever } from "../../objects/base";
+import { BaseTool } from "../../types";
+import { AgentRunner } from "../runner/base";
+import { OpenAIAgentWorker } from "./worker";
+
+type OpenAIAgentParams = {
+  tools: BaseTool[];
+  llm?: OpenAI;
+  memory?: any;
+  prefixMessages?: ChatMessage[];
+  verbose?: boolean;
+  maxFunctionCalls?: number;
+  defaultToolChoice?: string;
+  callbackManager?: CallbackManager;
+  toolRetriever?: ObjectRetriever<BaseTool>;
+};
+
+/**
+ * An agent that uses OpenAI's API to generate text.
+ *
+ * @category OpenAI
+ */
+export class OpenAIAgent extends AgentRunner {
+  constructor({
+    tools,
+    llm,
+    memory,
+    prefixMessages,
+    verbose,
+    maxFunctionCalls = 5,
+    defaultToolChoice = "auto",
+    callbackManager,
+    toolRetriever,
+  }: OpenAIAgentParams) {
+    const stepEngine = new OpenAIAgentWorker({
+      tools,
+      callbackManager,
+      llm,
+      prefixMessages,
+      maxFunctionCalls,
+      toolRetriever,
+      verbose,
+    });
+
+    super({
+      agentWorker: stepEngine,
+      memory,
+      callbackManager,
+      defaultToolChoice,
+      chatHistory: prefixMessages,
+    });
+  }
+}

packages/core/src/agent/openai/types/chat.ts

@@ -0,0 +1,13 @@
+export type OpenAIToolCall = ChatCompletionMessageToolCall;
+
+export interface Function {
+  arguments: string;
+  name: string;
+  type: "function";
+}
+
+export interface ChatCompletionMessageToolCall {
+  id: string;
+  function: Function;
+  type: "function";
+}

packages/core/src/agent/openai/utils.ts

@@ -0,0 +1,27 @@
+import { ToolMetadata } from "../../types";
+
+export type OpenAIFunction = {
+  type: "function";
+  function: ToolMetadata;
+};
+
+type OpenAiTool = {
+  name: string;
+  description: string;
+  parameters: ToolMetadata["parameters"];
+};
+
+export const toOpenAiTool = ({
+  name,
+  description,
+  parameters,
+}: OpenAiTool): OpenAIFunction => {
+  return {
+    type: "function",
+    function: {
+      name: name,
+      description: description,
+      parameters,
+    },
+  };
+};

packages/core/src/agent/openai/worker.ts

@@ -0,0 +1,405 @@
+// Assuming that the necessary interfaces and classes (like BaseTool, OpenAI, ChatMessage, CallbackManager, etc.) are defined elsewhere
+
+import { CallbackManager } from "../../callbacks/CallbackManager";
+import { AgentChatResponse, ChatResponseMode } from "../../engines/chat";
+import { randomUUID } from "../../env";
+import {
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  OpenAI,
+} from "../../llm";
+import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer";
+import { ObjectRetriever } from "../../objects/base";
+import { ToolOutput } from "../../tools/types";
+import { callToolWithErrorHandling } from "../../tools/utils";
+import { BaseTool } from "../../types";
+import { AgentWorker, Task, TaskStep, TaskStepOutput } from "../types";
+import { addUserStepToMemory, getFunctionByName } from "../utils";
+import { OpenAIToolCall } from "./types/chat";
+import { toOpenAiTool } from "./utils";
+
+const DEFAULT_MAX_FUNCTION_CALLS = 5;
+
+/**
+ * Call function.
+ * @param tools: tools
+ * @param toolCall: tool call
+ * @param verbose: verbose
+ * @returns: void
+ */
+async function callFunction(
+  tools: BaseTool[],
+  toolCall: OpenAIToolCall,
+  verbose: boolean = false,
+): Promise<[ChatMessage, ToolOutput]> {
+  const id_ = toolCall.id;
+  const functionCall = toolCall.function;
+  const name = toolCall.function.name;
+  const argumentsStr = toolCall.function.arguments;
+
+  if (verbose) {
+    console.log("=== Calling Function ===");
+    console.log(`Calling function: ${name} with args: ${argumentsStr}`);
+  }
+
+  const tool = getFunctionByName(tools, name);
+  const argumentDict = JSON.parse(argumentsStr);
+
+  // Call tool
+  // Use default error message
+  const output = await callToolWithErrorHandling(tool, argumentDict, null);
+
+  if (verbose) {
+    console.log(`Got output ${output}`);
+    console.log("==========================");
+  }
+
+  return [
+    {
+      content: String(output),
+      role: "tool",
+      additionalKwargs: {
+        name,
+        tool_call_id: id_,
+      },
+    },
+    output,
+  ];
+}
+
+type OpenAIAgentWorkerParams = {
+  tools: BaseTool[];
+  llm?: OpenAI;
+  prefixMessages?: ChatMessage[];
+  verbose?: boolean;
+  maxFunctionCalls?: number;
+  callbackManager?: CallbackManager | undefined;
+  toolRetriever?: ObjectRetriever<BaseTool>;
+};
+
+type CallFunctionOutput = {
+  message: ChatMessage;
+  toolOutput: ToolOutput;
+};
+
+/**
+ * OpenAI agent worker.
+ * This class is responsible for running the agent.
+ */
+export class OpenAIAgentWorker implements AgentWorker {
+  private _llm: OpenAI;
+  private _verbose: boolean;
+  private _maxFunctionCalls: number;
+
+  public prefixMessages: ChatMessage[];
+  public callbackManager: CallbackManager | undefined;
+
+  private _getTools: (input: string) => BaseTool[];
+
+  /**
+   * Initialize.
+   */
+  constructor({
+    tools,
+    llm,
+    prefixMessages,
+    verbose,
+    maxFunctionCalls = DEFAULT_MAX_FUNCTION_CALLS,
+    callbackManager,
+    toolRetriever,
+  }: OpenAIAgentWorkerParams) {
+    this._llm = llm ?? new OpenAI({ model: "gpt-3.5-turbo-0613" });
+    this._verbose = verbose || false;
+    this._maxFunctionCalls = maxFunctionCalls;
+    this.prefixMessages = prefixMessages || [];
+    this.callbackManager = callbackManager || this._llm.callbackManager;
+
+    if (tools.length > 0 && toolRetriever) {
+      throw new Error("Cannot specify both tools and tool_retriever");
+    } else if (tools.length > 0) {
+      this._getTools = () => tools;
+    } else if (toolRetriever) {
+      // @ts-ignore
+      this._getTools = (message: string) => toolRetriever.retrieve(message);
+    } else {
+      this._getTools = () => [];
+    }
+  }
+
+  /**
+   * Get all messages.
+   * @param task: task
+   * @returns: messages
+   */
+  public getAllMessages(task: Task): ChatMessage[] {
+    return [
+      ...this.prefixMessages,
+      ...task.memory.get(),
+      ...task.extraState.newMemory.get(),
+    ];
+  }
+
+  /**
+   * Get latest tool calls.
+   * @param task: task
+   * @returns: tool calls
+   */
+  public getLatestToolCalls(task: Task): OpenAIToolCall[] | null {
+    const chatHistory: ChatMessage[] = task.extraState.newMemory.getAll();
+
+    if (chatHistory.length === 0) {
+      return null;
+    }
+
+    return chatHistory[chatHistory.length - 1].additionalKwargs?.toolCalls;
+  }
+
+  /**
+   *
+   * @param task
+   * @param openaiTools
+   * @param toolChoice
+   * @returns
+   */
+  private _getLlmChatKwargs(
+    task: Task,
+    openaiTools: { [key: string]: any }[],
+    toolChoice: string | { [key: string]: any } = "auto",
+  ): { [key: string]: any } {
+    const llmChatKwargs: { [key: string]: any } = {
+      messages: this.getAllMessages(task),
+    };
+
+    if (openaiTools.length > 0) {
+      llmChatKwargs.tools = openaiTools;
+      llmChatKwargs.toolChoice = toolChoice;
+    }
+
+    return llmChatKwargs;
+  }
+
+  /**
+   * Process message.
+   * @param task: task
+   * @param chatResponse: chat response
+   * @returns: agent chat response
+   */
+  private _processMessage(
+    task: Task,
+    chatResponse: ChatResponse,
+  ): AgentChatResponse | AsyncIterable<ChatResponseChunk> {
+    const aiMessage = chatResponse.message;
+    task.extraState.newMemory.put(aiMessage);
+    return new AgentChatResponse(aiMessage.content, task.extraState.sources);
+  }
+
+  /**
+   * Get agent response.
+   * @param task: task
+   * @param mode: mode
+   * @param llmChatKwargs: llm chat kwargs
+   * @returns: agent chat response
+   */
+  private async _getAgentResponse(
+    task: Task,
+    mode: ChatResponseMode,
+    llmChatKwargs: any,
+  ): Promise<AgentChatResponse> {
+    if (mode === ChatResponseMode.WAIT) {
+      const chatResponse = (await this._llm.chat({
+        stream: false,
+        ...llmChatKwargs,
+      })) as unknown as ChatResponse;
+
+      return this._processMessage(task, chatResponse) as AgentChatResponse;
+    } else {
+      throw new Error("Not implemented");
+    }
+  }
+
+  /**
+   * Call function.
+   * @param tools: tools
+   * @param toolCall: tool call
+   * @param memory: memory
+   * @param sources: sources
+   * @returns: void
+   */
+  async callFunction(
+    tools: BaseTool[],
+    toolCall: OpenAIToolCall,
+  ): Promise<CallFunctionOutput> {
+    const functionCall = toolCall.function;
+
+    if (!functionCall) {
+      throw new Error("Invalid tool_call object");
+    }
+
+    const functionMessage = await callFunction(tools, toolCall, this._verbose);
+
+    const message = functionMessage[0];
+    const toolOutput = functionMessage[1];
+
+    return {
+      message,
+      toolOutput,
+    };
+  }
+
+  /**
+   * Initialize step.
+   * @param task: task
+   * @param kwargs: kwargs
+   * @returns: task step
+   */
+  initializeStep(task: Task, kwargs?: any): TaskStep {
+    const sources: ToolOutput[] = [];
+
+    const newMemory = new ChatMemoryBuffer();
+
+    const taskState = {
+      sources,
+      nFunctionCalls: 0,
+      newMemory,
+    };
+
+    task.extraState = {
+      ...task.extraState,
+      ...taskState,
+    };
+
+    return new TaskStep(task.taskId, randomUUID(), task.input);
+  }
+
+  /**
+   * Should continue.
+   * @param toolCalls: tool calls
+   * @param nFunctionCalls: number of function calls
+   * @returns: boolean
+   */
+  private _shouldContinue(
+    toolCalls: OpenAIToolCall[] | null,
+    nFunctionCalls: number,
+  ): boolean {
+    if (nFunctionCalls > this._maxFunctionCalls) {
+      return false;
+    }
+
+    if (toolCalls?.length === 0) {
+      return false;
+    }
+
+    return true;
+  }
+
+  /**
+   * Get tools.
+   * @param input: input
+   * @returns: tools
+   */
+  getTools(input: string): BaseTool[] {
+    return this._getTools(input);
+  }
+
+  private async _runStep(
+    step: TaskStep,
+    task: Task,
+    mode: ChatResponseMode = ChatResponseMode.WAIT,
+    toolChoice: string | { [key: string]: any } = "auto",
+  ): Promise<TaskStepOutput> {
+    const tools = this.getTools(task.input);
+
+    if (step.input) {
+      addUserStepToMemory(step, task.extraState.newMemory, this._verbose);
+    }
+
+    const openaiTools = tools.map((tool) =>
+      toOpenAiTool({
+        name: tool.metadata.name,
+        description: tool.metadata.description,
+        parameters: tool.metadata.parameters,
+      }),
+    );
+
+    const llmChatKwargs = this._getLlmChatKwargs(task, openaiTools, toolChoice);
+
+    const agentChatResponse = await this._getAgentResponse(
+      task,
+      mode,
+      llmChatKwargs,
+    );
+
+    const latestToolCalls = this.getLatestToolCalls(task) || [];
+
+    let isDone: boolean;
+    let newSteps: TaskStep[] = [];
+
+    if (
+      !this._shouldContinue(latestToolCalls, task.extraState.nFunctionCalls)
+    ) {
+      isDone = true;
+      newSteps = [];
+    } else {
+      isDone = false;
+      for (const toolCall of latestToolCalls) {
+        const { message, toolOutput } = await this.callFunction(
+          tools,
+          toolCall,
+        );
+
+        task.extraState.sources.push(toolOutput);
+        task.extraState.newMemory.put(message);
+
+        task.extraState.nFunctionCalls += 1;
+      }
+
+      newSteps = [step.getNextStep(randomUUID(), undefined)];
+    }
+
+    return new TaskStepOutput(agentChatResponse, step, newSteps, isDone);
+  }
+
+  /**
+   * Run step.
+   * @param step: step
+   * @param task: task
+   * @param kwargs: kwargs
+   * @returns: task step output
+   */
+  async runStep(
+    step: TaskStep,
+    task: Task,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
+    const toolChoice = kwargs?.toolChoice || "auto";
+    return this._runStep(step, task, ChatResponseMode.WAIT, toolChoice);
+  }
+
+  /**
+   * Stream step.
+   * @param step: step
+   * @param task: task
+   * @param kwargs: kwargs
+   * @returns: task step output
+   */
+  async streamStep(
+    step: TaskStep,
+    task: Task,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
+    const toolChoice = kwargs?.toolChoice || "auto";
+    return this._runStep(step, task, ChatResponseMode.STREAM, toolChoice);
+  }
+
+  /**
+   * Finalize task.
+   * @param task: task
+   * @param kwargs: kwargs
+   * @returns: void
+   */
+  finalizeTask(task: Task, kwargs?: any): void {
+    task.memory.set(task.memory.get().concat(task.extraState.newMemory.get()));
+    task.extraState.newMemory.reset();
+  }
+}

packages/core/src/agent/runner/base.ts

@@ -0,0 +1,343 @@
+import { randomUUID } from "crypto";
+import { CallbackManager } from "../../callbacks/CallbackManager";
+import {
+  AgentChatResponse,
+  ChatEngineAgentParams,
+  ChatResponseMode,
+} from "../../engines/chat";
+import { ChatMessage, LLM } from "../../llm";
+import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer";
+import { BaseMemory } from "../../memory/types";
+import { AgentWorker, Task, TaskStep, TaskStepOutput } from "../types";
+import { AgentState, BaseAgentRunner, TaskState } from "./types";
+
+const validateStepFromArgs = (
+  taskId: string,
+  input: string,
+  step?: any,
+  kwargs?: any,
+): TaskStep | undefined => {
+  if (step) {
+    if (input) {
+      throw new Error("Cannot specify both `step` and `input`");
+    }
+    return step;
+  } else {
+    return new TaskStep(taskId, step, input, kwargs);
+  }
+};
+
+type AgentRunnerParams = {
+  agentWorker: AgentWorker;
+  chatHistory?: ChatMessage[];
+  state?: AgentState;
+  memory?: BaseMemory;
+  llm?: LLM;
+  callbackManager?: CallbackManager;
+  initTaskStateKwargs?: Record<string, any>;
+  deleteTaskOnFinish?: boolean;
+  defaultToolChoice?: string;
+};
+
+export class AgentRunner extends BaseAgentRunner {
+  agentWorker: AgentWorker;
+  state: AgentState;
+  memory: BaseMemory;
+  callbackManager: CallbackManager;
+  initTaskStateKwargs: Record<string, any>;
+  deleteTaskOnFinish: boolean;
+  defaultToolChoice: string;
+
+  /**
+   * Creates an AgentRunner.
+   */
+  constructor(params: AgentRunnerParams) {
+    super();
+
+    this.agentWorker = params.agentWorker;
+    this.state = params.state ?? new AgentState();
+    this.memory =
+      params.memory ??
+      new ChatMemoryBuffer({
+        chatHistory: params.chatHistory,
+      });
+    this.callbackManager = params.callbackManager ?? new CallbackManager();
+    this.initTaskStateKwargs = params.initTaskStateKwargs ?? {};
+    this.deleteTaskOnFinish = params.deleteTaskOnFinish ?? false;
+    this.defaultToolChoice = params.defaultToolChoice ?? "auto";
+  }
+
+  /**
+   * Creates a task.
+   * @param input
+   * @param kwargs
+   */
+  createTask(input: string, kwargs?: any): Task {
+    let extraState;
+
+    if (!this.initTaskStateKwargs) {
+      if (kwargs && "extraState" in kwargs) {
+        if (extraState) {
+          delete extraState["extraState"];
+        }
+      }
+    } else {
+      if (kwargs && "extraState" in kwargs) {
+        throw new Error(
+          "Cannot specify both `extraState` and `initTaskStateKwargs`",
+        );
+      } else {
+        extraState = this.initTaskStateKwargs;
+      }
+    }
+
+    const task = new Task({
+      taskId: randomUUID(),
+      input,
+      memory: this.memory,
+      extraState,
+      ...kwargs,
+    });
+
+    const initialStep = this.agentWorker.initializeStep(task);
+
+    const taskState = new TaskState({
+      task,
+      stepQueue: [initialStep],
+    });
+
+    this.state.taskDict[task.taskId] = taskState;
+
+    return task;
+  }
+
+  /**
+   * Deletes the task.
+   * @param taskId
+   */
+  deleteTask(taskId: string): void {
+    delete this.state.taskDict[taskId];
+  }
+
+  /**
+   * Returns the list of tasks.
+   */
+  listTasks(): Task[] {
+    return Object.values(this.state.taskDict).map(
+      (taskState) => taskState.task,
+    );
+  }
+
+  /**
+   * Returns the task.
+   */
+  getTask(taskId: string): Task {
+    return this.state.taskDict[taskId].task;
+  }
+
+  /**
+   * Returns the completed steps in the task.
+   * @param taskId
+   * @param kwargs
+   */
+  getCompletedSteps(taskId: string): TaskStepOutput[] {
+    return this.state.taskDict[taskId].completedSteps;
+  }
+
+  /**
+   * Returns the next steps in the task.
+   * @param taskId
+   * @param kwargs
+   */
+  getUpcomingSteps(taskId: string, kwargs: any): TaskStep[] {
+    return this.state.taskDict[taskId].stepQueue;
+  }
+
+  private async _runStep(
+    taskId: string,
+    step?: TaskStep,
+    mode: ChatResponseMode = ChatResponseMode.WAIT,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
+    const task = this.state.getTask(taskId);
+    const curStep = step || this.state.getStepQueue(taskId).shift();
+
+    let curStepOutput;
+
+    if (!curStep) {
+      throw new Error(`No step found for task ${taskId}`);
+    }
+
+    if (mode === ChatResponseMode.WAIT) {
+      curStepOutput = await this.agentWorker.runStep(curStep, task, kwargs);
+    } else if (mode === ChatResponseMode.STREAM) {
+      curStepOutput = await this.agentWorker.streamStep(curStep, task, kwargs);
+    } else {
+      throw new Error(`Invalid mode: ${mode}`);
+    }
+
+    const nextSteps = curStepOutput.nextSteps;
+
+    this.state.addSteps(taskId, nextSteps);
+    this.state.addCompletedStep(taskId, [curStepOutput]);
+
+    return curStepOutput;
+  }
+
+  /**
+   * Runs the next step in the task.
+   * @param taskId
+   * @param kwargs
+   * @param step
+   * @returns
+   */
+  async runStep(
+    taskId: string,
+    input: string,
+    step?: TaskStep,
+    kwargs: any = {},
+  ): Promise<TaskStepOutput> {
+    const curStep = validateStepFromArgs(taskId, input, step, kwargs);
+    return this._runStep(taskId, curStep, ChatResponseMode.WAIT, kwargs);
+  }
+
+  /**
+   * Runs the step and returns the response.
+   * @param taskId
+   * @param input
+   * @param step
+   * @param kwargs
+   */
+  async streamStep(
+    taskId: string,
+    input: string,
+    step?: TaskStep,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
+    const curStep = validateStepFromArgs(taskId, input, step, kwargs);
+    return this._runStep(taskId, curStep, ChatResponseMode.STREAM, kwargs);
+  }
+
+  /**
+   * Finalizes the response and returns it.
+   * @param taskId
+   * @param kwargs
+   * @param stepOutput
+   * @returns
+   */
+  async finalizeResponse(
+    taskId: string,
+    stepOutput: TaskStepOutput,
+    kwargs?: any,
+  ): Promise<AgentChatResponse> {
+    if (!stepOutput) {
+      stepOutput =
+        this.getCompletedSteps(taskId)[
+          this.getCompletedSteps(taskId).length - 1
+        ];
+    }
+    if (!stepOutput.isLast) {
+      throw new Error(
+        "finalizeResponse can only be called on the last step output",
+      );
+    }
+
+    if (!(stepOutput.output instanceof AgentChatResponse)) {
+      throw new Error(
+        `When \`isLast\` is True, cur_step_output.output must be AGENT_CHAT_RESPONSE_TYPE: ${stepOutput.output}`,
+      );
+    }
+
+    this.agentWorker.finalizeTask(this.getTask(taskId), kwargs);
+
+    if (this.deleteTaskOnFinish) {
+      this.deleteTask(taskId);
+    }
+
+    return stepOutput.output;
+  }
+
+  protected async _chat({
+    message,
+    toolChoice,
+  }: ChatEngineAgentParams & { mode: ChatResponseMode }) {
+    const task = this.createTask(message as string);
+
+    let resultOutput;
+
+    while (true) {
+      const curStepOutput = await this._runStep(task.taskId);
+
+      if (curStepOutput.isLast) {
+        resultOutput = curStepOutput;
+        break;
+      }
+
+      toolChoice = "auto";
+    }
+
+    return this.finalizeResponse(task.taskId, resultOutput);
+  }
+
+  /**
+   * Sends a message to the LLM and returns the response.
+   * @param message
+   * @param chatHistory
+   * @param toolChoice
+   * @returns
+   */
+  public async chat({
+    message,
+    chatHistory,
+    toolChoice,
+  }: ChatEngineAgentParams): Promise<AgentChatResponse> {
+    if (!toolChoice) {
+      toolChoice = this.defaultToolChoice;
+    }
+
+    const chatResponse = await this._chat({
+      message,
+      chatHistory,
+      toolChoice,
+      mode: ChatResponseMode.WAIT,
+    });
+
+    return chatResponse;
+  }
+
+  protected _getPromptModules(): string[] {
+    return [];
+  }
+
+  protected _getPrompts(): string[] {
+    return [];
+  }
+
+  /**
+   * Resets the agent.
+   */
+  reset(): void {
+    this.state = new AgentState();
+  }
+
+  getCompletedStep(
+    taskId: string,
+    stepId: string,
+    kwargs: any,
+  ): TaskStepOutput {
+    const completedSteps = this.getCompletedSteps(taskId);
+    for (const stepOutput of completedSteps) {
+      if (stepOutput.taskStep.stepId === stepId) {
+        return stepOutput;
+      }
+    }
+
+    throw new Error(`Step ${stepId} not found in task ${taskId}`);
+  }
+
+  /**
+   * Undoes the step.
+   * @param taskId
+   */
+  undoStep(taskId: string): void {}
+}

packages/core/src/agent/runner/types.ts

@@ -0,0 +1,102 @@
+import { AgentChatResponse } from "../../engines/chat";
+import { BaseAgent, Task, TaskStep, TaskStepOutput } from "../types";
+
+export class TaskState {
+  task!: Task;
+  stepQueue!: TaskStep[];
+  completedSteps!: TaskStepOutput[];
+
+  constructor(init?: Partial<TaskState>) {
+    Object.assign(this, init);
+  }
+}
+
+export abstract class BaseAgentRunner extends BaseAgent {
+  constructor(init?: Partial<BaseAgentRunner>) {
+    super();
+  }
+
+  abstract createTask(input: string, kwargs: any): Task;
+  abstract deleteTask(taskId: string): void;
+  abstract getTask(taskId: string, kwargs: any): Task;
+  abstract listTasks(kwargs: any): Task[];
+  abstract getUpcomingSteps(taskId: string, kwargs: any): TaskStep[];
+  abstract getCompletedSteps(taskId: string, kwargs: any): TaskStepOutput[];
+
+  getCompletedStep(
+    taskId: string,
+    stepId: string,
+    kwargs: any,
+  ): TaskStepOutput {
+    const completedSteps = this.getCompletedSteps(taskId, kwargs);
+    for (const stepOutput of completedSteps) {
+      if (stepOutput.taskStep.stepId === stepId) {
+        return stepOutput;
+      }
+    }
+
+    throw new Error(`Step ${stepId} not found in task ${taskId}`);
+  }
+
+  abstract runStep(
+    taskId: string,
+    input: string,
+    step: TaskStep,
+    kwargs: any,
+  ): Promise<TaskStepOutput>;
+
+  abstract streamStep(
+    taskId: string,
+    input: string,
+    step: TaskStep,
+    kwargs?: any,
+  ): Promise<TaskStepOutput>;
+
+  abstract finalizeResponse(
+    taskId: string,
+    stepOutput: TaskStepOutput,
+    kwargs?: any,
+  ): Promise<AgentChatResponse>;
+
+  abstract undoStep(taskId: string): void;
+}
+
+export class AgentState {
+  taskDict!: Record<string, TaskState>;
+
+  constructor(init?: Partial<AgentState>) {
+    Object.assign(this, init);
+
+    if (!this.taskDict) {
+      this.taskDict = {};
+    }
+  }
+
+  getTask(taskId: string): Task {
+    return this.taskDict[taskId].task;
+  }
+
+  getCompletedSteps(taskId: string): TaskStepOutput[] {
+    return this.taskDict[taskId].completedSteps || [];
+  }
+
+  getStepQueue(taskId: string): TaskStep[] {
+    return this.taskDict[taskId].stepQueue || [];
+  }
+
+  addSteps(taskId: string, steps: TaskStep[]): void {
+    if (!this.taskDict[taskId].stepQueue) {
+      this.taskDict[taskId].stepQueue = [];
+    }
+
+    this.taskDict[taskId].stepQueue.push(...steps);
+  }
+
+  addCompletedStep(taskId: string, stepOutputs: TaskStepOutput[]): void {
+    if (!this.taskDict[taskId].completedSteps) {
+      this.taskDict[taskId].completedSteps = [];
+    }
+
+    this.taskDict[taskId].completedSteps.push(...stepOutputs);
+  }
+}

packages/core/src/agent/types.ts

@@ -0,0 +1,181 @@
+import { AgentChatResponse, ChatEngineAgentParams } from "../engines/chat";
+import { QueryEngineParamsNonStreaming } from "../types";
+
+export interface AgentWorker {
+  initializeStep(task: Task, kwargs?: any): TaskStep;
+  runStep(step: TaskStep, task: Task, kwargs?: any): Promise<TaskStepOutput>;
+  streamStep(step: TaskStep, task: Task, kwargs?: any): Promise<TaskStepOutput>;
+  finalizeTask(task: Task, kwargs?: any): void;
+}
+
+interface BaseChatEngine {
+  chat(params: ChatEngineAgentParams): Promise<AgentChatResponse>;
+}
+
+interface BaseQueryEngine {
+  query(params: QueryEngineParamsNonStreaming): Promise<AgentChatResponse>;
+}
+
+/**
+ * BaseAgent is the base class for all agents.
+ */
+export abstract class BaseAgent implements BaseChatEngine, BaseQueryEngine {
+  protected _getPrompts(): string[] {
+    return [];
+  }
+
+  protected _getPromptModules(): string[] {
+    return [];
+  }
+
+  abstract chat(params: ChatEngineAgentParams): Promise<AgentChatResponse>;
+  abstract reset(): void;
+
+  /**
+   * query is the main entrypoint for the agent. It takes a query and returns a response.
+   * @param params
+   * @returns
+   */
+  async query(
+    params: QueryEngineParamsNonStreaming,
+  ): Promise<AgentChatResponse> {
+    // Handle non-streaming query
+    const agentResponse = await this.chat({
+      message: params.query,
+      chatHistory: [],
+    });
+
+    return agentResponse;
+  }
+}
+
+type TaskParams = {
+  taskId: string;
+  input: string;
+  memory: any;
+  extraState: Record<string, any>;
+};
+
+/**
+ * Task is a unit of work for the agent.
+ * @param taskId: taskId
+ */
+export class Task {
+  taskId: string;
+  input: string;
+
+  memory: any;
+  extraState: Record<string, any>;
+
+  constructor({ taskId, input, memory, extraState }: TaskParams) {
+    this.taskId = taskId;
+    this.input = input;
+    this.memory = memory;
+    this.extraState = extraState ?? {};
+  }
+}
+
+interface ITaskStep {
+  taskId: string;
+  stepId: string;
+  input?: string | null;
+  stepState: Record<string, any>;
+  nextSteps: Record<string, TaskStep>;
+  prevSteps: Record<string, TaskStep>;
+  isReady: boolean;
+  getNextStep(
+    stepId: string,
+    input?: string,
+    stepState?: Record<string, any>,
+  ): TaskStep;
+  linkStep(nextStep: TaskStep): void;
+}
+
+/**
+ * TaskStep is a unit of work for the agent.
+ * @param taskId: taskId
+ * @param stepId: stepId
+ * @param input: input
+ * @param stepState: stepState
+ */
+export class TaskStep implements ITaskStep {
+  taskId: string;
+  stepId: string;
+  input?: string | null;
+  stepState: Record<string, any> = {};
+  nextSteps: Record<string, TaskStep> = {};
+  prevSteps: Record<string, TaskStep> = {};
+  isReady: boolean = true;
+
+  constructor(
+    taskId: string,
+    stepId: string,
+    input?: string | null,
+    stepState?: Record<string, any> | null,
+  ) {
+    this.taskId = taskId;
+    this.stepId = stepId;
+    this.input = input;
+    this.stepState = stepState ?? this.stepState;
+  }
+
+  /*
+   * getNextStep is a function that returns the next step.
+   * @param stepId: stepId
+   * @param input: input
+   * @param stepState: stepState
+   * @returns: TaskStep
+   */
+  getNextStep(
+    stepId: string,
+    input?: string,
+    stepState?: Record<string, unknown>,
+  ): TaskStep {
+    return new TaskStep(
+      this.taskId,
+      stepId,
+      input,
+      stepState ?? this.stepState,
+    );
+  }
+
+  /*
+   * linkStep is a function that links the next step.
+   * @param nextStep: nextStep
+   * @returns: void
+   */
+  linkStep(nextStep: TaskStep): void {
+    this.nextSteps[nextStep.stepId] = nextStep;
+    nextStep.prevSteps[this.stepId] = this;
+  }
+}
+
+/**
+ * TaskStepOutput is a unit of work for the agent.
+ * @param output: output
+ * @param taskStep: taskStep
+ * @param nextSteps: nextSteps
+ * @param isLast: isLast
+ */
+export class TaskStepOutput {
+  output: unknown;
+  taskStep: TaskStep;
+  nextSteps: TaskStep[];
+  isLast: boolean;
+
+  constructor(
+    output: unknown,
+    taskStep: TaskStep,
+    nextSteps: TaskStep[],
+    isLast: boolean = false,
+  ) {
+    this.output = output;
+    this.taskStep = taskStep;
+    this.nextSteps = nextSteps;
+    this.isLast = isLast;
+  }
+
+  toString(): string {
+    return String(this.output);
+  }
+}

packages/core/src/agent/utils.ts

@@ -0,0 +1,51 @@
+import { ChatMessage } from "../llm";
+import { ChatMemoryBuffer } from "../memory/ChatMemoryBuffer";
+import { BaseTool } from "../types";
+import { TaskStep } from "./types";
+
+/**
+ * Adds the user's input to the memory.
+ *
+ * @param step - The step to add to the memory.
+ * @param memory - The memory to add the step to.
+ * @param verbose - Whether to print debug messages.
+ */
+export function addUserStepToMemory(
+  step: TaskStep,
+  memory: ChatMemoryBuffer,
+  verbose: boolean = false,
+): void {
+  if (!step.input) {
+    return;
+  }
+
+  const userMessage: ChatMessage = {
+    content: step.input,
+    role: "user",
+  };
+
+  memory.put(userMessage);
+
+  if (verbose) {
+    console.log(`Added user message to memory!: ${userMessage.content}`);
+  }
+}
+
+/**
+ * Get function by name.
+ * @param tools: tools
+ * @param name: name
+ * @returns: tool
+ */
+export function getFunctionByName(tools: BaseTool[], name: string): BaseTool {
+  const nameToTool: { [key: string]: BaseTool } = {};
+  tools.forEach((tool) => {
+    nameToTool[tool.metadata.name] = tool;
+  });
+
+  if (!(name in nameToTool)) {
+    throw new Error(`Tool with name ${name} not found`);
+  }
+
+  return nameToTool[name];
+}

packages/core/src/embeddings/HuggingFaceEmbedding.ts

@@ -39,7 +39,7 @@ export class HuggingFaceEmbedding extends BaseEmbedding {
   async getTextEmbedding(text: string): Promise<number[]> {
     const extractor = await this.getExtractor();
     const output = await extractor(text, { pooling: "mean", normalize: true });
-    return output.data;
+    return Array.from(output.data);
   }
 
   async getQueryEmbedding(query: string): Promise<number[]> {

packages/core/src/engines/chat/types.ts

@@ -1,9 +1,10 @@
 import { ChatHistory } from "../../ChatHistory";
-import { NodeWithScore } from "../../Node";
+import { BaseNode, NodeWithScore } from "../../Node";
 import { Response } from "../../Response";
 import { Event } from "../../callbacks/CallbackManager";
 import { ChatMessage } from "../../llm";
 import { MessageContent } from "../../llm/types";
+import { ToolOutput } from "../../tools/types";
 
 /**
  * Represents the base parameters for ChatEngine.
@@ -24,6 +25,10 @@ export interface ChatEngineParamsNonStreaming extends ChatEngineParamsBase {
   stream?: false | null;
 }
 
+export interface ChatEngineAgentParams extends ChatEngineParamsBase {
+  toolChoice?: string | Record<string, any>;
+}
+
 /**
  * A ChatEngine is used to handle back and forth chats between the application and the LLM.
  */
@@ -52,3 +57,32 @@ export interface Context {
 export interface ContextGenerator {
   generate(message: string, parentEvent?: Event): Promise<Context>;
 }
+
+export enum ChatResponseMode {
+  WAIT = "wait",
+  STREAM = "stream",
+}
+
+export class AgentChatResponse {
+  response: string;
+  sources: ToolOutput[];
+  sourceNodes?: BaseNode[];
+
+  constructor(
+    response: string,
+    sources?: ToolOutput[],
+    sourceNodes?: BaseNode[],
+  ) {
+    this.response = response;
+    this.sources = sources || [];
+    this.sourceNodes = sourceNodes || [];
+  }
+
+  protected _getFormattedSources() {
+    throw new Error("Not implemented yet");
+  }
+
+  toString() {
+    return this.response ?? "";
+  }
+}

packages/core/src/engines/query/RetrieverQueryEngine.ts

@@ -0,0 +1,82 @@
+import { NodeWithScore } from "../../Node";
+import { Response } from "../../Response";
+import { BaseRetriever } from "../../Retriever";
+import { ServiceContext } from "../../ServiceContext";
+import { Event } from "../../callbacks/CallbackManager";
+import { randomUUID } from "../../env";
+import { BaseNodePostprocessor } from "../../postprocessors";
+import { BaseSynthesizer, ResponseSynthesizer } from "../../synthesizers";
+import {
+  BaseQueryEngine,
+  QueryEngineParamsNonStreaming,
+  QueryEngineParamsStreaming,
+} from "../../types";
+
+/**
+ * A query engine that uses a retriever to query an index and then synthesizes the response.
+ */
+export class RetrieverQueryEngine implements BaseQueryEngine {
+  retriever: BaseRetriever;
+  responseSynthesizer: BaseSynthesizer;
+  nodePostprocessors: BaseNodePostprocessor[];
+  preFilters?: unknown;
+
+  constructor(
+    retriever: BaseRetriever,
+    responseSynthesizer?: BaseSynthesizer,
+    preFilters?: unknown,
+    nodePostprocessors?: BaseNodePostprocessor[],
+  ) {
+    this.retriever = retriever;
+    const serviceContext: ServiceContext | undefined =
+      this.retriever.getServiceContext();
+    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);
+  }
+
+  query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
+  query(params: QueryEngineParamsNonStreaming): Promise<Response>;
+  async query(
+    params: QueryEngineParamsStreaming | QueryEngineParamsNonStreaming,
+  ): Promise<Response | AsyncIterable<Response>> {
+    const { query, stream } = params;
+    const parentEvent: Event = params.parentEvent || {
+      id: randomUUID(),
+      type: "wrapper",
+      tags: ["final"],
+    };
+    const nodesWithScore = await this.retrieve(query, parentEvent);
+    if (stream) {
+      return this.responseSynthesizer.synthesize({
+        query,
+        nodesWithScore,
+        parentEvent,
+        stream: true,
+      });
+    }
+    return this.responseSynthesizer.synthesize({
+      query,
+      nodesWithScore,
+      parentEvent,
+    });
+  }
+}

packages/core/src/engines/query/RouterQueryEngine.ts

@@ -0,0 +1,181 @@
+import { Response } from "../../Response";
+import {
+  ServiceContext,
+  serviceContextFromDefaults,
+} from "../../ServiceContext";
+import { BaseSelector, LLMSingleSelector } from "../../selectors";
+import { TreeSummarize } from "../../synthesizers";
+import {
+  BaseQueryEngine,
+  QueryBundle,
+  QueryEngineParamsNonStreaming,
+  QueryEngineParamsStreaming,
+} from "../../types";
+
+type RouterQueryEngineTool = {
+  queryEngine: BaseQueryEngine;
+  description: string;
+};
+
+type RouterQueryEngineMetadata = {
+  description: string;
+};
+
+async function combineResponses(
+  summarizer: TreeSummarize,
+  responses: Response[],
+  queryBundle: QueryBundle,
+  verbose: boolean = false,
+): Promise<Response> {
+  if (verbose) {
+    console.log("Combining responses from multiple query engines.");
+  }
+
+  const responseStrs = [];
+  const sourceNodes = [];
+
+  for (const response of responses) {
+    if (response?.sourceNodes) {
+      sourceNodes.push(...response.sourceNodes);
+    }
+
+    responseStrs.push(response.response);
+  }
+
+  const summary = await summarizer.getResponse({
+    query: queryBundle.queryStr,
+    textChunks: responseStrs,
+  });
+
+  return new Response(summary, sourceNodes);
+}
+
+/**
+ * A query engine that uses multiple query engines and selects the best one.
+ */
+export class RouterQueryEngine implements BaseQueryEngine {
+  serviceContext: ServiceContext;
+
+  private selector: BaseSelector;
+  private queryEngines: BaseQueryEngine[];
+  private metadatas: RouterQueryEngineMetadata[];
+  private summarizer: TreeSummarize;
+  private verbose: boolean;
+
+  constructor(init: {
+    selector: BaseSelector;
+    queryEngineTools: RouterQueryEngineTool[];
+    serviceContext?: ServiceContext;
+    summarizer?: TreeSummarize;
+    verbose?: boolean;
+  }) {
+    this.serviceContext = init.serviceContext || serviceContextFromDefaults({});
+    this.selector = init.selector;
+    this.queryEngines = init.queryEngineTools.map((tool) => tool.queryEngine);
+    this.metadatas = init.queryEngineTools.map((tool) => ({
+      description: tool.description,
+    }));
+    this.summarizer = init.summarizer || new TreeSummarize(this.serviceContext);
+    this.verbose = init.verbose ?? false;
+  }
+
+  static fromDefaults(init: {
+    queryEngineTools: RouterQueryEngineTool[];
+    selector?: BaseSelector;
+    serviceContext?: ServiceContext;
+    summarizer?: TreeSummarize;
+    verbose?: boolean;
+  }) {
+    const serviceContext =
+      init.serviceContext ?? serviceContextFromDefaults({});
+
+    return new RouterQueryEngine({
+      selector:
+        init.selector ?? new LLMSingleSelector({ llm: serviceContext.llm }),
+      queryEngineTools: init.queryEngineTools,
+      serviceContext,
+      summarizer: init.summarizer,
+      verbose: init.verbose,
+    });
+  }
+
+  query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
+  query(params: QueryEngineParamsNonStreaming): Promise<Response>;
+  async query(
+    params: QueryEngineParamsStreaming | QueryEngineParamsNonStreaming,
+  ): Promise<Response | AsyncIterable<Response>> {
+    const { query, stream } = params;
+
+    const response = await this.queryRoute({ queryStr: query });
+
+    if (stream) {
+      throw new Error("Streaming is not supported yet.");
+    }
+
+    return response;
+  }
+
+  private async queryRoute(queryBundle: QueryBundle): Promise<Response> {
+    const result = await this.selector.select(this.metadatas, queryBundle);
+
+    if (result.selections.length > 1) {
+      const responses = [];
+      for (let i = 0; i < result.selections.length; i++) {
+        const engineInd = result.selections[i];
+        const logStr = `Selecting query engine ${engineInd}: ${result.selections[i]}.`;
+
+        if (this.verbose) {
+          console.log(logStr + "\n");
+        }
+
+        const selectedQueryEngine = this.queryEngines[engineInd.index];
+        responses.push(
+          await selectedQueryEngine.query({
+            query: queryBundle.queryStr,
+          }),
+        );
+      }
+
+      if (responses.length > 1) {
+        const finalResponse = await combineResponses(
+          this.summarizer,
+          responses,
+          queryBundle,
+          this.verbose,
+        );
+
+        return finalResponse;
+      } else {
+        return responses[0];
+      }
+    } else {
+      let selectedQueryEngine;
+
+      try {
+        selectedQueryEngine = this.queryEngines[result.selections[0].index];
+
+        const logStr = `Selecting query engine ${result.selections[0].index}: ${result.selections[0].reason}`;
+
+        if (this.verbose) {
+          console.log(logStr + "\n");
+        }
+      } catch (e) {
+        throw new Error("Failed to select query engine");
+      }
+
+      if (!selectedQueryEngine) {
+        throw new Error("Selected query engine is null");
+      }
+
+      const finalResponse = await selectedQueryEngine.query({
+        query: queryBundle.queryStr,
+      });
+
+      // add selected result
+      finalResponse.metadata = finalResponse.metadata || {};
+      finalResponse.metadata["selectorResult"] = result;
+
+      return finalResponse;
+    }
+  }
+}

packages/core/src/engines/query/SubQuestionQueryEngine.ts

@@ -1,94 +1,25 @@
-import { NodeWithScore, TextNode } from "./Node";
-import { LLMQuestionGenerator } from "./QuestionGenerator";
-import { Response } from "./Response";
-import { BaseRetriever } from "./Retriever";
-import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
-import { Event } from "./callbacks/CallbackManager";
-import { randomUUID } from "./env";
-import { BaseNodePostprocessor } from "./postprocessors";
+import { NodeWithScore, TextNode } from "../../Node";
+import { LLMQuestionGenerator } from "../../QuestionGenerator";
+import { Response } from "../../Response";
+import {
+  ServiceContext,
+  serviceContextFromDefaults,
+} from "../../ServiceContext";
+import { Event } from "../../callbacks/CallbackManager";
+import { randomUUID } from "../../env";
 import {
   BaseSynthesizer,
   CompactAndRefine,
   ResponseSynthesizer,
-} from "./synthesizers";
+} from "../../synthesizers";
 import {
   BaseQueryEngine,
-  BaseQuestionGenerator,
   QueryEngineParamsNonStreaming,
   QueryEngineParamsStreaming,
   QueryEngineTool,
-  SubQuestion,
   ToolMetadata,
-} from "./types";
-
-/**
- * A query engine that uses a retriever to query an index and then synthesizes the response.
- */
-export class RetrieverQueryEngine implements BaseQueryEngine {
-  retriever: BaseRetriever;
-  responseSynthesizer: BaseSynthesizer;
-  nodePostprocessors: BaseNodePostprocessor[];
-  preFilters?: unknown;
-
-  constructor(
-    retriever: BaseRetriever,
-    responseSynthesizer?: BaseSynthesizer,
-    preFilters?: unknown,
-    nodePostprocessors?: BaseNodePostprocessor[],
-  ) {
-    this.retriever = retriever;
-    const serviceContext: ServiceContext | undefined =
-      this.retriever.getServiceContext();
-    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);
-  }
-
-  query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
-  query(params: QueryEngineParamsNonStreaming): Promise<Response>;
-  async query(
-    params: QueryEngineParamsStreaming | QueryEngineParamsNonStreaming,
-  ): Promise<Response | AsyncIterable<Response>> {
-    const { query, stream } = params;
-    const parentEvent: Event = params.parentEvent || {
-      id: randomUUID(),
-      type: "wrapper",
-      tags: ["final"],
-    };
-    const nodesWithScore = await this.retrieve(query, parentEvent);
-    if (stream) {
-      return this.responseSynthesizer.synthesize({
-        query,
-        nodesWithScore,
-        parentEvent,
-        stream: true,
-      });
-    }
-    return this.responseSynthesizer.synthesize({
-      query,
-      nodesWithScore,
-      parentEvent,
-    });
-  }
-}
+} from "../../types";
+import { BaseQuestionGenerator, SubQuestion } from "./types";
 
 /**
  * SubQuestionQueryEngine decomposes a question into subquestions and then

packages/core/src/engines/query/index.ts

@@ -0,0 +1,3 @@
+export * from "./RetrieverQueryEngine";
+export * from "./RouterQueryEngine";
+export * from "./SubQuestionQueryEngine";

packages/core/src/engines/query/types.ts

@@ -0,0 +1,13 @@
+import { ToolMetadata } from "../../types";
+
+/**
+ * QuestionGenerators generate new questions for the LLM using tools and a user query.
+ */
+export interface BaseQuestionGenerator {
+  generate(tools: ToolMetadata[], query: string): Promise<SubQuestion[]>;
+}
+
+export interface SubQuestion {
+  subQuestion: string;
+  toolName: string;
+}

packages/core/src/env/index.ts

@@ -18,6 +18,20 @@ export function createSHA256(): SHA256 {
   };
 }
 
-export const defaultFS: CompleteFileSystem = fs;
+export const defaultFS: CompleteFileSystem = {
+  writeFile: function (path: string, content: string) {
+    return fs.writeFile(path, content, "utf-8");
+  },
+  readRawFile(path: string): Promise<Buffer> {
+    return fs.readFile(path);
+  },
+  readFile: function (path: string) {
+    return fs.readFile(path, "utf-8");
+  },
+  access: fs.access,
+  mkdir: fs.mkdir,
+  readdir: fs.readdir,
+  stat: fs.stat,
+};
 
 export { EOL, ok, path, randomUUID };

packages/core/src/extractors/prompts.ts

@@ -20,26 +20,33 @@ export interface DefaultNodeTextTemplate {
 export const defaultKeywordExtractorPromptTemplate = ({
   contextStr = "",
   keywords = 5,
-}: DefaultKeywordExtractorPromptTemplate) => `
-  ${contextStr}. Give ${keywords} unique keywords for thiss
-  document. Format as comma separated. Keywords:
+}: DefaultKeywordExtractorPromptTemplate) => `${contextStr} 
+
+Give ${keywords} unique keywords for this document. 
+
+Format as comma separated. Keywords:
 `;
 
 export const defaultTitleExtractorPromptTemplate = (
   { contextStr = "" }: DefaultPromptTemplate = {
     contextStr: "",
   },
-) => `
-  ${contextStr}. Give a title that summarizes all of the unique entities, titles or themes found in the context. Title:
+) => `${contextStr}
+
+Give a title that summarizes all of the unique entities, titles or themes found in the context. 
+
+Title:
 `;
 
 export const defaultTitleCombinePromptTemplate = (
   { contextStr = "" }: DefaultPromptTemplate = {
     contextStr: "",
   },
-) => `
-  ${contextStr}. Based on the above candidate titles and content,s
-  what is the comprehensive title for this document? Title:
+) => `${contextStr} 
+
+Based on the above candidate titles and contents, what is the comprehensive title for this document? 
+
+Title:
 `;
 
 export const defaultQuestionAnswerPromptTemplate = (
@@ -47,23 +54,22 @@ export const defaultQuestionAnswerPromptTemplate = (
     contextStr: "",
     numQuestions: 5,
   },
-) => `
-  ${contextStr}. Given the contextual information,s
-  generate ${numQuestions} questions this context can provides
-  specific answers to which are unlikely to be found elsewhere.
+) => `${contextStr}
 
-  Higher-level summaries of surrounding context may be provideds
-  as well. Try using these summaries to generate better questionss
-  that this context can answer.
+Given the contextual informations, generate ${numQuestions} questions this context can provides specific answers to which are unlikely to be found elsewhere.Higher-level summaries of surrounding context may be provideds as well. 
+
+Try using these summaries to generate better questions that this context can answer.
 `;
 
 export const defaultSummaryExtractorPromptTemplate = (
   { contextStr = "" }: DefaultPromptTemplate = {
     contextStr: "",
   },
-) => `
-  ${contextStr}. Summarize the key topics and entities of the section.s
-  Summary:
+) => `${contextStr}
+
+Summarize the key topics and entities of the sections.
+
+Summary:
 `;
 
 export const defaultNodeTextTemplate = ({

packages/core/src/index.ts

@@ -4,16 +4,17 @@ export * from "./Node";
 export * from "./OutputParser";
 export * from "./Prompt";
 export * from "./PromptHelper";
-export * from "./QueryEngine";
 export * from "./QuestionGenerator";
 export * from "./Response";
 export * from "./Retriever";
 export * from "./ServiceContext";
 export * from "./TextSplitter";
+export * from "./agent";
 export * from "./callbacks/CallbackManager";
 export * from "./constants";
 export * from "./embeddings";
 export * from "./engines/chat";
+export * from "./engines/query";
 export * from "./extractors";
 export * from "./indices";
 export * from "./ingestion";
@@ -30,6 +31,7 @@ export * from "./readers/PDFReader";
 export * from "./readers/SimpleDirectoryReader";
 export * from "./readers/SimpleMongoReader";
 export * from "./readers/base";
+export * from "./selectors";
 export * from "./storage";
 export * from "./synthesizers";
-export type * from "./types";
+export * from "./tools";

packages/core/src/indices/BaseIndex.ts

@@ -2,6 +2,7 @@ import { BaseNode, Document, jsonToNode } from "../Node";
 import { BaseRetriever } from "../Retriever";
 import { ServiceContext } from "../ServiceContext";
 import { randomUUID } from "../env";
+import { runTransformations } from "../ingestion";
 import { StorageContext } from "../storage/StorageContext";
 import { BaseDocumentStore } from "../storage/docStore/types";
 import { BaseIndexStore } from "../storage/indexStore/types";
@@ -188,9 +189,10 @@ export abstract class BaseIndex<T> {
    * @param document
    */
   async insert(document: Document) {
-    const nodes = this.serviceContext.nodeParser.getNodesFromDocuments([
-      document,
-    ]);
+    const nodes = await runTransformations(
+      [document],
+      [this.serviceContext.nodeParser],
+    );
     await this.insertNodes(nodes);
     this.docStore.setDocumentHash(document.id_, document.hash);
   }

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

@@ -1,11 +1,11 @@
 import { BaseNode, Document, MetadataMode } from "../../Node";
 import { defaultKeywordExtractPrompt } from "../../Prompt";
-import { RetrieverQueryEngine } from "../../QueryEngine";
 import { BaseRetriever } from "../../Retriever";
 import {
   ServiceContext,
   serviceContextFromDefaults,
 } from "../../ServiceContext";
+import { RetrieverQueryEngine } from "../../engines/query";
 import { BaseNodePostprocessor } from "../../postprocessors";
 import {
   BaseDocumentStore,

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

@@ -1,11 +1,11 @@
 import _ from "lodash";
 import { BaseNode, Document } from "../../Node";
-import { RetrieverQueryEngine } from "../../QueryEngine";
 import { BaseRetriever } from "../../Retriever";
 import {
   ServiceContext,
   serviceContextFromDefaults,
 } from "../../ServiceContext";
+import { RetrieverQueryEngine } from "../../engines/query";
 import { BaseNodePostprocessor } from "../../postprocessors";
 import {
   BaseDocumentStore,

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

@@ -17,6 +17,12 @@ import { VectorStoreIndex } from "./VectorStoreIndex";
  * VectorIndexRetriever retrieves nodes from a VectorIndex.
  */
 
+export type VectorIndexRetrieverOptions = {
+  index: VectorStoreIndex;
+  similarityTopK?: number;
+  imageSimilarityTopK?: number;
+};
+
 export class VectorIndexRetriever implements BaseRetriever {
   index: VectorStoreIndex;
   similarityTopK: number;
@@ -27,11 +33,7 @@ export class VectorIndexRetriever implements BaseRetriever {
     index,
     similarityTopK,
     imageSimilarityTopK,
-  }: {
-    index: VectorStoreIndex;
-    similarityTopK?: number;
-    imageSimilarityTopK?: number;
-  }) {
+  }: VectorIndexRetrieverOptions) {
     this.index = index;
     this.serviceContext = this.index.serviceContext;
     this.similarityTopK = similarityTopK ?? DEFAULT_SIMILARITY_TOP_K;

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

@@ -6,7 +6,6 @@ import {
   ObjectType,
   splitNodesByType,
 } from "../../Node";
-import { RetrieverQueryEngine } from "../../QueryEngine";
 import { BaseRetriever } from "../../Retriever";
 import {
   ServiceContext,
@@ -17,6 +16,8 @@ import {
   ClipEmbedding,
   MultiModalEmbedding,
 } from "../../embeddings";
+import { RetrieverQueryEngine } from "../../engines/query";
+import { runTransformations } from "../../ingestion";
 import { BaseNodePostprocessor } from "../../postprocessors";
 import {
   BaseIndexStore,
@@ -33,7 +34,10 @@ import {
   IndexDict,
   IndexStructType,
 } from "../BaseIndex";
-import { VectorIndexRetriever } from "./VectorIndexRetriever";
+import {
+  VectorIndexRetriever,
+  VectorIndexRetrieverOptions,
+} from "./VectorIndexRetriever";
 
 interface IndexStructOptions {
   indexStruct?: IndexDict;
@@ -225,8 +229,9 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     if (args.logProgress) {
       console.log("Using node parser on documents...");
     }
-    args.nodes =
-      args.serviceContext.nodeParser.getNodesFromDocuments(documents);
+    args.nodes = await runTransformations(documents, [
+      args.serviceContext.nodeParser,
+    ]);
     if (args.logProgress) {
       console.log("Finished parsing documents.");
     }
@@ -258,7 +263,9 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     return index;
   }
 
-  asRetriever(options?: any): VectorIndexRetriever {
+  asRetriever(
+    options?: Omit<VectorIndexRetrieverOptions, "index">,
+  ): VectorIndexRetriever {
     return new VectorIndexRetriever({ index: this, ...options });
   }
 

packages/core/src/ingestion/IngestionCache.ts

@@ -0,0 +1,68 @@
+import { BaseNode, MetadataMode } from "../Node";
+import { createSHA256 } from "../env";
+import { docToJson, jsonToDoc } from "../storage/docStore/utils";
+import { SimpleKVStore } from "../storage/kvStore/SimpleKVStore";
+import { BaseKVStore } from "../storage/kvStore/types";
+import { TransformComponent } from "./types";
+
+const transformToJSON = (obj: TransformComponent) => {
+  let seen: any[] = [];
+
+  const replacer = (key: string, value: any) => {
+    if (value != null && typeof value == "object") {
+      if (seen.indexOf(value) >= 0) {
+        return;
+      }
+      seen.push(value);
+    }
+    return value;
+  };
+
+  // this is a custom replacer function that will allow us to handle circular references
+  const jsonStr = JSON.stringify(obj, replacer);
+
+  return jsonStr;
+};
+
+export function getTransformationHash(
+  nodes: BaseNode[],
+  transform: TransformComponent,
+) {
+  const nodesStr: string = nodes
+    .map((node) => node.getContent(MetadataMode.ALL))
+    .join("");
+
+  const transformString: string = transformToJSON(transform);
+
+  const hash = createSHA256();
+  hash.update(nodesStr + transformString);
+  return hash.digest();
+}
+
+export class IngestionCache {
+  collection: string = "llama_cache";
+  cache: BaseKVStore;
+  nodesKey = "nodes";
+
+  constructor(collection?: string) {
+    if (collection) {
+      this.collection = collection;
+    }
+    this.cache = new SimpleKVStore();
+  }
+
+  async put(hash: string, nodes: BaseNode[]) {
+    const val = {
+      [this.nodesKey]: nodes.map((node) => docToJson(node)),
+    };
+    await this.cache.put(hash, val, this.collection);
+  }
+
+  async get(hash: string): Promise<BaseNode[] | undefined> {
+    const json = await this.cache.get(hash, this.collection);
+    if (!json || !json[this.nodesKey] || !Array.isArray(json[this.nodesKey])) {
+      return undefined;
+    }
+    return json[this.nodesKey].map((doc: any) => jsonToDoc(doc));
+  }
+}

packages/core/src/ingestion/IngestionPipeline.ts

@@ -1,32 +1,47 @@
 import { BaseNode, Document } from "../Node";
 import { BaseReader } from "../readers/base";
 import { BaseDocumentStore, VectorStore } from "../storage";
+import { IngestionCache, getTransformationHash } from "./IngestionCache";
 import { DocStoreStrategy, createDocStoreStrategy } from "./strategies";
 import { TransformComponent } from "./types";
 
-interface IngestionRunArgs {
+type IngestionRunArgs = {
   documents?: Document[];
   nodes?: BaseNode[];
+};
+
+type TransformRunArgs = {
   inPlace?: boolean;
-}
+  cache?: IngestionCache;
+};
 
 export async function runTransformations(
   nodesToRun: BaseNode[],
   transformations: TransformComponent[],
   transformOptions: any = {},
-  { inPlace = true }: IngestionRunArgs,
+  { inPlace = true, cache }: TransformRunArgs = {},
 ): Promise<BaseNode[]> {
   let nodes = nodesToRun;
   if (!inPlace) {
     nodes = [...nodesToRun];
   }
   for (const transform of transformations) {
-    nodes = await transform.transform(nodes, transformOptions);
+    if (cache) {
+      const hash = getTransformationHash(nodes, transform);
+      const cachedNodes = await cache.get(hash);
+      if (cachedNodes) {
+        nodes = cachedNodes;
+      } else {
+        nodes = await transform.transform(nodes, transformOptions);
+        await cache.put(hash, nodes);
+      }
+    } else {
+      nodes = await transform.transform(nodes, transformOptions);
+    }
   }
   return nodes;
 }
 
-// TODO: add caching, add concurrency
 export class IngestionPipeline {
   transformations: TransformComponent[] = [];
   documents?: Document[];
@@ -34,7 +49,8 @@ export class IngestionPipeline {
   vectorStore?: VectorStore;
   docStore?: BaseDocumentStore;
   docStoreStrategy: DocStoreStrategy = DocStoreStrategy.UPSERTS;
-  disableCache: boolean = true;
+  cache?: IngestionCache;
+  disableCache: boolean = false;
 
   private _docStoreStrategy?: TransformComponent;
 
@@ -45,6 +61,9 @@ export class IngestionPipeline {
       this.docStore,
       this.vectorStore,
     );
+    if (!this.disableCache) {
+      this.cache = new IngestionCache();
+    }
   }
 
   async prepareInput(
@@ -68,9 +87,10 @@ export class IngestionPipeline {
   }
 
   async run(
-    args: IngestionRunArgs = {},
+    args: IngestionRunArgs & TransformRunArgs = {},
     transformOptions?: any,
   ): Promise<BaseNode[]> {
+    args.cache = args.cache ?? this.cache;
     const inputNodes = await this.prepareInput(args.documents, args.nodes);
     let nodesToRun;
     if (this._docStoreStrategy) {

packages/core/src/llm/LLM.ts

@@ -77,7 +77,14 @@ export class OpenAI extends BaseLLM {
   maxTokens?: number;
   additionalChatOptions?: Omit<
     Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
-    "max_tokens" | "messages" | "model" | "temperature" | "top_p" | "stream"
+    | "max_tokens"
+    | "messages"
+    | "model"
+    | "temperature"
+    | "top_p"
+    | "stream"
+    | "tools"
+    | "toolChoice"
   >;
 
   // OpenAI session params
@@ -179,7 +186,7 @@ export class OpenAI extends BaseLLM {
 
   mapMessageType(
     messageType: MessageType,
-  ): "user" | "assistant" | "system" | "function" {
+  ): "user" | "assistant" | "system" | "function" | "tool" {
     switch (messageType) {
       case "user":
         return "user";
@@ -189,11 +196,30 @@ export class OpenAI extends BaseLLM {
         return "system";
       case "function":
         return "function";
+      case "tool":
+        return "tool";
       default:
         return "user";
     }
   }
 
+  toOpenAIMessage(messages: ChatMessage[]) {
+    return messages.map((message) => {
+      const additionalKwargs = message.additionalKwargs ?? {};
+
+      if (message.additionalKwargs?.toolCalls) {
+        additionalKwargs.tool_calls = message.additionalKwargs.toolCalls;
+        delete additionalKwargs.toolCalls;
+      }
+
+      return {
+        role: this.mapMessageType(message.role),
+        content: message.content,
+        ...additionalKwargs,
+      };
+    });
+  }
+
   chat(
     params: LLMChatParamsStreaming,
   ): Promise<AsyncIterable<ChatResponseChunk>>;
@@ -201,25 +227,24 @@ export class OpenAI extends BaseLLM {
   async chat(
     params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
   ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, parentEvent, stream } = params;
-    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
+    const { messages, parentEvent, stream, tools, toolChoice } = params;
+
+    let baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
       model: this.model,
       temperature: this.temperature,
       max_tokens: this.maxTokens,
-      messages: messages.map(
-        (message) =>
-          ({
-            role: this.mapMessageType(message.role),
-            content: message.content,
-          }) as ChatCompletionMessageParam,
-      ),
+      tools: tools,
+      tool_choice: toolChoice,
+      messages: this.toOpenAIMessage(messages) as ChatCompletionMessageParam[],
       top_p: this.topP,
       ...this.additionalChatOptions,
     };
+
     // Streaming
     if (stream) {
       return this.streamChat(params);
     }
+
     // Non-streaming
     const response = await this.session.openai.chat.completions.create({
       ...baseRequestParams,
@@ -227,8 +252,19 @@ export class OpenAI extends BaseLLM {
     });
 
     const content = response.choices[0].message?.content ?? "";
+
+    const kwargsOutput: Record<string, any> = {};
+
+    if (response.choices[0].message?.tool_calls) {
+      kwargsOutput.toolCalls = response.choices[0].message.tool_calls;
+    }
+
     return {
-      message: { content, role: response.choices[0].message.role },
+      message: {
+        content,
+        role: response.choices[0].message.role,
+        additionalKwargs: kwargsOutput,
+      },
     };
   }
 

packages/core/src/llm/types.ts

@@ -39,17 +39,20 @@ export type MessageType =
   | "system"
   | "generic"
   | "function"
-  | "memory";
+  | "memory"
+  | "tool";
 
 export interface ChatMessage {
   // TODO: use MessageContent
   content: any;
   role: MessageType;
+  additionalKwargs?: Record<string, any>;
 }
 
 export interface ChatResponse {
   message: ChatMessage;
   raw?: Record<string, any>;
+  additionalKwargs?: Record<string, any>;
 }
 
 export interface ChatResponseChunk {
@@ -74,6 +77,9 @@ export interface LLMChatParamsBase {
   messages: ChatMessage[];
   parentEvent?: Event;
   extraParams?: Record<string, any>;
+  tools?: any;
+  toolChoice?: any;
+  additionalKwargs?: Record<string, any>;
 }
 
 export interface LLMChatParamsStreaming extends LLMChatParamsBase {

packages/core/src/memory/ChatMemoryBuffer.ts

@@ -0,0 +1,119 @@
+import { ChatMessage } from "../llm";
+import { SimpleChatStore } from "../storage/chatStore/SimpleChatStore";
+import { BaseChatStore } from "../storage/chatStore/types";
+import { BaseMemory } from "./types";
+
+type ChatMemoryBufferParams = {
+  tokenLimit?: number;
+  chatStore?: BaseChatStore;
+  chatStoreKey?: string;
+  chatHistory?: ChatMessage[];
+};
+
+/**
+ * Chat memory buffer.
+ */
+export class ChatMemoryBuffer implements BaseMemory {
+  tokenLimit: number;
+
+  chatStore: BaseChatStore;
+  chatStoreKey: string;
+
+  /**
+   * Initialize.
+   */
+  constructor(init?: Partial<ChatMemoryBufferParams>) {
+    this.tokenLimit = init?.tokenLimit ?? 3000;
+    this.chatStore = init?.chatStore ?? new SimpleChatStore();
+    this.chatStoreKey = init?.chatStoreKey ?? "chat_history";
+
+    if (init?.chatHistory) {
+      this.chatStore.setMessages(this.chatStoreKey, init.chatHistory);
+    }
+  }
+
+  /**
+    Get chat history.
+    @param initialTokenCount: number of tokens to start with
+  */
+  get(initialTokenCount: number = 0): ChatMessage[] {
+    const chatHistory = this.getAll();
+
+    if (initialTokenCount > this.tokenLimit) {
+      throw new Error("Initial token count exceeds token limit");
+    }
+
+    let messageCount = chatHistory.length;
+    let tokenCount =
+      this._tokenCountForMessageCount(messageCount) + initialTokenCount;
+
+    while (tokenCount > this.tokenLimit && messageCount > 1) {
+      messageCount -= 1;
+      if (chatHistory[-messageCount].role === "assistant") {
+        // we cannot have an assistant message at the start of the chat history
+        // if after removal of the first, we have an assistant message,
+        // we need to remove the assistant message too
+        messageCount -= 1;
+      }
+
+      tokenCount =
+        this._tokenCountForMessageCount(messageCount) + initialTokenCount;
+    }
+
+    // catch one message longer than token limit
+    if (tokenCount > this.tokenLimit || messageCount <= 0) {
+      return [];
+    }
+
+    return chatHistory.slice(-messageCount);
+  }
+
+  /**
+   * Get all chat history.
+   * @returns {ChatMessage[]} chat history
+   */
+  getAll(): ChatMessage[] {
+    return this.chatStore.getMessages(this.chatStoreKey);
+  }
+
+  /**
+   * Put chat history.
+   * @param message
+   */
+  put(message: ChatMessage): void {
+    this.chatStore.addMessage(this.chatStoreKey, message);
+  }
+
+  /**
+   * Set chat history.
+   * @param messages
+   */
+  set(messages: ChatMessage[]): void {
+    this.chatStore.setMessages(this.chatStoreKey, messages);
+  }
+
+  /**
+   * Reset chat history.
+   */
+  reset(): void {
+    this.chatStore.deleteMessages(this.chatStoreKey);
+  }
+
+  /**
+   * Get token count for message count.
+   * @param messageCount
+   * @returns {number} token count
+   */
+  private _tokenCountForMessageCount(messageCount: number): number {
+    if (messageCount <= 0) {
+      return 0;
+    }
+
+    const chatHistory = this.getAll();
+    const msgStr = chatHistory
+      .slice(-messageCount)
+      .map((m) => m.content)
+      .join(" ");
+    return msgStr.split(" ").length;
+  }
+}

packages/core/src/memory/types.ts

@@ -0,0 +1,24 @@
+import { ChatMessage } from "../llm";
+
+export interface BaseMemory {
+  /* 
+   Get chat history.
+  */
+  get(...args: any): ChatMessage[];
+  /*
+    Get all chat history.
+  */
+  getAll(): ChatMessage[];
+  /*
+    Put chat history.
+  */
+  put(message: ChatMessage): void;
+  /*
+    Set chat history.
+  */
+  set(messages: ChatMessage[]): void;
+  /*
+    Reset chat history.
+  */
+  reset(): void;
+}

packages/core/src/objects/base.ts

@@ -0,0 +1,80 @@
+import { BaseNode, TextNode } from "../Node";
+import { BaseRetriever } from "../Retriever";
+
+// Assuming that necessary interfaces and classes (like OT, TextNode, BaseNode, etc.) are defined elsewhere
+// Import statements (e.g., for TextNode, BaseNode) should be added based on your project's structure
+
+export abstract class BaseObjectNodeMapping<OT> {
+  // TypeScript doesn't support Python's classmethod directly, but we can use static methods as an alternative
+  abstract fromObjects<OT>(
+    objs: OT[],
+    ...args: any[]
+  ): BaseObjectNodeMapping<OT>;
+
+  // Abstract methods in TypeScript
+  abstract objNodeMapping(): Record<any, any>;
+  abstract toNode(obj: OT): TextNode;
+
+  // Concrete methods can be defined as usual
+  validateObject(obj: OT): void {}
+
+  // Implementing the add object logic
+  addObj(obj: OT): void {
+    this.validateObject(obj);
+    this._addObj(obj);
+  }
+
+  // Abstract method for internal add object logic
+  protected abstract _addObj(obj: OT): void;
+
+  // Implementing toNodes method
+  toNodes(objs: OT[]): TextNode[] {
+    return objs.map((obj) => this.toNode(obj));
+  }
+
+  // Abstract method for internal from node logic
+  protected abstract _fromNode(node: BaseNode): OT;
+
+  // Implementing fromNode method
+  fromNode(node: BaseNode): OT {
+    const obj = this._fromNode(node);
+    this.validateObject(obj);
+    return obj;
+  }
+
+  // Abstract methods for persistence
+  abstract persist(persistDir: string, objNodeMappingFilename: string): void;
+}
+
+// You will need to implement specific subclasses of BaseObjectNodeMapping as per your project requirements.
+
+type QueryType = string;
+
+export class ObjectRetriever<OT> {
+  private _retriever: BaseRetriever;
+  private _objectNodeMapping: BaseObjectNodeMapping<OT>;
+
+  constructor(
+    retriever: BaseRetriever,
+    objectNodeMapping: BaseObjectNodeMapping<OT>,
+  ) {
+    this._retriever = retriever;
+    this._objectNodeMapping = objectNodeMapping;
+  }
+
+  // In TypeScript, getters are defined like this.
+  get retriever(): BaseRetriever {
+    return this._retriever;
+  }
+
+  // Translating the retrieve method
+  async retrieve(strOrQueryBundle: QueryType): Promise<OT[]> {
+    const nodes = await this._retriever.retrieve(strOrQueryBundle);
+    return nodes.map((node) => this._objectNodeMapping.fromNode(node.node));
+  }
+
+  // // Translating the _asQueryComponent method
+  // public asQueryComponent(kwargs: any): any {
+  //     return new ObjectRetrieverComponent(this);
+  // }
+}

packages/core/src/outputParsers/selectors.ts

@@ -0,0 +1,52 @@
+import { parseJsonMarkdown } from "../OutputParser";
+import { BaseOutputParser, StructuredOutput } from "../types";
+
+export type Answer = {
+  choice: number;
+  reason: string;
+};
+
+const formatStr = `The output should be ONLY JSON formatted as a JSON instance.
+
+Here is an example:
+[
+    {
+        choice: 1,
+        reason: "<insert reason for choice>"
+    },
+    ...
+]
+`;
+
+/*
+ * An OutputParser is used to extract structured data from the raw output of the LLM.
+ */
+export class SelectionOutputParser
+  implements BaseOutputParser<StructuredOutput<Answer[]>>
+{
+  /**
+   *
+   * @param output
+   */
+  parse(output: string): StructuredOutput<Answer[]> {
+    let parsed;
+
+    try {
+      parsed = parseJsonMarkdown(output);
+    } catch (e) {
+      try {
+        parsed = JSON.parse(output);
+      } catch (e) {
+        throw new Error(
+          `Got invalid JSON object. Error: ${e}. Got JSON string: ${output}`,
+        );
+      }
+    }
+
+    return { rawOutput: output, parsedOutput: parsed };
+  }
+
+  format(output: string): string {
+    return output + "\n\n" + formatStr;
+  }
+}

packages/core/src/readers/CSVReader.ts

@@ -43,7 +43,7 @@ export class PapaCSVReader implements BaseReader {
     file: string,
     fs: GenericFileSystem = defaultFS,
   ): Promise<Document[]> {
-    const fileContent: string = await fs.readFile(file, "utf-8");
+    const fileContent = await fs.readFile(file);
     const result = Papa.parse(fileContent, this.papaConfig);
     const textList = result.data.map((row: any) => {
       // Compatible with header row mode

packages/core/src/readers/DocxReader.ts

@@ -10,7 +10,7 @@ export class DocxReader implements BaseReader {
     file: string,
     fs: GenericFileSystem = defaultFS,
   ): Promise<Document[]> {
-    const dataBuffer = (await fs.readFile(file)) as any;
+    const dataBuffer = await fs.readRawFile(file);
     const { value } = await mammoth.extractRawText({ buffer: dataBuffer });
     return [new Document({ text: value, id_: file })];
   }

packages/core/src/readers/HTMLReader.ts

@@ -22,7 +22,7 @@ export class HTMLReader implements BaseReader {
     file: string,
     fs: GenericFileSystem = defaultFS,
   ): Promise<Document[]> {
-    const dataBuffer = await fs.readFile(file, "utf-8");
+    const dataBuffer = await fs.readFile(file);
     const htmlOptions = this.getOptions();
     const content = await this.parseContent(dataBuffer, htmlOptions);
     return [new Document({ text: content, id_: file })];

packages/core/src/readers/MarkdownReader.ts

@@ -92,7 +92,7 @@ export class MarkdownReader implements BaseReader {
     file: string,
     fs: GenericFileSystem = defaultFS,
   ): Promise<Document[]> {
-    const content = await fs.readFile(file, { encoding: "utf-8" });
+    const content = await fs.readFile(file);
     const tups = this.parseTups(content);
     const results: Document[] = [];
     for (const [header, value] of tups) {

packages/core/src/readers/PDFReader.ts

@@ -1,5 +1,5 @@
 import { Document } from "../Node";
-import { defaultFS } from "../env";
+import { createSHA256, defaultFS } from "../env";
 import { GenericFileSystem } from "../storage/FileSystem";
 import { BaseReader } from "./base";
 
@@ -11,78 +11,31 @@ export class PDFReader implements BaseReader {
     file: string,
     fs: GenericFileSystem = defaultFS,
   ): Promise<Document[]> {
-    const content = (await fs.readFile(file)) as any;
-    if (!(content instanceof Buffer)) {
-      console.warn(`PDF File ${file} can only be loaded using the Node FS`);
-      return [];
-    }
-    const data = new Uint8Array(
-      content.buffer,
-      content.byteOffset,
-      content.byteLength,
-    );
-    const pdf = await readPDF(data);
-    return [new Document({ text: pdf.text, id_: file })];
+    const content = await fs.readRawFile(file);
+    const text = await readPDF(content);
+    return text.map((text) => {
+      const sha256 = createSHA256();
+      sha256.update(text);
+      return new Document({
+        text,
+        id_: sha256.digest(),
+      });
+    });
   }
 }
 
-// NOTE: the following code is taken from https://www.npmjs.com/package/pdf-parse and modified
-async function readPage(pageData: any) {
-  //check documents https://mozilla.github.io/pdf.js/
-  const textContent = await pageData.getTextContent({
-    includeMarkedContent: false,
+async function readPDF(data: Buffer): Promise<string[]> {
+  const parser = await import("pdf2json").then(
+    ({ default: Pdfparser }) => new Pdfparser(null, 1),
+  );
+  const text = await new Promise<string>((resolve, reject) => {
+    parser.on("pdfParser_dataError", (error) => {
+      reject(error);
+    });
+    parser.on("pdfParser_dataReady", () => {
+      resolve((parser as any).getRawTextContent() as string);
+    });
+    parser.parseBuffer(data);
   });
-
-  let lastY = null,
-    text = "";
-  //https://github.com/mozilla/pdf.js/issues/8963
-  //https://github.com/mozilla/pdf.js/issues/2140
-  //https://gist.github.com/hubgit/600ec0c224481e910d2a0f883a7b98e3
-  //https://gist.github.com/hubgit/600ec0c224481e910d2a0f883a7b98e3
-  for (const item of textContent.items) {
-    if (lastY == item.transform[5] || !lastY) {
-      text += item.str;
-    } else {
-      text += "\n" + item.str;
-    }
-    lastY = item.transform[5];
-  }
-  return text;
-}
-
-const PDF_DEFAULT_OPTIONS = {
-  max: 0,
-};
-
-async function readPDF(data: Uint8Array, options = PDF_DEFAULT_OPTIONS) {
-  const { getDocument, version } = await import("pdfjs-dist");
-
-  const doc = await getDocument({ data }).promise;
-  const metaData = await doc.getMetadata().catch(() => null);
-  const counter =
-    options.max === 0 ? doc.numPages : Math.max(options.max, doc.numPages);
-
-  let text = "";
-
-  for (let i = 1; i <= counter; i++) {
-    try {
-      const pageData = await doc.getPage(i);
-      const pageText = await readPage(pageData);
-
-      text += `\n\n${pageText}`;
-    } catch (err) {
-      console.log(err);
-    }
-  }
-
-  await doc.destroy();
-
-  return {
-    numpages: doc.numPages,
-    numrender: counter,
-    info: metaData?.info,
-    metadata: metaData?.metadata,
-    text,
-    version,
-  };
+  return text.split(/----------------Page \(\d+\) Break----------------/g);
 }

packages/core/src/readers/SimpleDirectoryReader.ts

@@ -30,7 +30,7 @@ export class TextFileReader implements BaseReader {
     file: string,
     fs: CompleteFileSystem = defaultFS,
   ): Promise<Document[]> {
-    const dataBuffer = await fs.readFile(file, "utf-8");
+    const dataBuffer = await fs.readFile(file);
     return [new Document({ text: dataBuffer, id_: file })];
   }
 }
@@ -49,7 +49,7 @@ export const FILE_EXT_TO_READER: Record<string, BaseReader> = {
   gif: new ImageReader(),
 };
 
-export type SimpleDirectoryReaderLoadDataProps = {
+export type SimpleDirectoryReaderLoadDataParams = {
   directoryPath: string;
   fs?: CompleteFileSystem;
   defaultReader?: BaseReader | null;
@@ -64,12 +64,20 @@ export type SimpleDirectoryReaderLoadDataProps = {
 export class SimpleDirectoryReader implements BaseReader {
   constructor(private observer?: ReaderCallback) {}
 
-  async loadData({
-    directoryPath,
-    fs = defaultFS,
-    defaultReader = new TextFileReader(),
-    fileExtToReader = FILE_EXT_TO_READER,
-  }: SimpleDirectoryReaderLoadDataProps): Promise<Document[]> {
+  async loadData(
+    params: SimpleDirectoryReaderLoadDataParams | string,
+  ): Promise<Document[]> {
+    if (typeof params === "string") {
+      params = { directoryPath: params };
+    }
+
+    const {
+      directoryPath,
+      fs = defaultFS,
+      defaultReader = new TextFileReader(),
+      fileExtToReader = FILE_EXT_TO_READER,
+    } = params;
+
     // Observer can decide to skip the directory
     if (
       !this.doObserverCheck("directory", directoryPath, ReaderStatus.STARTED)

packages/core/src/selectors/base.ts

@@ -0,0 +1,45 @@
+import { QueryBundle, ToolMetadataOnlyDescription } from "../types";
+
+export interface SingleSelection {
+  index: number;
+  reason: string;
+}
+
+export type SelectorResult = {
+  selections: SingleSelection[];
+};
+
+type QueryType = string | QueryBundle;
+
+function wrapChoice(
+  choice: string | ToolMetadataOnlyDescription,
+): ToolMetadataOnlyDescription {
+  if (typeof choice === "string") {
+    return { description: choice };
+  } else {
+    return choice;
+  }
+}
+
+function wrapQuery(query: QueryType): QueryBundle {
+  if (typeof query === "string") {
+    return { queryStr: query };
+  }
+
+  return query;
+}
+
+type MetadataType = string | ToolMetadataOnlyDescription;
+
+export abstract class BaseSelector {
+  async select(choices: MetadataType[], query: QueryType) {
+    const metadatas = choices.map((choice) => wrapChoice(choice));
+    const queryBundle = wrapQuery(query);
+    return await this._select(metadatas, queryBundle);
+  }
+
+  abstract _select(
+    choices: ToolMetadataOnlyDescription[],
+    query: QueryBundle,
+  ): Promise<SelectorResult>;
+}

packages/core/src/selectors/index.ts

@@ -0,0 +1,3 @@
+export * from "./base";
+export * from "./llmSelectors";
+export * from "./utils";

packages/core/src/selectors/llmSelectors.ts

@@ -0,0 +1,168 @@
+import { DefaultPromptTemplate } from "../extractors/prompts";
+import { LLM } from "../llm";
+import { Answer, SelectionOutputParser } from "../outputParsers/selectors";
+import {
+  BaseOutputParser,
+  QueryBundle,
+  StructuredOutput,
+  ToolMetadataOnlyDescription,
+} from "../types";
+import { BaseSelector, SelectorResult } from "./base";
+import { defaultSingleSelectPrompt } from "./prompts";
+
+function buildChoicesText(choices: ToolMetadataOnlyDescription[]): string {
+  const texts: string[] = [];
+  for (const [ind, choice] of choices.entries()) {
+    let text = choice.description.split("\n").join(" ");
+    text = `(${ind + 1}) ${text}`; // to one indexing
+    texts.push(text);
+  }
+  return texts.join("");
+}
+
+function _structuredOutputToSelectorResult(
+  output: StructuredOutput<Answer[]>,
+): SelectorResult {
+  const structuredOutput = output;
+  const answers = structuredOutput.parsedOutput;
+
+  // adjust for zero indexing
+  const selections = answers.map((answer: any) => {
+    return { index: answer.choice - 1, reason: answer.reason };
+  });
+
+  return { selections };
+}
+
+type LLMPredictorType = LLM;
+
+/**
+ * A selector that uses the LLM to select a single or multiple choices from a list of choices.
+ */
+export class LLMMultiSelector extends BaseSelector {
+  _llm: LLMPredictorType;
+  _prompt: DefaultPromptTemplate | undefined;
+  _maxOutputs: number | null;
+  _outputParser: BaseOutputParser<any> | null;
+
+  constructor(init: {
+    llm: LLMPredictorType;
+    prompt?: DefaultPromptTemplate;
+    maxOutputs?: number;
+    outputParser?: BaseOutputParser<any>;
+  }) {
+    super();
+    this._llm = init.llm;
+    this._prompt = init.prompt;
+    this._maxOutputs = init.maxOutputs ?? null;
+
+    this._outputParser = init.outputParser ?? new SelectionOutputParser();
+  }
+
+  _getPrompts(): Record<string, any> {
+    return { prompt: this._prompt };
+  }
+
+  _updatePrompts(prompts: Record<string, any>): void {
+    if ("prompt" in prompts) {
+      this._prompt = prompts.prompt;
+    }
+  }
+
+  /**
+   * Selects a single choice from a list of choices.
+   * @param choices
+   * @param query
+   */
+  async _select(
+    choices: ToolMetadataOnlyDescription[],
+    query: QueryBundle,
+  ): Promise<SelectorResult> {
+    const choicesText = buildChoicesText(choices);
+
+    const prompt =
+      this._prompt?.contextStr ??
+      defaultSingleSelectPrompt(
+        choicesText.length,
+        choicesText,
+        query.queryStr,
+      );
+    const formattedPrompt = this._outputParser?.format(prompt);
+
+    const prediction = await this._llm.complete({
+      prompt: formattedPrompt,
+    });
+
+    const parsed = this._outputParser?.parse(prediction.text);
+
+    return _structuredOutputToSelectorResult(parsed);
+  }
+
+  asQueryComponent(): unknown {
+    throw new Error("Method not implemented.");
+  }
+}
+
+/**
+ * A selector that uses the LLM to select a single choice from a list of choices.
+ */
+export class LLMSingleSelector extends BaseSelector {
+  _llm: LLMPredictorType;
+  _prompt: DefaultPromptTemplate | undefined;
+  _outputParser: BaseOutputParser<any> | null;
+
+  constructor(init: {
+    llm: LLMPredictorType;
+    prompt?: DefaultPromptTemplate;
+    outputParser?: BaseOutputParser<any>;
+  }) {
+    super();
+    this._llm = init.llm;
+    this._prompt = init.prompt;
+    this._outputParser = init.outputParser ?? new SelectionOutputParser();
+  }
+
+  _getPrompts(): Record<string, any> {
+    return { prompt: this._prompt };
+  }
+
+  _updatePrompts(prompts: Record<string, any>): void {
+    if ("prompt" in prompts) {
+      this._prompt = prompts.prompt;
+    }
+  }
+
+  /**
+   * Selects a single choice from a list of choices.
+   * @param choices
+   * @param query
+   */
+  async _select(
+    choices: ToolMetadataOnlyDescription[],
+    query: QueryBundle,
+  ): Promise<SelectorResult> {
+    const choicesText = buildChoicesText(choices);
+
+    const prompt =
+      this._prompt?.contextStr ??
+      defaultSingleSelectPrompt(
+        choicesText.length,
+        choicesText,
+        query.queryStr,
+      );
+
+    const formattedPrompt = this._outputParser?.format(prompt);
+
+    const prediction = await this._llm.complete({
+      prompt: formattedPrompt,
+    });
+
+    const parsed = this._outputParser?.parse(prediction.text);
+
+    return _structuredOutputToSelectorResult(parsed);
+  }
+
+  asQueryComponent(): unknown {
+    throw new Error("Method not implemented.");
+  }
+}

packages/core/src/selectors/prompts.ts

@@ -0,0 +1,30 @@
+export const defaultSingleSelectPrompt = (
+  numChoices: number,
+  contextList: string,
+  queryStr: string,
+): string => {
+  return `Some choices are given below. It is provided in a numbered list (1 to ${numChoices}), where each item in the list corresponds to a summary.
+---------------------
+${contextList}
+---------------------
+Using only the choices above and not prior knowledge, return the choice that is most relevant to the question: '${queryStr}'
+`;
+};
+
+export type SingleSelectPrompt = typeof defaultSingleSelectPrompt;
+
+export const defaultMultiSelectPrompt = (
+  numChoices: number,
+  contextList: string,
+  queryStr: string,
+  maxOutputs: number,
+) => {
+  return `Some choices are given below. It is provided in a numbered list (1 to ${numChoices}), where each item in the list corresponds to a summary.
+---------------------
+${contextList}
+---------------------
+Using only the choices above and not prior knowledge, return the top choices (no more than ${maxOutputs}, but only select what is needed) that are most relevant to the question: '${queryStr}'
+`;
+};
+
+export type MultiSelectPrompt = typeof defaultMultiSelectPrompt;

packages/core/src/selectors/utils.ts

@@ -0,0 +1,24 @@
+import { ServiceContext } from "../ServiceContext";
+import { BaseSelector } from "./base";
+import { LLMMultiSelector, LLMSingleSelector } from "./llmSelectors";
+
+export const getSelectorFromContext = (
+  serviceContext: ServiceContext,
+  isMulti: boolean = false,
+): BaseSelector => {
+  let selector: BaseSelector | null = null;
+
+  const llm = serviceContext.llm;
+
+  if (isMulti) {
+    selector = new LLMMultiSelector({ llm });
+  } else {
+    selector = new LLMSingleSelector({ llm });
+  }
+
+  if (selector === null) {
+    throw new Error("Selector is null");
+  }
+
+  return selector;
+};

packages/core/src/storage/FileSystem.ts

@@ -1,5 +1,4 @@
 import _ from "lodash";
-import type nodeFS from "node:fs/promises";
 
 /**
  * A filesystem interface that is meant to be compatible with
@@ -8,20 +7,23 @@ import type nodeFS from "node:fs/promises";
  * browsers.
  */
 export type GenericFileSystem = {
-  writeFile(
-    path: string,
-    content: string,
-    options?: Parameters<typeof nodeFS.writeFile>[2],
-  ): Promise<void>;
-  readFile(
-    path: string,
-    options?: Parameters<typeof nodeFS.readFile>[1],
-  ): Promise<string>;
+  writeFile(path: string, content: string): Promise<void>;
+  /**
+   * Reads a file and returns its content as a raw buffer.
+   */
+  readRawFile(path: string): Promise<Buffer>;
+  /**
+   * Reads a file and returns its content as an utf-8 string.
+   */
+  readFile(path: string): Promise<string>;
   access(path: string): Promise<void>;
   mkdir(
     path: string,
-    options?: Parameters<typeof nodeFS.mkdir>[1],
-  ): Promise<void>;
+    options: {
+      recursive: boolean;
+    },
+  ): Promise<string | undefined>;
+  mkdir(path: string): Promise<void>;
 };
 
 export type WalkableFileSystem = {
@@ -45,7 +47,7 @@ export class InMemoryFileSystem implements CompleteFileSystem {
     this.files[path] = _.cloneDeep(content);
   }
 
-  async readFile(path: string, options?: unknown): Promise<string> {
+  async readFile(path: string): Promise<string> {
     if (!(path in this.files)) {
       throw new Error(`File ${path} does not exist`);
     }
@@ -58,8 +60,9 @@ export class InMemoryFileSystem implements CompleteFileSystem {
     }
   }
 
-  async mkdir(path: string, options?: unknown): Promise<void> {
+  async mkdir(path: string) {
     this.files[path] = _.get(this.files, path, null);
+    return undefined;
   }
 
   async readdir(path: string): Promise<string[]> {
@@ -69,6 +72,10 @@ export class InMemoryFileSystem implements CompleteFileSystem {
   async stat(path: string): Promise<any> {
     throw new Error("Not implemented");
   }
+
+  async readRawFile(path: string): Promise<Buffer> {
+    throw new Error("Not implemented");
+  }
 }
 
 // FS utility functions

packages/core/src/storage/chatStore/SimpleChatStore.ts

@@ -0,0 +1,92 @@
+import { ChatMessage } from "../../llm";
+import { BaseChatStore } from "./types";
+
+/**
+ * Simple chat store.
+ */
+export class SimpleChatStore implements BaseChatStore {
+  store: { [key: string]: ChatMessage[] } = {};
+
+  /**
+   * Set messages.
+   * @param key: key
+   * @param messages: messages
+   * @returns: void
+   */
+  public setMessages(key: string, messages: ChatMessage[]): void {
+    this.store[key] = messages;
+  }
+
+  /**
+   * Get messages.
+   * @param key: key
+   * @returns: messages
+   */
+  public getMessages(key: string): ChatMessage[] {
+    return this.store[key] || [];
+  }
+
+  /**
+   * Add message.
+   * @param key: key
+   * @param message: message
+   * @returns: void
+   */
+  public addMessage(key: string, message: ChatMessage): void {
+    this.store[key] = this.store[key] || [];
+    this.store[key].push(message);
+  }
+
+  /**
+   * Delete messages.
+   * @param key: key
+   * @returns: messages
+   */
+  public deleteMessages(key: string): ChatMessage[] | null {
+    if (!(key in this.store)) {
+      return null;
+    }
+    const messages = this.store[key];
+    delete this.store[key];
+    return messages;
+  }
+
+  /**
+   * Delete message.
+   * @param key: key
+   * @param idx: idx
+   * @returns: message
+   */
+  public deleteMessage(key: string, idx: number): ChatMessage | null {
+    if (!(key in this.store)) {
+      return null;
+    }
+    if (idx >= this.store[key].length) {
+      return null;
+    }
+    return this.store[key].splice(idx, 1)[0];
+  }
+
+  /**
+   * Delete last message.
+   * @param key: key
+   * @returns: message
+   */
+  public deleteLastMessage(key: string): ChatMessage | null {
+    if (!(key in this.store)) {
+      return null;
+    }
+
+    const lastMessage = this.store[key].pop();
+
+    return lastMessage || null;
+  }
+
+  /**
+   * Get keys.
+   * @returns: keys
+   */
+  public getKeys(): string[] {
+    return Object.keys(this.store);
+  }
+}

packages/core/src/storage/chatStore/types.ts

@@ -0,0 +1,11 @@
+import { ChatMessage } from "../../llm";
+
+export interface BaseChatStore {
+  setMessages(key: string, messages: ChatMessage[]): void;
+  getMessages(key: string): ChatMessage[];
+  addMessage(key: string, message: ChatMessage): void;
+  deleteMessages(key: string): ChatMessage[] | null;
+  deleteMessage(key: string, idx: number): ChatMessage | null;
+  deleteLastMessage(key: string): ChatMessage | null;
+  getKeys(): string[];
+}

packages/core/src/storage/index.ts

@@ -1,5 +1,7 @@
 export * from "./FileSystem";
 export * from "./StorageContext";
+export { SimpleChatStore } from "./chatStore/SimpleChatStore";
+export * from "./chatStore/types";
 export * from "./constants";
 export { SimpleDocumentStore } from "./docStore/SimpleDocumentStore";
 export * from "./docStore/types";

packages/core/src/storage/vectorStore/PineconeVectorStore.ts

@@ -142,8 +142,8 @@ export class PineconeVectorStore implements VectorStore {
     var options: any = {
       vector: query.queryEmbedding,
       topK: query.similarityTopK,
-      include_values: true,
-      include_metadara: true,
+      includeValues: true,
+      includeMetadata: true,
       filter: filter,
     };
 

packages/core/src/synthesizers/builders.ts

@@ -33,11 +33,11 @@ enum ResponseMode {
  */
 export class SimpleResponseBuilder implements ResponseBuilder {
   llm: LLM;
-  textQATemplate: SimplePrompt;
+  textQATemplate: TextQaPrompt;
 
-  constructor(serviceContext: ServiceContext) {
+  constructor(serviceContext: ServiceContext, textQATemplate?: TextQaPrompt) {
     this.llm = serviceContext.llm;
-    this.textQATemplate = defaultTextQaPrompt;
+    this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
   }
 
   getResponse(

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

@@ -1,6 +1,7 @@
 import nodeFS from "node:fs/promises";
 import os from "os";
 import path from "path";
+import { defaultFS } from "../env";
 import {
   GenericFileSystem,
   InMemoryFileSystem,
@@ -16,7 +17,7 @@ type FileSystemUnderTest = {
   tempDir: string;
 };
 
-describe.each<FileSystemUnderTest>([
+const cases: FileSystemUnderTest[] = [
   {
     name: "InMemoryFileSystem",
     prepare: async () => {},
@@ -27,17 +28,19 @@ describe.each<FileSystemUnderTest>([
     tempDir: "./",
   },
   {
-    name: "Node.js fs",
+    name: "Default fs",
     prepare: async function () {
       this.tempDir = await nodeFS.mkdtemp(path.join(os.tmpdir(), "jest-"));
     },
     cleanup: async function () {
       await nodeFS.rm(this.tempDir, { recursive: true });
     },
-    implementation: nodeFS,
+    implementation: defaultFS,
     tempDir: "./",
   },
-])("Test %s", (testParams) => {
+];
+
+describe.each<FileSystemUnderTest>(cases)("Test %s", (testParams) => {
   let testFS: GenericFileSystem;
   let tempDir: string;
 
@@ -58,7 +61,7 @@ describe.each<FileSystemUnderTest>([
   describe("writeFile", () => {
     it("writes file to memory", async () => {
       await testFS.writeFile(`${tempDir}/test.txt`, "Hello, world!");
-      expect(await testFS.readFile(`${tempDir}/test.txt`, "utf-8")).toBe(
+      expect(await testFS.readFile(`${tempDir}/test.txt`)).toBe(
         "Hello, world!",
       );
     });
@@ -66,7 +69,7 @@ describe.each<FileSystemUnderTest>([
     it("overwrites existing file", async () => {
       await testFS.writeFile(`${tempDir}/test.txt`, "Hello, world!");
       await testFS.writeFile(`${tempDir}/test.txt`, "Hello, again!");
-      expect(await testFS.readFile(`${tempDir}/test.txt`, "utf-8")).toBe(
+      expect(await testFS.readFile(`${tempDir}/test.txt`)).toBe(
         "Hello, again!",
       );
     });
@@ -75,7 +78,7 @@ describe.each<FileSystemUnderTest>([
   describe("readFile", () => {
     it("throws error for non-existing file", async () => {
       await expect(
-        testFS.readFile(`${tempDir}/not_exist.txt`, "utf-8"),
+        testFS.readFile(`${tempDir}/not_exist.txt`),
       ).rejects.toThrow();
     });
   });