changeset

Add observer/callback feature to SimpleDirectoryReader

.changeset/good-badgers-train.md

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

.changeset/long-feet-smell.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Added DocxReader for Word documents (thanks @jayantasamaddar)

.changeset/shiny-lions-clap.md

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

.changeset/slimy-points-walk.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Updated OpenAI streaming (thanks @kkang2097)

.changeset/wet-buckets-compete.md

@@ -1,5 +0,0 @@
----
-"llamaindex": patch
----
-
-Migrated to Tiktoken lite, which hopefully fixes the Windows issue

.gitignore

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

README.md

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

apps/docs/docs/environments.md

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

apps/docs/package.json

@@ -15,24 +15,24 @@
     "typecheck": "tsc"
   },
   "dependencies": {
-    "@docusaurus/core": "2.4.1",
-    "@docusaurus/preset-classic": "2.4.1",
-    "@docusaurus/remark-plugin-npm2yarn": "^2.4.1",
+    "@docusaurus/core": "2.4.3",
+    "@docusaurus/preset-classic": "2.4.3",
+    "@docusaurus/remark-plugin-npm2yarn": "^2.4.3",
     "@mdx-js/react": "^1.6.22",
     "clsx": "^1.2.1",
-    "postcss": "^8.4.28",
+    "postcss": "^8.4.31",
     "prism-react-renderer": "^1.3.5",
     "raw-loader": "^4.0.2",
     "react": "^17.0.2",
     "react-dom": "^17.0.2"
   },
   "devDependencies": {
-    "@docusaurus/module-type-aliases": "2.4.1",
-    "@docusaurus/types": "^2.4.1",
-    "@tsconfig/docusaurus": "^1.0.7",
+    "@docusaurus/module-type-aliases": "2.4.3",
+    "@docusaurus/types": "^2.4.3",
+    "@tsconfig/docusaurus": "^2.0.1",
     "docusaurus-plugin-typedoc": "^0.19.2",
     "typedoc": "^0.24.8",
-    "typedoc-plugin-markdown": "^3.15.4",
+    "typedoc-plugin-markdown": "^3.16.0",
     "typescript": "^4.9.5"
   },
   "browserslist": {

apps/simple/CHANGELOG.md

@@ -1,5 +1,34 @@
 # simple
 
+## 0.0.29
+
+### Patch Changes
+
+- Updated dependencies [6c55b2d]
+- Updated dependencies [8aa8c65]
+- Updated dependencies [6c55b2d]
+  - llamaindex@0.0.31
+
+## 0.0.28
+
+### Patch Changes
+
+- Updated dependencies [139abad]
+- Updated dependencies [139abad]
+- Updated dependencies [eb0e994]
+- Updated dependencies [eb0e994]
+- Updated dependencies [139abad]
+  - llamaindex@0.0.30
+
+## 0.0.27
+
+### Patch Changes
+
+- Updated dependencies [a52143b]
+- Updated dependencies [1b7fd95]
+- Updated dependencies [0db3f41]
+  - llamaindex@0.0.29
+
 ## 0.0.26
 
 ### Patch Changes

apps/simple/directory.ts

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

apps/simple/html.ts

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

apps/simple/llmStream.ts

@@ -0,0 +1,47 @@
+import { ChatMessage, OpenAI, SimpleChatEngine } from "llamaindex";
+import {Anthropic} from "../../packages/core/src/llm/LLM";
+import { stdin as input, stdout as output } from "node:process";
+import readline from "node:readline/promises";
+
+async function main() {
+  const query: string = `
+Where is Istanbul?
+  `;
+
+  // const llm = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 });
+  const llm = new Anthropic();
+  const message: ChatMessage = { content: query, role: "user" };
+
+  //TODO: Add callbacks later
+
+  //Stream Complete
+  //Note: Setting streaming flag to true or false will auto-set your return type to
+  //either an AsyncGenerator or a Response.
+  // Omitting the streaming flag automatically sets streaming to false
+
+  const chatEngine: SimpleChatEngine = new SimpleChatEngine({
+    chatHistory: undefined,
+    llm: llm,
+  });
+
+  const rl = readline.createInterface({ input, output });
+  while (true) {
+    const query = await rl.question("Query: ");
+
+    if (!query) {
+      break;
+    }
+
+    //Case 1: .chat(query, undefined, true) => Stream
+    //Case 2: .chat(query, undefined, false) => Response object
+    //Case 3: .chat(query, undefined) => Response object
+    const chatStream = await chatEngine.chat(query, undefined, true);
+    var accumulated_result = "";
+    for await (const part of chatStream) {
+      accumulated_result += part;
+      process.stdout.write(part);
+    }
+  }
+}
+
+main();

apps/simple/package.json

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

apps/simple/portkey.ts

@@ -0,0 +1,23 @@
+import { Portkey } from "llamaindex";
+
+(async () => {
+  const llms = [{
+
+  }]
+  const portkey = new Portkey({
+    mode: "single",
+    llms: [{
+      provider:"anyscale",
+      virtual_key:"anyscale-3b3c04",
+      model: "meta-llama/Llama-2-13b-chat-hf",
+      max_tokens: 2000
+    }]
+  });
+  const result = portkey.stream_chat([
+    { role: "system", content: "You are a helpful assistant." },
+    { role: "user", content: "Tell me a joke." }
+  ]);
+  for await (const res of result) {
+    process.stdout.write(res)
+  }
+})();

apps/simple/vectorIndexCustomize.ts

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

apps/simple/vectorIndexFromVectorStore.ts

@@ -0,0 +1,197 @@
+import {
+  OpenAI,
+  ResponseSynthesizer,
+  RetrieverQueryEngine,
+  serviceContextFromDefaults,
+  TextNode,
+  TreeSummarize,
+  VectorIndexRetriever,
+  VectorStore,
+  VectorStoreIndex,
+  VectorStoreQuery,
+  VectorStoreQueryResult,
+} from "llamaindex";
+
+import { Index, Pinecone, RecordMetadata } from "@pinecone-database/pinecone";
+
+/**
+ * Please do not use this class in production; it's only for demonstration purposes.
+ */
+class PineconeVectorStore<T extends RecordMetadata = RecordMetadata>
+  implements VectorStore
+{
+  storesText = true;
+  isEmbeddingQuery = false;
+
+  indexName!: string;
+  pineconeClient!: Pinecone;
+  index!: Index<T>;
+
+  constructor({ indexName, client }: { indexName: string; client: Pinecone }) {
+    this.indexName = indexName;
+    this.pineconeClient = client;
+    this.index = client.index<T>(indexName);
+  }
+
+  client() {
+    return this.pineconeClient;
+  }
+
+  async query(
+    query: VectorStoreQuery,
+    kwargs?: any,
+  ): Promise<VectorStoreQueryResult> {
+    let queryEmbedding: number[] = [];
+    if (query.queryEmbedding) {
+      if (typeof query.alpha === "number") {
+        const alpha = query.alpha;
+        queryEmbedding = query.queryEmbedding.map((v) => v * alpha);
+      } else {
+        queryEmbedding = query.queryEmbedding;
+      }
+    }
+
+    // Current LlamaIndexTS implementation only support exact match filter, so we use kwargs instead.
+    const filter = kwargs?.filter || {};
+
+    const response = await this.index.query({
+      filter,
+      vector: queryEmbedding,
+      topK: query.similarityTopK,
+      includeValues: true,
+      includeMetadata: true,
+    });
+
+    console.log(
+      `Numbers of vectors returned by Pinecone after preFilters are applied: ${
+        response?.matches?.length || 0
+      }.`,
+    );
+
+    const topKIds: string[] = [];
+    const topKNodes: TextNode[] = [];
+    const topKScores: number[] = [];
+
+    const metadataToNode = (metadata?: T): Partial<TextNode> => {
+      if (!metadata) {
+        throw new Error("metadata is undefined.");
+      }
+
+      const nodeContent = metadata["_node_content"];
+      if (!nodeContent) {
+        throw new Error("nodeContent is undefined.");
+      }
+
+      if (typeof nodeContent !== "string") {
+        throw new Error("nodeContent is not a string.");
+      }
+
+      return JSON.parse(nodeContent);
+    };
+
+    if (response.matches) {
+      for (const match of response.matches) {
+        const node = new TextNode({
+          ...metadataToNode(match.metadata),
+          embedding: match.values,
+        });
+
+        topKIds.push(match.id);
+        topKNodes.push(node);
+        topKScores.push(match.score ?? 0);
+      }
+    }
+
+    const result = {
+      ids: topKIds,
+      nodes: topKNodes,
+      similarities: topKScores,
+    };
+
+    return result;
+  }
+
+  add(): Promise<string[]> {
+    return Promise.resolve([]);
+  }
+
+  delete(): Promise<void> {
+    throw new Error("Method `delete` not implemented.");
+  }
+
+  persist(): Promise<void> {
+    throw new Error("Method `persist` not implemented.");
+  }
+}
+
+/**
+ * The goal of this example is to show how to use Pinecone as a vector store
+ * for LlamaIndexTS with(out) preFilters.
+ *
+ * It should not be used in production like that,
+ * as you might want to find a proper PineconeVectorStore implementation.
+ */
+async function main() {
+  process.env.PINECONE_API_KEY = "Your Pinecone API Key.";
+  process.env.PINECONE_ENVIRONMENT = "Your Pinecone Environment.";
+  process.env.PINECONE_PROJECT_ID = "Your Pinecone Project ID.";
+  process.env.PINECONE_INDEX_NAME = "Your Pinecone Index Name.";
+  process.env.OPENAI_API_KEY = "Your OpenAI API Key.";
+  process.env.OPENAI_API_ORGANIZATION = "Your OpenAI API Organization.";
+
+  const getPineconeVectorStore = async () => {
+    return new PineconeVectorStore({
+      indexName: process.env.PINECONE_INDEX_NAME || "index-name",
+      client: new Pinecone(),
+    });
+  };
+
+  const getServiceContext = () => {
+    const openAI = new OpenAI({
+      model: "gpt-4",
+      apiKey: process.env.OPENAI_API_KEY,
+    });
+
+    return serviceContextFromDefaults({
+      llm: openAI,
+    });
+  };
+
+  const getQueryEngine = async (filter: unknown) => {
+    const vectorStore = await getPineconeVectorStore();
+    const serviceContext = getServiceContext();
+
+    const vectorStoreIndex = await VectorStoreIndex.fromVectorStore(
+      vectorStore,
+      serviceContext,
+    );
+
+    const retriever = new VectorIndexRetriever({
+      index: vectorStoreIndex,
+      similarityTopK: 500,
+    });
+
+    const responseSynthesizer = new ResponseSynthesizer({
+      serviceContext,
+      responseBuilder: new TreeSummarize(serviceContext),
+    });
+
+    return new RetrieverQueryEngine(retriever, responseSynthesizer, {
+      filter,
+    });
+  };
+
+  // whatever is a key from your metadata
+  const queryEngine = await getQueryEngine({
+    whatever: {
+      $gte: 1,
+      $lte: 100,
+    },
+  });
+
+  const response = await queryEngine.query("How many results do you have?");
+
+  console.log(response.toString());
+}
+
+main().catch(console.error);

examples/llm_stream.ts

@@ -1,61 +0,0 @@
-import * as tiktoken from "tiktoken-node";
-import {
-  CallbackManager,
-  Event,
-  EventType,
-} from "../packages/core/src/callbacks/CallbackManager";
-import { ChatMessage, MessageType, OpenAI } from "../packages/core/src/llm/LLM";
-
-async function main() {
-  const query: string = "Where is Istanbul?";
-
-  const llm = new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.1 });
-  const message: ChatMessage = { content: query, role: "user" };
-
-  var accumulated_result: string = "";
-  var total_tokens: number = 0;
-
-  //Callback stuff, like logging token usage.
-  //GPT 3.5 Turbo uses CL100K_Base encodings, check your LLM to see which tokenizer it uses.
-  const encoding = tiktoken.getEncoding("cl100k_base");
-
-  const callback: CallbackManager = new CallbackManager();
-  callback.onLLMStream = (callback_response) => {
-    //Token text
-    const text = callback_response.token.choices[0].delta.content
-      ? callback_response.token.choices[0].delta.content
-      : "";
-    //Increment total number of tokens
-    total_tokens += encoding.encode(text).length;
-  };
-
-  llm.callbackManager = callback;
-
-  //Create a dummy event to trigger our Stream Callback
-  const dummy_event: Event = {
-    id: "something",
-    type: "intermediate" as EventType,
-  };
-
-  //Stream Complete
-  const stream = llm.stream_complete(query, dummy_event);
-
-  for await (const part of stream) {
-    //This only gives you the string part of a stream
-    console.log(part);
-    accumulated_result += part;
-  }
-
-  const correct_total_tokens: number =
-    encoding.encode(accumulated_result).length;
-
-  //Check if our stream token counter works
-  console.log(
-    `Output token total using tokenizer on accumulated output: ${correct_total_tokens}`,
-  );
-  console.log(
-    `Output token total using tokenizer on stream output: ${total_tokens}`,
-  );
-}
-
-main();

examples/portkey.ts

@@ -0,0 +1,23 @@
+import { Portkey } from "llamaindex";
+
+(async () => {
+  const llms = [{
+
+  }]
+  const portkey = new Portkey({
+    mode: "single",
+    llms: [{
+      provider:"anyscale",
+      virtual_key:"anyscale-3b3c04",
+      model: "meta-llama/Llama-2-13b-chat-hf",
+      max_tokens: 2000
+    }]
+  });
+  const result = portkey.stream_chat([
+    { role: "system", content: "You are a helpful assistant." },
+    { role: "user", content: "Tell me a joke." }
+  ]);
+  for await (const res of result) {
+    process.stdout.write(res)
+  }
+})();

examples/readonly.ts

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

package.json

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

packages/core/CHANGELOG.md

@@ -1,5 +1,31 @@
 # llamaindex
 
+## 0.0.31
+
+### Patch Changes
+
+- 6c55b2d: Give HistoryChatEngine pluggable options (thanks @marcusschiesser)
+- 8aa8c65: Add SimilarityPostProcessor (thanks @TomPenguin)
+- 6c55b2d: Added LLMMetadata (thanks @marcusschiesser)
+
+## 0.0.30
+
+### Patch Changes
+
+- 139abad: Streaming improvements including Anthropic (thanks @kkang2097)
+- 139abad: Portkey integration (Thank you @noble-varghese)
+- eb0e994: Add export for PromptHelper (thanks @zigamall)
+- eb0e994: Publish ESM module again
+- 139abad: Pinecone demo (thanks @Einsenhorn)
+
+## 0.0.29
+
+### Patch Changes
+
+- a52143b: Added DocxReader for Word documents (thanks @jayantasamaddar)
+- 1b7fd95: Updated OpenAI streaming (thanks @kkang2097)
+- 0db3f41: Migrated to Tiktoken lite, which hopefully fixes the Windows issue
+
 ## 0.0.28
 
 ### Patch Changes

packages/core/package.json

@@ -1,40 +1,47 @@
 {
   "name": "llamaindex",
-  "version": "0.0.28",
+  "version": "0.0.31",
+  "license": "MIT",
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.6.2",
+    "@anthropic-ai/sdk": "^0.8.1",
     "@notionhq/client": "^2.2.13",
+    "js-tiktoken": "^1.0.7",
     "lodash": "^4.17.21",
     "mammoth": "^1.6.0",
     "md-utils-ts": "^2.0.0",
-    "mongodb": "^6.1.0",
+    "mongodb": "^6.2.0",
     "notion-md-crawler": "^0.0.2",
-    "openai": "^4.10.0",
+    "openai": "^4.14.0",
     "papaparse": "^5.4.1",
     "pdf-parse": "^1.1.1",
+    "portkey-ai": "^0.1.13",
     "rake-modified": "^1.0.8",
-    "replicate": "^0.18.1",
-    "tiktoken": "^1.0.10",
+    "replicate": "^0.20.1",
+    "string-strip-html": "^13.4.3",
     "uuid": "^9.0.1",
     "wink-nlp": "^1.14.3"
   },
   "devDependencies": {
-    "@types/lodash": "^4.14.199",
-    "@types/node": "^18.18.0",
-    "@types/papaparse": "^5.3.9",
-    "@types/pdf-parse": "^1.1.2",
-    "@types/uuid": "^9.0.4",
+    "@types/lodash": "^4.14.200",
+    "@types/node": "^18.18.7",
+    "@types/papaparse": "^5.3.10",
+    "@types/pdf-parse": "^1.1.3",
+    "@types/uuid": "^9.0.6",
     "node-stdlib-browser": "^1.2.0",
-    "tsup": "^7.2.0"
+    "tsup": "^7.2.0",
+    "typescript": "^5.2.2"
   },
   "engines": {
     "node": ">=18.0.0"
   },
   "types": "./dist/index.d.ts",
   "main": "./dist/index.js",
+  "module": "./dist/index.mjs",
+  "repository": "run-llama/LlamaIndexTS",
   "scripts": {
     "lint": "eslint .",
     "test": "jest",
-    "build": "tsup src/index.ts --format esm,cjs --dts"
+    "build": "tsup src/index.ts --format esm,cjs --dts",
+    "dev": "tsup src/index.ts --format esm,cjs --dts --watch"
   }
 }

packages/core/src/ChatEngine.ts

@@ -1,8 +1,9 @@
 import { v4 as uuidv4 } from "uuid";
 import { Event } from "./callbacks/CallbackManager";
-import { ChatHistory, SimpleChatHistory } from "./ChatHistory";
+import { ChatHistory } from "./ChatHistory";
+import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
 import { ChatMessage, LLM, OpenAI } from "./llm/LLM";
-import { TextNode } from "./Node";
+import { NodeWithScore, TextNode } from "./Node";
 import {
   CondenseQuestionPrompt,
   ContextSystemPrompt,
@@ -23,8 +24,16 @@ export interface ChatEngine {
    * Send message along with the class's current chat history to the LLM.
    * @param message
    * @param chatHistory optional chat history if you want to customize the chat history
+   * @param streaming optional streaming flag, which auto-sets the return value if True.
    */
-  chat(message: string, chatHistory?: ChatMessage[]): Promise<Response>;
+  chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
+  >(
+    message: string,
+    chatHistory?: ChatMessage[],
+    streaming?: T,
+  ): Promise<R>;
 
   /**
    * Resets the chat history so that it's empty.
@@ -44,13 +53,45 @@ export class SimpleChatEngine implements ChatEngine {
     this.llm = init?.llm ?? new OpenAI();
   }
 
-  async chat(message: string, chatHistory?: ChatMessage[]): Promise<Response> {
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
+  >(message: string, chatHistory?: ChatMessage[], streaming?: T): Promise<R> {
+    //Streaming option
+    if (streaming) {
+      return this.streamChat(message, chatHistory) as R;
+    }
+
+    //Non-streaming option
     chatHistory = chatHistory ?? this.chatHistory;
     chatHistory.push({ content: message, role: "user" });
-    const response = await this.llm.chat(chatHistory);
+    const response = await this.llm.chat(chatHistory, undefined);
     chatHistory.push(response.message);
     this.chatHistory = chatHistory;
-    return new Response(response.message.content);
+    return new Response(response.message.content) as R;
+  }
+
+  protected async *streamChat(
+    message: string,
+    chatHistory?: ChatMessage[],
+  ): AsyncGenerator<string, void, unknown> {
+    chatHistory = chatHistory ?? this.chatHistory;
+    chatHistory.push({ content: message, role: "user" });
+    const response_generator = await this.llm.chat(
+      chatHistory,
+      undefined,
+      true,
+    );
+
+    var accumulator: string = "";
+    for await (const part of response_generator) {
+      accumulator += part;
+      yield part;
+    }
+
+    chatHistory.push({ content: accumulator, role: "assistant" });
+    this.chatHistory = chatHistory;
+    return;
   }
 
   reset() {
@@ -99,10 +140,14 @@ export class CondenseQuestionChatEngine implements ChatEngine {
     );
   }
 
-  async chat(
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
+  >(
     message: string,
     chatHistory?: ChatMessage[] | undefined,
-  ): Promise<Response> {
+    streaming?: T,
+  ): Promise<R> {
     chatHistory = chatHistory ?? this.chatHistory;
 
     const condensedQuestion = (
@@ -114,7 +159,7 @@ export class CondenseQuestionChatEngine implements ChatEngine {
     chatHistory.push({ content: message, role: "user" });
     chatHistory.push({ content: response.response, role: "assistant" });
 
-    return response;
+    return response as R;
   }
 
   reset() {
@@ -122,57 +167,117 @@ export class CondenseQuestionChatEngine implements ChatEngine {
   }
 }
 
+export interface Context {
+  message: ChatMessage;
+  nodes: NodeWithScore[];
+}
+
+export interface ContextGenerator {
+  generate(message: string, parentEvent?: Event): Promise<Context>;
+}
+
+export class DefaultContextGenerator implements ContextGenerator {
+  retriever: BaseRetriever;
+  contextSystemPrompt: ContextSystemPrompt;
+  nodePostprocessors: BaseNodePostprocessor[];
+
+  constructor(init: {
+    retriever: BaseRetriever;
+    contextSystemPrompt?: ContextSystemPrompt;
+    nodePostprocessors?: BaseNodePostprocessor[];
+  }) {
+    this.retriever = init.retriever;
+    this.contextSystemPrompt =
+      init?.contextSystemPrompt ?? defaultContextSystemPrompt;
+    this.nodePostprocessors = init.nodePostprocessors || [];
+  }
+
+  private applyNodePostprocessors(nodes: NodeWithScore[]) {
+    return this.nodePostprocessors.reduce(
+      (nodes, nodePostprocessor) => nodePostprocessor.postprocessNodes(nodes),
+      nodes,
+    );
+  }
+
+  async generate(message: string, parentEvent?: Event): Promise<Context> {
+    if (!parentEvent) {
+      parentEvent = {
+        id: uuidv4(),
+        type: "wrapper",
+        tags: ["final"],
+      };
+    }
+    const sourceNodesWithScore = await this.retriever.retrieve(
+      message,
+      parentEvent,
+    );
+
+    const nodes = this.applyNodePostprocessors(sourceNodesWithScore);
+
+    return {
+      message: {
+        content: this.contextSystemPrompt({
+          context: nodes.map((r) => (r.node as TextNode).text).join("\n\n"),
+        }),
+        role: "system",
+      },
+      nodes,
+    };
+  }
+}
+
 /**
  * ContextChatEngine uses the Index to get the appropriate context for each query.
  * The context is stored in the system prompt, and the chat history is preserved,
  * ideally allowing the appropriate context to be surfaced for each query.
  */
 export class ContextChatEngine implements ChatEngine {
-  retriever: BaseRetriever;
-  chatModel: OpenAI;
+  chatModel: LLM;
   chatHistory: ChatMessage[];
-  contextSystemPrompt: ContextSystemPrompt;
+  contextGenerator: ContextGenerator;
 
   constructor(init: {
     retriever: BaseRetriever;
-    chatModel?: OpenAI;
+    chatModel?: LLM;
     chatHistory?: ChatMessage[];
     contextSystemPrompt?: ContextSystemPrompt;
+    nodePostprocessors?: BaseNodePostprocessor[];
   }) {
-    this.retriever = init.retriever;
     this.chatModel =
       init.chatModel ?? new OpenAI({ model: "gpt-3.5-turbo-16k" });
     this.chatHistory = init?.chatHistory ?? [];
-    this.contextSystemPrompt =
-      init?.contextSystemPrompt ?? defaultContextSystemPrompt;
+    this.contextGenerator = new DefaultContextGenerator({
+      retriever: init.retriever,
+      contextSystemPrompt: init?.contextSystemPrompt,
+    });
   }
 
-  async chat(message: string, chatHistory?: ChatMessage[] | undefined) {
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
+  >(
+    message: string,
+    chatHistory?: ChatMessage[] | undefined,
+    streaming?: T,
+  ): Promise<R> {
     chatHistory = chatHistory ?? this.chatHistory;
 
+    //Streaming option
+    if (streaming) {
+      return this.streamChat(message, chatHistory) as R;
+    }
+
     const parentEvent: Event = {
       id: uuidv4(),
       type: "wrapper",
       tags: ["final"],
     };
-    const sourceNodesWithScore = await this.retriever.retrieve(
-      message,
-      parentEvent,
-    );
-
-    const systemMessage: ChatMessage = {
-      content: this.contextSystemPrompt({
-        context: sourceNodesWithScore
-          .map((r) => (r.node as TextNode).text)
-          .join("\n\n"),
-      }),
-      role: "system",
-    };
+    const context = await this.contextGenerator.generate(message, parentEvent);
 
     chatHistory.push({ content: message, role: "user" });
 
     const response = await this.chatModel.chat(
-      [systemMessage, ...chatHistory],
+      [context.message, ...chatHistory],
       parentEvent,
     );
     chatHistory.push(response.message);
@@ -181,8 +286,41 @@ export class ContextChatEngine implements ChatEngine {
 
     return new Response(
       response.message.content,
-      sourceNodesWithScore.map((r) => r.node),
+      context.nodes.map((r) => r.node),
+    ) as R;
+  }
+
+  protected async *streamChat(
+    message: string,
+    chatHistory?: ChatMessage[] | undefined,
+  ): AsyncGenerator<string, void, unknown> {
+    chatHistory = chatHistory ?? this.chatHistory;
+
+    const parentEvent: Event = {
+      id: uuidv4(),
+      type: "wrapper",
+      tags: ["final"],
+    };
+    const context = await this.contextGenerator.generate(message, parentEvent);
+
+    chatHistory.push({ content: message, role: "user" });
+
+    const response_stream = await this.chatModel.chat(
+      [context.message, ...chatHistory],
+      parentEvent,
+      true,
     );
+    var accumulator: string = "";
+    for await (const part of response_stream) {
+      accumulator += part;
+      yield part;
+    }
+
+    chatHistory.push({ content: accumulator, role: "assistant" });
+
+    this.chatHistory = chatHistory;
+
+    return;
   }
 
   reset() {
@@ -191,26 +329,69 @@ export class ContextChatEngine implements ChatEngine {
 }
 
 /**
- * HistoryChatEngine is a ChatEngine that uses a ChatHistory to keep track of the chat history. This is an example with the same behavior as SimpleChatEngine
- * TODO: generally use the ChatHistory instead of ChatMessage[] - breaking change
+ * HistoryChatEngine is a ChatEngine that uses a `ChatHistory` object
+ * to keeps track of chat's message history.
+ * A `ChatHistory` object is passed as a parameter for each call to the `chat` method,
+ * so the state of the chat engine is preserved between calls.
+ * Optionally, a `ContextGenerator` can be used to generate an additional context for each call to `chat`.
  */
-export class HistoryChatEngine implements ChatEngine {
-  chatHistory: ChatHistory;
+export class HistoryChatEngine {
   llm: LLM;
+  contextGenerator?: ContextGenerator;
 
   constructor(init?: Partial<HistoryChatEngine>) {
-    this.chatHistory = init?.chatHistory ?? new SimpleChatHistory();
     this.llm = init?.llm ?? new OpenAI();
+    this.contextGenerator = init?.contextGenerator;
   }
 
-  async chat(message: string): Promise<Response> {
-    this.chatHistory.addMessage({ content: message, role: "user" });
-    const response = await this.llm.chat(this.chatHistory.messages);
-    this.chatHistory.addMessage(response.message);
-    return new Response(response.message.content);
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : Response,
+  >(message: string, chatHistory: ChatHistory, streaming?: T): Promise<R> {
+    //Streaming option
+    if (streaming) {
+      return this.streamChat(message, chatHistory) as R;
+    }
+    const context = await this.contextGenerator?.generate(message);
+    chatHistory.addMessage({
+      content: message,
+      role: "user",
+    });
+    const response = await this.llm.chat(
+      await chatHistory.requestMessages(
+        context ? [context.message] : undefined,
+      ),
+    );
+    chatHistory.addMessage(response.message);
+    return new Response(response.message.content) as R;
   }
 
-  reset() {
-    this.chatHistory.reset();
+  protected async *streamChat(
+    message: string,
+    chatHistory: ChatHistory,
+  ): AsyncGenerator<string, void, unknown> {
+    const context = await this.contextGenerator?.generate(message);
+    chatHistory.addMessage({
+      content: message,
+      role: "user",
+    });
+    const response_stream = await this.llm.chat(
+      await chatHistory.requestMessages(
+        context ? [context.message] : undefined,
+      ),
+      undefined,
+      true,
+    );
+
+    var accumulator = "";
+    for await (const part of response_stream) {
+      accumulator += part;
+      yield part;
+    }
+    chatHistory.addMessage({
+      content: accumulator,
+      role: "assistant",
+    });
+    return;
   }
 }

packages/core/src/ChatHistory.ts

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

packages/core/src/GlobalsHelper.ts

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

packages/core/src/QueryEngine.ts

@@ -1,4 +1,6 @@
 import { v4 as uuidv4 } from "uuid";
+import { Event } from "./callbacks/CallbackManager";
+import { BaseNodePostprocessor } from "./indices/BaseNodePostprocessor";
 import { NodeWithScore, TextNode } from "./Node";
 import {
   BaseQuestionGenerator,
@@ -10,7 +12,6 @@ import { CompactAndRefine, ResponseSynthesizer } from "./ResponseSynthesizer";
 import { BaseRetriever } from "./Retriever";
 import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
 import { QueryEngineTool, ToolMetadata } from "./Tool";
-import { Event } from "./callbacks/CallbackManager";
 
 /**
  * A query engine is a question answerer that can use one or more steps.
@@ -30,16 +31,39 @@ export interface BaseQueryEngine {
 export class RetrieverQueryEngine implements BaseQueryEngine {
   retriever: BaseRetriever;
   responseSynthesizer: ResponseSynthesizer;
+  nodePostprocessors: BaseNodePostprocessor[];
+  preFilters?: unknown;
 
   constructor(
     retriever: BaseRetriever,
     responseSynthesizer?: ResponseSynthesizer,
+    preFilters?: unknown,
+    nodePostprocessors?: BaseNodePostprocessor[],
   ) {
     this.retriever = retriever;
     const serviceContext: ServiceContext | undefined =
       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);
   }
 
   async query(query: string, parentEvent?: Event) {
@@ -48,7 +72,7 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
       type: "wrapper",
       tags: ["final"],
     };
-    const nodes = await this.retriever.retrieve(query, _parentEvent);
+    const nodes = await this.retrieve(query, _parentEvent);
     return this.responseSynthesizer.synthesize(query, nodes, _parentEvent);
   }
 }

packages/core/src/ResponseSynthesizer.ts

@@ -1,18 +1,18 @@
+import { Event } from "./callbacks/CallbackManager";
+import { LLM } from "./llm/LLM";
 import { MetadataMode, NodeWithScore } from "./Node";
 import {
+  defaultRefinePrompt,
+  defaultTextQaPrompt,
+  defaultTreeSummarizePrompt,
   RefinePrompt,
   SimplePrompt,
   TextQaPrompt,
   TreeSummarizePrompt,
-  defaultRefinePrompt,
-  defaultTextQaPrompt,
-  defaultTreeSummarizePrompt,
 } from "./Prompt";
 import { getBiggestPrompt } from "./PromptHelper";
 import { Response } from "./Response";
 import { ServiceContext, serviceContextFromDefaults } from "./ServiceContext";
-import { Event } from "./callbacks/CallbackManager";
-import { LLM } from "./llm/LLM";
 
 /**
  * Response modes of the response synthesizer
@@ -231,6 +231,7 @@ export class TreeSummarize implements BaseResponseBuilder {
       throw new Error("Must have at least one text chunk");
     }
 
+    // Should we send the query here too?
     const packedTextChunks = this.serviceContext.promptHelper.repack(
       this.summaryTemplate,
       textChunks,
@@ -241,6 +242,7 @@ export class TreeSummarize implements BaseResponseBuilder {
         await this.serviceContext.llm.complete(
           this.summaryTemplate({
             context: packedTextChunks[0],
+            query,
           }),
           parentEvent,
         )
@@ -251,6 +253,7 @@ export class TreeSummarize implements BaseResponseBuilder {
           this.serviceContext.llm.complete(
             this.summaryTemplate({
               context: chunk,
+              query,
             }),
             parentEvent,
           ),

packages/core/src/Retriever.ts

@@ -1,11 +1,15 @@
+import { Event } from "./callbacks/CallbackManager";
 import { NodeWithScore } from "./Node";
 import { ServiceContext } from "./ServiceContext";
-import { Event } from "./callbacks/CallbackManager";
 
 /**
  * Retrievers retrieve the nodes that most closely match our query in similarity.
  */
 export interface BaseRetriever {
-  retrieve(query: string, parentEvent?: Event): Promise<NodeWithScore[]>;
+  retrieve(
+    query: string,
+    parentEvent?: Event,
+    preFilters?: unknown,
+  ): Promise<NodeWithScore[]>;
   getServiceContext(): ServiceContext;
 }

packages/core/src/callbacks/CallbackManager.ts

@@ -39,6 +39,13 @@ export interface DefaultStreamToken {
 //OpenAI stream token schema is the default.
 //Note: Anthropic and Replicate also use similar token schemas.
 export type OpenAIStreamToken = DefaultStreamToken;
+export type AnthropicStreamToken = {
+  completion: string;
+  model: string;
+  stop_reason: string | undefined;
+  stop?: boolean | undefined;
+  log_id?: string;
+};
 
 //
 //Callback Responses

packages/core/src/index.ts

@@ -1,30 +1,29 @@
+export * from "./callbacks/CallbackManager";
 export * from "./ChatEngine";
+export * from "./ChatHistory";
+export * from "./constants";
 export * from "./Embedding";
 export * from "./GlobalsHelper";
+export * from "./indices";
+export * from "./llm/LLM";
 export * from "./Node";
 export * from "./NodeParser";
 export * from "./OutputParser";
 export * from "./Prompt";
+export * from "./PromptHelper";
 export * from "./QueryEngine";
 export * from "./QuestionGenerator";
-export * from "./Response";
-export * from "./ResponseSynthesizer";
-export * from "./Retriever";
-export * from "./ServiceContext";
-export * from "./TextSplitter";
-export * from "./Tool";
-export * from "./constants";
-export * from "./llm/LLM";
-
-export * from "./indices";
-
-export * from "./callbacks/CallbackManager";
-
+export * from "./readers/base";
 export * from "./readers/CSVReader";
 export * from "./readers/MarkdownReader";
 export * from "./readers/NotionReader";
 export * from "./readers/PDFReader";
+export * from "./readers/HTMLReader";
 export * from "./readers/SimpleDirectoryReader";
-export * from "./readers/base";
-
+export * from "./Response";
+export * from "./ResponseSynthesizer";
+export * from "./Retriever";
+export * from "./ServiceContext";
 export * from "./storage";
+export * from "./TextSplitter";
+export * from "./Tool";

packages/core/src/indices/BaseNodePostprocessor.ts

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

packages/core/src/indices/index.ts

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

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

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

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

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

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

@@ -1,9 +1,9 @@
+import { Event } from "../../callbacks/CallbackManager";
+import { DEFAULT_SIMILARITY_TOP_K } from "../../constants";
 import { globalsHelper } from "../../GlobalsHelper";
 import { NodeWithScore } from "../../Node";
 import { BaseRetriever } from "../../Retriever";
 import { ServiceContext } from "../../ServiceContext";
-import { Event } from "../../callbacks/CallbackManager";
-import { DEFAULT_SIMILARITY_TOP_K } from "../../constants";
 import {
   VectorStoreQuery,
   VectorStoreQueryMode,
@@ -32,7 +32,7 @@ export class VectorIndexRetriever implements BaseRetriever {
     this.similarityTopK = similarityTopK ?? DEFAULT_SIMILARITY_TOP_K;
   }
 
-  async retrieve(query: string, parentEvent?: Event): Promise<NodeWithScore[]> {
+  async retrieve(query: string, parentEvent?: Event, preFilters?: unknown): Promise<NodeWithScore[]> {
     const queryEmbedding =
       await this.serviceContext.embedModel.getQueryEmbedding(query);
 
@@ -41,10 +41,15 @@ export class VectorIndexRetriever implements BaseRetriever {
       mode: VectorStoreQueryMode.DEFAULT,
       similarityTopK: this.similarityTopK,
     };
-    const result = await this.index.vectorStore.query(q);
+    const result = await this.index.vectorStore.query(q, preFilters);
 
     let nodesWithScores: NodeWithScore[] = [];
     for (let i = 0; i < result.ids.length; i++) {
+      const nodeFromResult = result.nodes?.[i];
+      if (!this.index.indexStruct.nodesDict[result.ids[i]] && nodeFromResult) {
+        this.index.indexStruct.nodesDict[result.ids[i]] = nodeFromResult;
+      }
+
       const node = this.index.indexStruct.nodesDict[result.ids[i]];
       nodesWithScores.push({
         node: node,

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

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

packages/core/src/llm/LLM.ts

@@ -1,5 +1,6 @@
 import OpenAILLM, { ClientOptions as OpenAIClientOptions } from "openai";
 import {
+  AnthropicStreamToken,
   CallbackManager,
   Event,
   EventType,
@@ -7,6 +8,8 @@ import {
   StreamCallbackResponse,
 } from "../callbacks/CallbackManager";
 
+import { LLMOptions } from "portkey-ai";
+import { globalsHelper, Tokenizers } from "../GlobalsHelper";
 import {
   AnthropicSession,
   ANTHROPIC_AI_PROMPT,
@@ -21,6 +24,7 @@ import {
   shouldUseAzure,
 } from "./azure";
 import { getOpenAISession, OpenAISession } from "./openai";
+import { getPortkeySession, PortkeySession } from "./portkey";
 import { ReplicateSession } from "./replicate";
 
 export type MessageType =
@@ -28,7 +32,8 @@ export type MessageType =
   | "assistant"
   | "system"
   | "generic"
-  | "function";
+  | "function"
+  | "memory";
 
 export interface ChatMessage {
   content: string;
@@ -44,31 +49,54 @@ export interface ChatResponse {
 // NOTE in case we need CompletionResponse to diverge from ChatResponse in the future
 export type CompletionResponse = ChatResponse;
 
+export interface LLMMetadata {
+  model: string;
+  temperature: number;
+  topP: number;
+  maxTokens?: number;
+  contextWindow: number;
+  tokenizer: Tokenizers | undefined;
+}
+
 /**
  * Unified language model interface
  */
 export interface LLM {
+  metadata: LLMMetadata;
+  // Whether a LLM has streaming support
+  hasStreaming: boolean;
   /**
    * Get a chat response from the LLM
    * @param messages
+   *
+   * The return type of chat() and complete() are set by the "streaming" parameter being set to True.
    */
-  chat(messages: ChatMessage[], parentEvent?: Event): Promise<ChatResponse>;
+  chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(
+    messages: ChatMessage[],
+    parentEvent?: Event,
+    streaming?: T,
+  ): Promise<R>;
 
   /**
    * Get a prompt completion from the LLM
    * @param prompt the prompt to complete
    */
-  complete(prompt: string, parentEvent?: Event): Promise<CompletionResponse>;
-
-  stream_chat?(
-    messages: ChatMessage[],
+  complete<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(
+    prompt: string,
     parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown>;
+    streaming?: T,
+  ): Promise<R>;
 
-  stream_complete?(
-    query: string,
-    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown>;
+  /**
+   * Calculates the number of tokens needed for the given chat messages
+   */
+  tokens(messages: ChatMessage[]): number;
 }
 
 export const GPT4_MODELS = {
@@ -93,13 +121,15 @@ export const ALL_AVAILABLE_OPENAI_MODELS = {
  * OpenAI LLM implementation
  */
 export class OpenAI implements LLM {
+  hasStreaming: boolean = true;
+
   // Per completion OpenAI params
   model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS;
   temperature: number;
   topP: number;
   maxTokens?: number;
   additionalChatOptions?: Omit<
-    Partial<OpenAILLM.Chat.CompletionCreateParams>,
+    Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
     "max_tokens" | "messages" | "model" | "temperature" | "top_p" | "streaming"
   >;
 
@@ -169,6 +199,32 @@ export class OpenAI implements LLM {
     this.callbackManager = init?.callbackManager;
   }
 
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_OPENAI_MODELS[this.model].contextWindow,
+      tokenizer: Tokenizers.CL100K_BASE,
+    };
+  }
+
+  tokens(messages: ChatMessage[]): number {
+    // for latest OpenAI models, see https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+    const tokenizer = globalsHelper.tokenizer(this.metadata.tokenizer);
+    const tokensPerMessage = 3;
+    let numTokens = 0;
+    for (const message of messages) {
+      numTokens += tokensPerMessage;
+      for (const value of Object.values(message)) {
+        numTokens += tokenizer(value).length;
+      }
+    }
+    numTokens += 3; // every reply is primed with <|im_start|>assistant<|im_sep|>
+    return numTokens;
+  }
+
   mapMessageType(
     messageType: MessageType,
   ): "user" | "assistant" | "system" | "function" {
@@ -186,11 +242,11 @@ export class OpenAI implements LLM {
     }
   }
 
-  async chat(
-    messages: ChatMessage[],
-    parentEvent?: Event,
-  ): Promise<ChatResponse> {
-    const baseRequestParams: OpenAILLM.Chat.CompletionCreateParams = {
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(messages: ChatMessage[], parentEvent?: Event, streaming?: T): Promise<R> {
+    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
       model: this.model,
       temperature: this.temperature,
       max_tokens: this.maxTokens,
@@ -201,6 +257,13 @@ export class OpenAI implements LLM {
       top_p: this.topP,
       ...this.additionalChatOptions,
     };
+    // Streaming
+    if (streaming) {
+      if (!this.hasStreaming) {
+        throw Error("No streaming support for this LLM.");
+      }
+      return this.streamChat(messages, parentEvent) as R;
+    }
     // Non-streaming
     const response = await this.session.openai.chat.completions.create({
       ...baseRequestParams,
@@ -208,24 +271,30 @@ export class OpenAI implements LLM {
     });
 
     const content = response.choices[0].message?.content ?? "";
-    return { message: { content, role: response.choices[0].message.role } };
+    return {
+      message: { content, role: response.choices[0].message.role },
+    } as R;
   }
 
-  async complete(
-    prompt: string,
-    parentEvent?: Event,
-  ): Promise<CompletionResponse> {
-    return this.chat([{ content: prompt, role: "user" }], parentEvent);
+  async complete<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
+    return this.chat(
+      [{ content: prompt, role: "user" }],
+      parentEvent,
+      streaming,
+    );
   }
 
   //We can wrap a stream in a generator to add some additional logging behavior
   //For future edits: syntax for generator type is <typeof Yield, typeof Return, typeof Accept>
   //"typeof Accept" refers to what types you'll accept when you manually call generator.next(<AcceptType>)
-  async *stream_chat(
+  protected async *streamChat(
     messages: ChatMessage[],
     parentEvent?: Event,
   ): AsyncGenerator<string, void, unknown> {
-    const baseRequestParams: OpenAILLM.Chat.CompletionCreateParams = {
+    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
       model: this.model,
       temperature: this.temperature,
       max_tokens: this.maxTokens,
@@ -279,12 +348,12 @@ export class OpenAI implements LLM {
     return;
   }
 
-  //Stream_complete doesn't need to be async because it's child function is already async
-  stream_complete(
+  //streamComplete doesn't need to be async because it's child function is already async
+  protected streamComplete(
     query: string,
     parentEvent?: Event,
   ): AsyncGenerator<string, void, unknown> {
-    return this.stream_chat([{ content: query, role: "user" }], parentEvent);
+    return this.streamChat([{ content: query, role: "user" }], parentEvent);
   }
 }
 
@@ -348,6 +417,7 @@ export class LlamaDeuce implements LLM {
   topP: number;
   maxTokens?: number;
   replicateSession: ReplicateSession;
+  hasStreaming: boolean;
 
   constructor(init?: Partial<LlamaDeuce>) {
     this.model = init?.model ?? "Llama-2-70b-chat-4bit";
@@ -362,6 +432,22 @@ export class LlamaDeuce implements LLM {
       init?.maxTokens ??
       ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow; // For Replicate, the default is 500 tokens which is too low.
     this.replicateSession = init?.replicateSession ?? new ReplicateSession();
+    this.hasStreaming = init?.hasStreaming ?? false;
+  }
+
+  tokens(messages: ChatMessage[]): number {
+    throw new Error("Method not implemented.");
+  }
+
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow,
+      tokenizer: undefined,
+    };
   }
 
   mapMessagesToPrompt(messages: ChatMessage[]) {
@@ -468,10 +554,10 @@ If a question does not make any sense, or is not factually coherent, explain why
     };
   }
 
-  async chat(
-    messages: ChatMessage[],
-    _parentEvent?: Event,
-  ): Promise<ChatResponse> {
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(messages: ChatMessage[], _parentEvent?: Event, streaming?: T): Promise<R> {
     const api = ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model]
       .replicateApi as `${string}/${string}:${string}`;
 
@@ -492,6 +578,9 @@ If a question does not make any sense, or is not factually coherent, explain why
       replicateOptions.input.max_length = this.maxTokens;
     }
 
+    //TODO: Add streaming for this
+
+    //Non-streaming
     const response = await this.replicateSession.replicate.run(
       api,
       replicateOptions,
@@ -502,24 +591,32 @@ If a question does not make any sense, or is not factually coherent, explain why
         //^ We need to do this because Replicate returns a list of strings (for streaming functionality which is not exposed by the run function)
         role: "assistant",
       },
-    };
+    } as R;
   }
 
-  async complete(
-    prompt: string,
-    parentEvent?: Event,
-  ): Promise<CompletionResponse> {
+  async complete<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
     return this.chat([{ content: prompt, role: "user" }], parentEvent);
   }
 }
 
+export const ALL_AVAILABLE_ANTHROPIC_MODELS = {
+  // both models have 100k context window, see https://docs.anthropic.com/claude/reference/selecting-a-model
+  "claude-2": { contextWindow: 100000 },
+  "claude-instant-1": { contextWindow: 100000 },
+};
+
 /**
  * Anthropic LLM implementation
  */
 
 export class Anthropic implements LLM {
+  hasStreaming: boolean = true;
+
   // Per completion Anthropic params
-  model: string;
+  model: keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS;
   temperature: number;
   topP: number;
   maxTokens?: number;
@@ -551,6 +648,21 @@ export class Anthropic implements LLM {
 
     this.callbackManager = init?.callbackManager;
   }
+  
+  tokens(messages: ChatMessage[]): number {
+    throw new Error("Method not implemented.");
+  }
+
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow: ALL_AVAILABLE_ANTHROPIC_MODELS[this.model].contextWindow,
+      tokenizer: undefined,
+    };
+  }
 
   mapMessagesToPrompt(messages: ChatMessage[]) {
     return (
@@ -567,10 +679,22 @@ export class Anthropic implements LLM {
     );
   }
 
-  async chat(
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(
     messages: ChatMessage[],
     parentEvent?: Event | undefined,
-  ): Promise<ChatResponse> {
+    streaming?: T,
+  ): Promise<R> {
+    //Streaming
+    if (streaming) {
+      if (!this.hasStreaming) {
+        throw Error("No streaming support for this LLM.");
+      }
+      return this.streamChat(messages, parentEvent) as R;
+    }
+    //Non-streaming
     const response = await this.session.anthropic.completions.create({
       model: this.model,
       prompt: this.mapMessagesToPrompt(messages),
@@ -583,12 +707,182 @@ export class Anthropic implements LLM {
       message: { content: response.completion.trimStart(), role: "assistant" },
       //^ We're trimming the start because Anthropic often starts with a space in the response
       // That space will be re-added when we generate the next prompt.
-    };
+    } as R;
   }
-  async complete(
+
+  protected async *streamChat(
+    messages: ChatMessage[],
+    parentEvent?: Event | undefined,
+  ): AsyncGenerator<string, void, unknown> {
+    // AsyncIterable<AnthropicStreamToken>
+    const stream: AsyncIterable<AnthropicStreamToken> =
+      await this.session.anthropic.completions.create({
+        model: this.model,
+        prompt: this.mapMessagesToPrompt(messages),
+        max_tokens_to_sample: this.maxTokens ?? 100000,
+        temperature: this.temperature,
+        top_p: this.topP,
+        stream: true,
+      });
+
+    var idx_counter: number = 0;
+    for await (const part of stream) {
+      //TODO: LLM Stream Callback, pending re-work.
+
+      idx_counter++;
+      yield part.completion;
+    }
+    return;
+  }
+
+  async complete<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(
     prompt: string,
     parentEvent?: Event | undefined,
-  ): Promise<CompletionResponse> {
-    return this.chat([{ content: prompt, role: "user" }], parentEvent);
+    streaming?: T,
+  ): Promise<R> {
+    if (streaming) {
+      return this.streamComplete(prompt, parentEvent) as R;
+    }
+    return this.chat(
+      [{ content: prompt, role: "user" }],
+      parentEvent,
+      streaming,
+    ) as R;
+  }
+
+  protected streamComplete(
+    prompt: string,
+    parentEvent?: Event | undefined,
+  ): AsyncGenerator<string, void, unknown> {
+    return this.streamChat([{ content: prompt, role: "user" }], parentEvent);
+  }
+}
+
+export class Portkey implements LLM {
+  hasStreaming: boolean = true;
+
+  apiKey?: string = undefined;
+  baseURL?: string = undefined;
+  mode?: string = undefined;
+  llms?: [LLMOptions] | null = undefined;
+  session: PortkeySession;
+  callbackManager?: CallbackManager;
+
+  constructor(init?: Partial<Portkey>) {
+    this.apiKey = init?.apiKey;
+    this.baseURL = init?.baseURL;
+    this.mode = init?.mode;
+    this.llms = init?.llms;
+    this.session = getPortkeySession({
+      apiKey: this.apiKey,
+      baseURL: this.baseURL,
+      llms: this.llms,
+      mode: this.mode,
+    });
+    this.callbackManager = init?.callbackManager;
+  }
+
+  tokens(messages: ChatMessage[]): number {
+    throw new Error("Method not implemented.");
+  }
+
+  get metadata(): LLMMetadata {
+    throw new Error("metadata not implemented for Portkey");
+  }
+
+  async chat<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(
+    messages: ChatMessage[],
+    parentEvent?: Event | undefined,
+    streaming?: T,
+    params?: Record<string, any>,
+  ): Promise<R> {
+    if (streaming) {
+      return this.streamChat(messages, parentEvent, params) as R;
+    } else {
+      const resolvedParams = params || {};
+      const response = await this.session.portkey.chatCompletions.create({
+        messages,
+        ...resolvedParams,
+      });
+
+      const content = response.choices[0].message?.content ?? "";
+      const role = response.choices[0].message?.role || "assistant";
+      return { message: { content, role: role as MessageType } } as R;
+    }
+  }
+
+  async complete<
+    T extends boolean | undefined = undefined,
+    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
+  >(
+    prompt: string,
+    parentEvent?: Event | undefined,
+    streaming?: T,
+  ): Promise<R> {
+    return this.chat(
+      [{ content: prompt, role: "user" }],
+      parentEvent,
+      streaming,
+    );
+  }
+
+  async *streamChat(
+    messages: ChatMessage[],
+    parentEvent?: Event,
+    params?: Record<string, any>,
+  ): AsyncGenerator<string, void, unknown> {
+    // Wrapping the stream in a callback.
+    const onLLMStream = this.callbackManager?.onLLMStream
+      ? this.callbackManager.onLLMStream
+      : () => {};
+
+    const chunkStream = await this.session.portkey.chatCompletions.create({
+      messages,
+      ...params,
+      stream: true,
+    });
+
+    const event: Event = parentEvent
+      ? parentEvent
+      : {
+          id: "unspecified",
+          type: "llmPredict" as EventType,
+        };
+
+    //Indices
+    var idx_counter: number = 0;
+    for await (const part of chunkStream) {
+      //Increment
+      part.choices[0].index = idx_counter;
+      const is_done: boolean =
+        part.choices[0].finish_reason === "stop" ? true : false;
+      //onLLMStream Callback
+
+      const stream_callback: StreamCallbackResponse = {
+        event: event,
+        index: idx_counter,
+        isDone: is_done,
+        // token: part,
+      };
+      onLLMStream(stream_callback);
+
+      idx_counter++;
+
+      yield part.choices[0].delta?.content ?? "";
+    }
+    return;
+  }
+
+  streamComplete(
+    query: string,
+    parentEvent?: Event,
+  ): AsyncGenerator<string, void, unknown> {
+    return this.streamChat([{ content: query, role: "user" }], parentEvent);
   }
 }

packages/core/src/llm/portkey.ts

@@ -0,0 +1,62 @@
+import _ from "lodash";
+import { LLMOptions, Portkey } from "portkey-ai";
+
+export const readEnv = (env: string, default_val?: string): string | undefined => {
+  if (typeof process !== 'undefined') {
+      return process.env?.[env] ?? default_val;
+  }
+  return default_val;
+};
+
+interface PortkeyOptions {
+  apiKey?: string;
+  baseURL?: string;
+  mode?: string;
+  llms?: [LLMOptions] | null
+}
+
+export class PortkeySession {
+  portkey: Portkey;
+
+  constructor(options:PortkeyOptions = {}) {
+    if (!options.apiKey) {
+      options.apiKey = readEnv('PORTKEY_API_KEY')
+    }
+
+    if (!options.baseURL) {
+      options.baseURL = readEnv('PORTKEY_BASE_URL', "https://api.portkey.ai")
+    }
+
+    this.portkey = new Portkey({});
+    this.portkey.llms = [{}]
+    if (!options.apiKey) {
+      throw new Error("Set Portkey ApiKey in PORTKEY_API_KEY env variable");
+    }
+
+    this.portkey = new Portkey(options);
+  }
+}
+
+let defaultPortkeySession: {
+  session: PortkeySession;
+  options: PortkeyOptions;
+}[] = [];
+
+/**
+ * Get a session for the Portkey API. If one already exists with the same options,
+ * it will be returned. Otherwise, a new session will be created.
+ * @param options
+ * @returns
+ */
+export function getPortkeySession(options: PortkeyOptions = {}) {
+  let session = defaultPortkeySession.find((session) => {
+    return _.isEqual(session.options, options);
+  })?.session;
+
+  if (!session) {
+    session = new PortkeySession(options);
+    defaultPortkeySession.push({ session, options });
+  }
+  return session;
+}
+

packages/core/src/readers/HTMLReader.ts

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

packages/core/src/readers/SimpleDirectoryReader.ts

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

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

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

packages/core/tsconfig.json

@@ -3,10 +3,10 @@
     "esModuleInterop": true,
     "forceConsistentCasingInFileNames": true,
     "isolatedModules": true,
+    "module": "esnext",
     "moduleResolution": "node",
     "preserveWatchOutput": true,
     "skipLibCheck": true,
-    "noEmit": true,
     "strict": true,
     "lib": ["es2015", "dom"],
     "target": "ES2015",

packages/eslint-config-custom/index.js

@@ -18,6 +18,12 @@ module.exports = {
           "OPENAI_API_BASE",
           "OPENAI_API_VERSION",
           "OPENAI_API_TYPE",
+          "OPENAI_API_ORGANIZATION",
+
+          "PINECONE_API_KEY",
+          "PINECONE_ENVIRONMENT",
+          "PINECONE_PROJECT_ID",
+          "PINECONE_INDEX_NAME",
 
           "AZURE_OPENAI_API_KEY",
           "AZURE_OPENAI_API_INSTANCE_NAME",