fix: only export settings

# Conflicts:
#	examples/readers/src/csv.ts
#	examples/recipes/cost-analysis.ts
#	packages/core/package.json
#	packages/core/src/indices/vectorStore/index.ts

.changeset/fresh-pugs-explain.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Add pipeline.register to create a managed index in LlamaCloud

.changeset/green-beers-listen.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: make edge run build after core

.changeset/lemon-windows-fetch.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add tool factory

.changeset/long-mirrors-battle.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: throw error when no pipelines exist for the retriever

.changeset/nasty-pillows-breathe.md

@@ -0,0 +1,7 @@
+---
+"llamaindex": patch
+"@llamaindex/env": patch
+"@llamaindex/edge": patch
+---
+
+feat: improve CallbackManager

.changeset/nine-fireants-bow.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Update the list of supported Azure OpenAI API versions as of 2024-04-02.

.changeset/strange-ducks-kiss.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: use claude3 with react agent

.changeset/violet-bags-hide.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add wikipedia tool

.changeset/warm-fans-guess.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add result type json

.changeset/young-bananas-wave.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Add support for doc store strategies to VectorStoreIndex.fromDocuments

.github/workflows/test.yml

@@ -2,13 +2,10 @@ name: Run Tests
 
 on: [push, pull_request]
 
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
-
 jobs:
-  e2e:
+  test:
     runs-on: ubuntu-latest
+
     steps:
       - uses: actions/checkout@v4
       - uses: pnpm/action-setup@v2
@@ -19,27 +16,6 @@ jobs:
           cache: "pnpm"
       - name: Install dependencies
         run: pnpm install
-      - name: Run E2E Tests
-        run: pnpm run e2e
-
-  test:
-    strategy:
-      fail-fast: false
-      matrix:
-        node-version: [18.x, 20.x, 21.x]
-    name: Test on Node.js ${{ matrix.node-version }}
-    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: ${{ matrix.node-version }}
-          cache: "pnpm"
-      - name: Install dependencies
-        run: pnpm install
       - name: Run tests
         run: pnpm run test
   typecheck:

apps/docs/package.json

@@ -15,12 +15,12 @@
     "typecheck": "tsc"
   },
   "dependencies": {
-    "@docusaurus/core": "^3.2.1",
-    "@docusaurus/remark-plugin-npm2yarn": "^3.2.1",
-    "@llamaindex/examples": "workspace:*",
-    "@mdx-js/react": "^3.0.1",
+    "@docusaurus/core": "^3.2.0",
+    "@llamaindex/env": "workspace:*",
+    "@docusaurus/remark-plugin-npm2yarn": "^3.2.0",
+    "@mdx-js/react": "^3.0.0",
     "clsx": "^2.1.0",
-    "postcss": "^8.4.38",
+    "postcss": "^8.4.33",
     "prism-react-renderer": "^2.3.1",
     "raw-loader": "^4.0.2",
     "react": "^18.2.0",
@@ -28,15 +28,15 @@
   },
   "devDependencies": {
     "@docusaurus/module-type-aliases": "3.2.0",
-    "@docusaurus/preset-classic": "^3.2.1",
-    "@docusaurus/theme-classic": "^3.2.1",
-    "@docusaurus/types": "^3.2.1",
+    "@docusaurus/preset-classic": "^3.2.0",
+    "@docusaurus/theme-classic": "^3.2.0",
+    "@docusaurus/types": "^3.2.0",
     "@tsconfig/docusaurus": "^2.0.3",
-    "@types/node": "^18.19.31",
+    "@types/node": "^18.19.10",
     "docusaurus-plugin-typedoc": "^0.22.0",
-    "typedoc": "^0.25.13",
+    "typedoc": "^0.25.12",
     "typedoc-plugin-markdown": "^3.17.1",
-    "typescript": "^5.4.4"
+    "typescript": "^5.4.3"
   },
   "browserslist": {
     "production": [

examples/agent/multi_document_agent.ts

@@ -86,6 +86,7 @@ async function main() {
     const agent = new OpenAIAgent({
       tools: queryEngineTools,
       llm: new OpenAI({ model: "gpt-4" }),
+      verbose: true,
     });
 
     documentAgents[title] = agent;
@@ -125,6 +126,7 @@ async function main() {
   const topAgent = new OpenAIAgent({
     toolRetriever: await objectIndex.asRetriever({}),
     llm: new OpenAI({ model: "gpt-4" }),
+    verbose: true,
     prefixMessages: [
       {
         content:
@@ -143,4 +145,4 @@ async function main() {
   });
 }
 
-void main();
+main();

examples/agent/openai.ts

@@ -59,6 +59,7 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [functionTool, functionTool2],
+    verbose: true,
   });
 
   // Chat with the agent
@@ -70,6 +71,6 @@ async function main() {
   console.log(String(response));
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("Done");
 });

examples/agent/query_openai_agent.ts

@@ -29,6 +29,7 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [queryEngineTool],
+    verbose: true,
   });
 
   // Chat with the agent
@@ -40,6 +41,6 @@ async function main() {
   console.log(String(response));
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("Done");
 });

examples/agent/react_agent.ts

@@ -65,6 +65,7 @@ async function main() {
   const agent = new ReActAgent({
     llm: anthropic,
     tools: [functionTool, functionTool2],
+    verbose: true,
   });
 
   // Chat with the agent
@@ -76,6 +77,6 @@ async function main() {
   console.log(String(response));
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("Done");
 });

examples/agent/step_wise_openai.ts

@@ -59,6 +59,7 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [functionTool, functionTool2],
+    verbose: true,
   });
 
   // Create a task to sum and divide numbers
@@ -89,6 +90,6 @@ async function main() {
   }
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("Done");
 });

examples/agent/step_wise_query_tool.ts

@@ -29,6 +29,7 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [queryEngineTool],
+    verbose: true,
   });
 
   const task = agent.createTask("What was his salary?");
@@ -58,6 +59,6 @@ async function main() {
   }
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("Done");
 });

examples/agent/step_wise_react.ts

@@ -59,6 +59,7 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new ReActAgent({
     tools: [functionTool, functionTool2],
+    verbose: true,
   });
 
   const task = agent.createTask("Divide 16 by 2 then add 20");
@@ -84,6 +85,6 @@ async function main() {
   }
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("Done");
 });

examples/agent/stream_openai_agent.ts

@@ -59,6 +59,7 @@ async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
     tools: [functionTool, functionTool2],
+    verbose: false,
   });
 
   const stream = await agent.chat({
@@ -71,6 +72,6 @@ async function main() {
   }
 }
 
-void main().then(() => {
+main().then(() => {
   console.log("\nDone");
 });

examples/agent/wiki.ts

@@ -1,27 +0,0 @@
-import { OpenAI, OpenAIAgent, WikipediaTool } from "llamaindex";
-
-async function main() {
-  const llm = new OpenAI({ model: "gpt-4-turbo" });
-  const wikiTool = new WikipediaTool();
-
-  // Create an OpenAIAgent with the Wikipedia tool
-  const agent = new OpenAIAgent({
-    llm,
-    tools: [wikiTool],
-  });
-
-  // Chat with the agent
-  const response = await agent.chat({
-    message: "Who was Goethe?",
-    stream: true,
-  });
-
-  for await (const chunk of response.response) {
-    process.stdout.write(chunk.response);
-  }
-}
-
-(async function () {
-  await main();
-  console.log("\nDone");
-})();

examples/agent/wikipedia-tool.ts

@@ -0,0 +1,23 @@
+import { OpenAIAgent, WikipediaTool } from "llamaindex";
+
+async function main() {
+  const wikipediaTool = new WikipediaTool();
+
+  // Create an OpenAIAgent with the function tools
+  const agent = new OpenAIAgent({
+    tools: [wikipediaTool],
+    verbose: true,
+  });
+
+  // Chat with the agent
+  const response = await agent.chat({
+    message: "Where is Ho Chi Minh City?",
+  });
+
+  // Print the response
+  console.log(response);
+}
+
+main().then(() => {
+  console.log("Done");
+});

examples/astradb/example.ts

@@ -55,4 +55,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/astradb/load.ts

@@ -27,4 +27,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/astradb/query.ts

@@ -23,4 +23,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/chatHistory.ts

@@ -1,18 +1,7 @@
 import { stdin as input, stdout as output } from "node:process";
 import readline from "node:readline/promises";
 
-import {
-  OpenAI,
-  Settings,
-  SimpleChatEngine,
-  SummaryChatHistory,
-} from "llamaindex";
-
-if (process.env.NODE_ENV === "development") {
-  Settings.callbackManager.on("llm-end", (event) => {
-    console.log("callers chain", event.reason?.computedCallers);
-  });
-}
+import { OpenAI, SimpleChatEngine, SummaryChatHistory } from "llamaindex";
 
 async function main() {
   // Set maxTokens to 75% of the context window size of 4096

examples/chromadb/preFilters.ts

@@ -54,4 +54,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/chromadb/test.ts

@@ -37,4 +37,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/cloud/from_documents.ts

@@ -1,44 +0,0 @@
-import fs from "node:fs/promises";
-
-import { stdin as input, stdout as output } from "node:process";
-
-import readline from "node:readline/promises";
-
-import { Document, LlamaCloudIndex } from "llamaindex";
-
-async function main() {
-  const path = "node_modules/llamaindex/examples/abramov.txt";
-
-  const essay = await fs.readFile(path, "utf-8");
-
-  // Create Document object with essay
-  const document = new Document({ text: essay, id_: path });
-
-  const index = await LlamaCloudIndex.fromDocuments({
-    documents: [document],
-    name: "test",
-    projectName: "default",
-    apiKey: process.env.LLAMA_CLOUD_API_KEY,
-    baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
-  });
-
-  const queryEngine = index.asQueryEngine({
-    denseSimilarityTopK: 5,
-  });
-
-  const rl = readline.createInterface({ input, output });
-
-  while (true) {
-    const query = await rl.question("Query: ");
-    const stream = await queryEngine.query({
-      query,
-      stream: true,
-    });
-    console.log();
-    for await (const chunk of stream) {
-      process.stdout.write(chunk.response);
-    }
-  }
-}
-
-main().catch(console.error);

examples/evaluation/correctness.ts

@@ -22,4 +22,4 @@ However, general relativity, published in 1915, extended these ideas to include
   console.log(result);
 }
 
-void main();
+main();

examples/evaluation/faithfulness.ts

@@ -36,4 +36,4 @@ async function main() {
   console.log(result);
 }
 
-void main();
+main();

examples/evaluation/relevancy.ts

@@ -37,4 +37,4 @@ async function main() {
   console.log(result);
 }
 
-void main();
+main();

examples/jsonExtract.ts

@@ -36,7 +36,9 @@ async function main() {
     ],
   });
 
-  console.log(response.message.content);
+  const json = JSON.parse(response.message.content);
+
+  console.log(json);
 }
 
 main().catch(console.error);

examples/milvus/load.ts

@@ -23,4 +23,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/milvus/query.ts

@@ -22,4 +22,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/mongo.ts

@@ -61,4 +61,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/mongodb/1_import.ts

@@ -31,4 +31,4 @@ async function importJsonToMongo() {
 }
 
 // Run the import function
-void importJsonToMongo();
+importJsonToMongo();

examples/mongodb/3_query.ts

@@ -27,4 +27,4 @@ async function query() {
   await client.close();
 }
 
-void query();
+query();

examples/multimodal/clip.ts

@@ -30,4 +30,4 @@ async function main() {
   console.log(`Similarity between "${text2}" and the image is ${sim2}`);
 }
 
-void main();
+main();

examples/multimodal/rag.ts

@@ -13,7 +13,7 @@ Settings.chunkSize = 512;
 Settings.chunkOverlap = 20;
 
 // Update llm
-Settings.llm = new OpenAI({ model: "gpt-4-turbo", maxTokens: 512 });
+Settings.llm = new OpenAI({ model: "gpt-4-vision-preview", maxTokens: 512 });
 
 // Update callbackManager
 Settings.callbackManager = new CallbackManager({

examples/nodeParser/MarkdownNodeParser.ts

@@ -21,4 +21,4 @@ Sub-header content
   console.log(splits);
 }
 
-void main();
+main();

examples/package.json

@@ -1,5 +1,5 @@
 {
-  "name": "@llamaindex/examples",
+  "name": "examples",
   "private": true,
   "version": "0.0.4",
   "dependencies": {
@@ -10,16 +10,15 @@
     "@zilliz/milvus2-sdk-node": "^2.3.5",
     "chromadb": "^1.8.1",
     "commander": "^11.1.0",
-    "dotenv": "^16.4.5",
-    "js-tiktoken": "^1.0.10",
-    "llamaindex": "workspace:latest",
-    "mongodb": "^6.5.0",
+    "dotenv": "^16.4.1",
+    "llamaindex": "latest",
+    "mongodb": "^6.2.0",
     "pathe": "^1.1.2"
   },
   "devDependencies": {
-    "@types/node": "^18.19.31",
+    "@types/node": "^18.19.10",
     "ts-node": "^10.9.2",
-    "typescript": "^5.4.4"
+    "typescript": "^5.4.3"
   },
   "scripts": {
     "lint": "eslint ."

examples/pg-vector-store/load-docs.ts

@@ -32,7 +32,7 @@ async function main(args: any) {
   console.log(`Found ${count} files`);
 
   console.log(`Importing contents from ${count} files in ${sourceDir}`);
-  const fileName = "";
+  var fileName = "";
   try {
     // Passing callback fn to the ctor here
     // will enable looging to console.
@@ -42,7 +42,7 @@ async function main(args: any) {
 
     const pgvs = new PGVectorStore();
     pgvs.setCollection(sourceDir);
-    await pgvs.clearCollection();
+    pgvs.clearCollection();
 
     const ctx = await storageContextFromDefaults({ vectorStore: pgvs });
 
@@ -65,4 +65,4 @@ async function main(args: any) {
   process.exit(0);
 }
 
-void main(process.argv).catch((err) => console.error(err));
+main(process.argv).catch((err) => console.error(err));

examples/pinecone-vector-store/load-docs.ts

@@ -32,7 +32,7 @@ async function main(args: any) {
   console.log(`Found ${count} files`);
 
   console.log(`Importing contents from ${count} files in ${sourceDir}`);
-  const fileName = "";
+  var fileName = "";
   try {
     // Passing callback fn to the ctor here
     // will enable looging to console.
@@ -63,4 +63,4 @@ async function main(args: any) {
   process.exit(0);
 }
 
-void main(process.argv).catch((err) => console.error(err));
+main(process.argv).catch((err) => console.error(err));

examples/prompts/promptMixin.ts

@@ -45,4 +45,4 @@ async function main() {
   await queryEngine.query({ query });
 }
 
-void main();
+main();

examples/qdrantdb/preFilters.ts

@@ -79,4 +79,4 @@ async function main() {
   }
 }
 
-void main();
+main();

examples/readers/src/docx.ts

@@ -20,4 +20,4 @@ async function main() {
   console.log(`Test query > ${SAMPLE_QUERY}:\n`, response.toString());
 }
 
-void main();
+main();

examples/readers/src/markdown.ts

@@ -20,4 +20,4 @@ async function main() {
   console.log(`Test query > ${SAMPLE_QUERY}:\n`, response.toString());
 }
 
-void main();
+main();

examples/recipes/cost-analysis.ts

@@ -1,9 +1,5 @@
-import { encodingForModel } from "js-tiktoken";
 import { OpenAI } from "llamaindex";
 import { Settings } from "llamaindex/Settings";
-import { extractText } from "llamaindex/llm/utils";
-
-const encoding = encodingForModel("gpt-4-0125-preview");
 
 const llm = new OpenAI({
   model: "gpt-4-0125-preview",
@@ -13,26 +9,16 @@ let tokenCount = 0;
 
 Settings.callbackManager.on("llm-start", (event) => {
   const { messages } = event.detail.payload;
-  tokenCount += messages.reduce((count, message) => {
-    return count + encoding.encode(extractText(message.content)).length;
-  }, 0);
+  tokenCount += llm.tokens(messages);
   console.log("Token count:", tokenCount);
   // https://openai.com/pricing
   // $10.00 / 1M tokens
-  console.log(`Price: $${(tokenCount / 1_000_000) * 10}`);
-});
-Settings.callbackManager.on("llm-end", (event) => {
-  const { response } = event.detail.payload;
-  tokenCount += encoding.encode(extractText(response.message.content)).length;
-  console.log("Token count:", tokenCount);
-  // https://openai.com/pricing
-  // $30.00 / 1M tokens
-  console.log(`Price: $${(tokenCount / 1_000_000) * 30}`);
+  console.log(`Price: $${(tokenCount / 1000000) * 10}`);
 });
 
 const question = "Hello, how are you?";
 console.log("Question:", question);
-void llm
+llm
   .chat({
     stream: true,
     messages: [

examples/routerQueryEngine.ts

@@ -65,4 +65,4 @@ async function main() {
   });
 }
 
-void main().then(() => console.log("Done"));
+main().then(() => console.log("Done"));

examples/split.ts

@@ -13,4 +13,4 @@ async function main() {
   console.log(chunks);
 }
 
-void main();
+main();

examples/toolsStream.ts

@@ -1,47 +0,0 @@
-import { ChatResponseChunk, OpenAI } from "llamaindex";
-
-async function main() {
-  const llm = new OpenAI({ model: "gpt-4-turbo" });
-
-  const args: Parameters<typeof llm.chat>[0] = {
-    additionalChatOptions: {
-      tool_choice: "auto",
-    },
-    messages: [
-      {
-        content: "Who was Goethe?",
-        role: "user",
-      },
-    ],
-    tools: [
-      {
-        metadata: {
-          name: "wikipedia_tool",
-          description: "A tool that uses a query engine to search Wikipedia.",
-          parameters: {
-            type: "object",
-            properties: {
-              query: {
-                type: "string",
-                description: "The query to search for",
-              },
-            },
-            required: ["query"],
-          },
-        },
-      },
-    ],
-  };
-
-  const stream = await llm.chat({ ...args, stream: true });
-  let chunk: ChatResponseChunk | null = null;
-  for await (chunk of stream) {
-    process.stdout.write(chunk.delta);
-  }
-  console.log(chunk?.additionalKwargs?.toolCalls[0]);
-}
-
-(async function () {
-  await main();
-  console.log("Done");
-})();

examples/vision.ts

@@ -1,7 +1,7 @@
 import { OpenAI } from "llamaindex";
 
 (async () => {
-  const llm = new OpenAI({ model: "gpt-4-turbo", temperature: 0.1 });
+  const llm = new OpenAI({ model: "gpt-4-vision-preview", temperature: 0.1 });
 
   // complete api
   const response1 = await llm.complete({ prompt: "How are you?" });

package.json

@@ -9,27 +9,27 @@
     "format:write": "prettier --ignore-unknown --write .",
     "lint": "turbo run lint",
     "prepare": "husky",
-    "e2e": "turbo run e2e",
     "test": "turbo run test",
     "type-check": "tsc -b --diagnostics",
     "release": "pnpm run check-minor-version && pnpm run build:release && changeset publish",
     "release-snapshot": "pnpm run check-minor-version && pnpm run build:release && changeset publish --tag snapshot",
     "check-minor-version": "node ./scripts/check-minor-version",
-    "new-version": "pnpm run build:release && changeset version && pnpm run check-minor-version",
-    "new-snapshot": "pnpm run build:release && changeset version --snapshot"
+    "update-version": "node ./scripts/update-version",
+    "new-version": "pnpm run build:release && changeset version && pnpm run check-minor-version && pnpm run update-version",
+    "new-snapshot": "pnpm run build:release && changeset version --snapshot && pnpm run update-version"
   },
   "devDependencies": {
     "@changesets/cli": "^2.27.1",
-    "eslint": "^8.57.0",
+    "eslint": "^8.56.0",
     "eslint-config-custom": "workspace:*",
-    "husky": "^9.0.11",
+    "husky": "^9.0.10",
     "lint-staged": "^15.2.2",
     "prettier": "^3.2.5",
     "prettier-plugin-organize-imports": "^3.2.4",
-    "turbo": "^1.13.2",
-    "typescript": "^5.4.4"
+    "turbo": "^1.12.3",
+    "typescript": "^5.4.3"
   },
-  "packageManager": "pnpm@8.15.6+sha256.01c01eeb990e379b31ef19c03e9d06a14afa5250b82e81303f88721c99ff2e6f",
+  "packageManager": "pnpm@8.15.1",
   "pnpm": {
     "overrides": {
       "trim": "1.0.1",

packages/core/CHANGELOG.md

@@ -1,37 +1,5 @@
 # llamaindex
 
-## 0.2.4
-
-### Patch Changes
-
-- 3bc77f7: gpt-4-turbo GA
-- 8d2b21e: Mistral 0.1.3
-
-## 0.2.3
-
-### Patch Changes
-
-- f0704ec: Support streaming for OpenAI agent (and OpenAI tool calls)
-- Removed 'parentEvent' - Use 'event.reason?.computedCallers' instead
-- 3cbfa98: Added LlamaCloudIndex.fromDocuments
-
-## 0.2.2
-
-### Patch Changes
-
-- 3f8407c: Add pipeline.register to create a managed index in LlamaCloud
-- 60a1603: fix: make edge run build after core
-- fececd8: feat: add tool factory
-- 1115f83: fix: throw error when no pipelines exist for the retriever
-- 7a23cc6: feat: improve CallbackManager
-- ea467fa: Update the list of supported Azure OpenAI API versions as of 2024-04-02.
-- 6d9e015: feat: use claude3 with react agent
-- 0b665bd: feat: add wikipedia tool
-- 24b4033: feat: add result type json
-- 8b28092: Add support for doc store strategies to VectorStoreIndex.fromDocuments
-- Updated dependencies [7a23cc6]
-  - @llamaindex/env@0.0.6
-
 ## 0.2.1
 
 ### Patch Changes

packages/core/e2e/.gitignore

@@ -1 +0,0 @@
-logs

packages/core/e2e/README.md

@@ -1,38 +0,0 @@
-# LlamaIndexTS Core E2E Tests
-
-## Overview
-
-We are using Node.js Test Runner to run E2E tests for LlamaIndexTS Core.
-
-It supports the following features:
-
-- Run tests in parallel
-- Pure Node.js Environment
-- Switch between mock and real LLM API
-- Customizable logics
-
-## Usage
-
-- Run with mock register:
-
-```shell
-node --import tsx --import ./mock-register.js --test ./node/basic.e2e.ts
-```
-
-- Run without mock register:
-
-```shell
-node --import tsx --test ./node/basic.e2e.ts
-```
-
-- Run with specific test:
-
-```shell
-node --import tsx --import ./mock-register.js --test-name-pattern=agent --test ./node/basic.e2e.ts
-```
-
-- Run with debug logs:
-
-```shell
-CONSOLA_LEVEL=5 node --import tsx --import ./mock-register.js --test-name-pattern=agent --test ./node/basic.e2e.ts
-```

packages/core/e2e/fixtures/llm/open_ai.ts

@@ -1,68 +0,0 @@
-import { faker } from "@faker-js/faker";
-import type {
-  ChatResponse,
-  ChatResponseChunk,
-  CompletionResponse,
-  LLM,
-  LLMChatParamsNonStreaming,
-  LLMChatParamsStreaming,
-  LLMCompletionParamsNonStreaming,
-  LLMCompletionParamsStreaming,
-} from "llamaindex/llm/types";
-
-export function getOpenAISession() {
-  return {};
-}
-
-export function isFunctionCallingModel() {
-  return true;
-}
-
-export class OpenAI implements LLM {
-  get metadata() {
-    return {
-      model: "mock-model",
-      temperature: 0.1,
-      topP: 1,
-      contextWindow: 2048,
-      tokenizer: undefined,
-      isFunctionCallingModel: true,
-    };
-  }
-  chat(
-    params: LLMChatParamsStreaming<Record<string, unknown>>,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
-  chat(
-    params: LLMChatParamsNonStreaming<Record<string, unknown>>,
-  ): Promise<ChatResponse>;
-  chat(
-    params:
-      | LLMChatParamsStreaming<Record<string, unknown>>
-      | LLMChatParamsNonStreaming<Record<string, unknown>>,
-  ): unknown {
-    if (params.stream) {
-      return {
-        [Symbol.asyncIterator]: async function* () {
-          yield {
-            delta: faker.word.words(),
-          } satisfies ChatResponseChunk;
-        },
-      };
-    }
-    return {
-      message: {
-        content: faker.lorem.paragraph(),
-        role: "assistant",
-      },
-    } satisfies ChatResponse;
-  }
-  complete(
-    params: LLMCompletionParamsStreaming,
-  ): Promise<AsyncIterable<CompletionResponse>>;
-  complete(
-    params: LLMCompletionParamsNonStreaming,
-  ): Promise<CompletionResponse>;
-  async complete(params: unknown): Promise<unknown> {
-    throw new Error("Method not implemented.");
-  }
-}

packages/core/e2e/mock-module.js

@@ -1,36 +0,0 @@
-/**
- * This script will replace the resolved module with the corresponding fixture file.
- */
-import { stat } from "node:fs/promises";
-import { join, relative } from "node:path";
-import { fileURLToPath, pathToFileURL } from "node:url";
-const packageDistDir = fileURLToPath(new URL("../dist", import.meta.url));
-const fixturesDir = fileURLToPath(new URL("./fixtures", import.meta.url));
-
-export async function resolve(specifier, context, nextResolve) {
-  const result = await nextResolve(specifier, context);
-  if (result.format === "builtin" || result.url.startsWith("node:")) {
-    return result;
-  }
-  const targetUrl = fileURLToPath(result.url).replace(/\.js$/, ".ts");
-  const relativePath = relative(packageDistDir, targetUrl);
-  if (relativePath.startsWith(".") || relativePath.startsWith("/")) {
-    return result;
-  }
-  const url = pathToFileURL(join(fixturesDir, relativePath)).toString();
-  const exist = await stat(fileURLToPath(url))
-    .then((stat) => stat.isFile())
-    .catch((err) => {
-      if (err.code === "ENOENT") {
-        return false;
-      }
-      throw err;
-    });
-  if (!exist) {
-    return result;
-  }
-  return {
-    url,
-    format: "module",
-  };
-}

packages/core/e2e/mock-register.js

@@ -1,3 +0,0 @@
-import { register } from "node:module";
-
-register("./mock-module.js", import.meta.url);

packages/core/e2e/node/basic.e2e.ts

@@ -1,139 +0,0 @@
-/* eslint-disable @typescript-eslint/no-floating-promises */
-import { consola } from "consola";
-import {
-  OpenAI,
-  OpenAIAgent,
-  Settings,
-  type LLM,
-  type LLMEndEvent,
-  type LLMStartEvent,
-} from "llamaindex";
-import { ok } from "node:assert";
-import type { WriteStream } from "node:fs";
-import { createWriteStream } from "node:fs";
-import { mkdir } from "node:fs/promises";
-import { join } from "node:path";
-import { after, before, beforeEach, describe, test } from "node:test";
-import { inspect } from "node:util";
-
-let llm: LLM;
-let fsStream: WriteStream;
-before(async () => {
-  const logUrl = new URL(
-    join(
-      "..",
-      "logs",
-      `basic.e2e.${new Date().toISOString().replace(/:/g, "-").replace(/\./g, "-")}.log`,
-    ),
-    import.meta.url,
-  );
-  await mkdir(new URL(".", logUrl), { recursive: true });
-  fsStream = createWriteStream(logUrl, {
-    encoding: "utf-8",
-  });
-});
-
-after(() => {
-  fsStream.end();
-});
-
-beforeEach((s) => {
-  fsStream.write("start: " + s.name + "\n");
-});
-
-const llmEventStartHandler = (event: LLMStartEvent) => {
-  const { payload } = event.detail;
-  fsStream.write(
-    "llmEventStart: " +
-      inspect(payload, {
-        depth: Infinity,
-      }) +
-      "\n",
-  );
-};
-
-const llmEventEndHandler = (event: LLMEndEvent) => {
-  const { payload } = event.detail;
-  fsStream.write(
-    "llmEventEnd: " +
-      inspect(payload, {
-        depth: Infinity,
-      }) +
-      "\n",
-  );
-};
-
-before(() => {
-  Settings.llm = new OpenAI({
-    model: "gpt-3.5-turbo",
-  });
-  llm = Settings.llm;
-  Settings.callbackManager.on("llm-start", llmEventStartHandler);
-  Settings.callbackManager.on("llm-end", llmEventEndHandler);
-});
-
-after(() => {
-  Settings.callbackManager.off("llm-start", llmEventStartHandler);
-  Settings.callbackManager.off("llm-end", llmEventEndHandler);
-});
-
-describe("llm", () => {
-  test("llm.chat", async () => {
-    const response = await llm.chat({
-      messages: [
-        {
-          content: "Hello",
-          role: "user",
-        },
-      ],
-    });
-    consola.debug("response:", response);
-    ok(typeof response.message.content === "string");
-  });
-
-  test("stream llm.chat", async () => {
-    const iter = await llm.chat({
-      stream: true,
-      messages: [
-        {
-          content: "hello",
-          role: "user",
-        },
-      ],
-    });
-    for await (const chunk of iter) {
-      consola.debug("chunk:", chunk);
-      ok(typeof chunk.delta === "string");
-    }
-  });
-});
-
-describe("agent", () => {
-  test("agent.chat", async () => {
-    const agent = new OpenAIAgent({
-      tools: [
-        {
-          call: async () => {
-            return "35 degrees and sunny in San Francisco";
-          },
-          metadata: {
-            name: "Weather",
-            description: "Get the weather",
-            parameters: {
-              type: "object",
-              properties: {
-                location: { type: "string" },
-              },
-              required: ["location"],
-            },
-          },
-        },
-      ],
-    });
-    const result = await agent.chat({
-      message: "What is the weather in San Francisco?",
-    });
-    consola.debug("response:", result.response);
-    ok(typeof result.response === "string");
-  });
-});

packages/core/e2e/package.json

@@ -1,16 +0,0 @@
-{
-  "name": "@llamaindex/core-e2e",
-  "private": true,
-  "version": "0.0.2",
-  "type": "module",
-  "scripts": {
-    "e2e": "node --import tsx --import ./mock-register.js --test ./node/*.e2e.ts",
-    "e2e:nomock": "node --import tsx --test ./node/*.e2e.ts"
-  },
-  "devDependencies": {
-    "@faker-js/faker": "^8.4.1",
-    "consola": "^3.2.3",
-    "llamaindex": "workspace:*",
-    "tsx": "^4.7.2"
-  }
-}

packages/core/e2e/tsconfig.json

@@ -1,23 +0,0 @@
-{
-  "extends": "../../../tsconfig.json",
-  "compilerOptions": {
-    "outDir": "./lib",
-    "module": "node16",
-    "moduleResolution": "node16",
-    "target": "ESNext"
-  },
-  "include": [
-    "./**/*.ts",
-    "./mock-module.js",
-    "./mock-register.js",
-    "./fixtures"
-  ],
-  "references": [
-    {
-      "path": "../../core/tsconfig.json"
-    },
-    {
-      "path": "../../env/tsconfig.json"
-    }
-  ]
-}

packages/core/jsr.json

@@ -1,8 +1,8 @@
 {
   "name": "@llamaindex/core",
-  "version": "0.2.3",
+  "version": "0.1.21",
   "exports": "./src/index.ts",
   "imports": {
-    "@llamaindex/env": "jsr:@llamaindex/env@0.0.6"
+    "@llamaindex/env": "jsr:@llamaindex/env@0.0.5"
   }
 }

packages/core/package.json

@@ -1,6 +1,6 @@
 {
   "name": "llamaindex",
-  "version": "0.2.4",
+  "version": "0.2.1",
   "expectedMinorVersion": "2",
   "license": "MIT",
   "type": "module",
@@ -8,49 +8,49 @@
     "@anthropic-ai/sdk": "^0.18.0",
     "@aws-crypto/sha256-js": "^5.2.0",
     "@datastax/astra-db-ts": "^0.1.4",
-    "@grpc/grpc-js": "^1.10.6",
-    "@llamaindex/cloud": "0.0.5",
+    "@grpc/grpc-js": "^1.10.2",
+    "@llamaindex/cloud": "0.0.4",
     "@llamaindex/env": "workspace:*",
-    "@mistralai/mistralai": "^0.1.3",
+    "@mistralai/mistralai": "^0.0.10",
     "@notionhq/client": "^2.2.14",
-    "@pinecone-database/pinecone": "^2.2.0",
-    "@qdrant/js-client-rest": "^1.8.2",
-    "@types/lodash": "^4.17.0",
-    "@types/node": "^18.19.31",
+    "@pinecone-database/pinecone": "^2.0.1",
+    "@qdrant/js-client-rest": "^1.7.0",
+    "@types/lodash": "^4.14.202",
+    "@types/node": "^18.19.14",
     "@types/papaparse": "^5.3.14",
-    "@types/pg": "^8.11.5",
-    "@xenova/transformers": "^2.16.1",
+    "@types/pg": "^8.11.0",
+    "@xenova/transformers": "^2.15.0",
     "@zilliz/milvus2-sdk-node": "^2.3.5",
-    "assemblyai": "^4.3.4",
+    "assemblyai": "^4.2.2",
     "chromadb": "~1.7.3",
-    "cohere-ai": "^7.9.2",
+    "cohere-ai": "^7.7.5",
     "js-tiktoken": "^1.0.10",
     "lodash": "^4.17.21",
     "magic-bytes.js": "^1.10.0",
-    "mammoth": "^1.7.1",
+    "mammoth": "^1.6.0",
     "md-utils-ts": "^2.0.0",
-    "mongodb": "^6.5.0",
+    "mongodb": "^6.3.0",
     "notion-md-crawler": "^0.0.2",
-    "openai": "^4.33.0",
+    "openai": "^4.26.1",
     "papaparse": "^5.4.1",
     "pathe": "^1.1.2",
     "pdf2json": "^3.0.5",
-    "pg": "^8.11.5",
-    "pgvector": "^0.1.8",
+    "pg": "^8.11.3",
+    "pgvector": "^0.1.7",
     "portkey-ai": "^0.1.16",
     "rake-modified": "^1.0.8",
     "replicate": "^0.25.2",
-    "string-strip-html": "^13.4.8",
-    "wikipedia": "^2.1.2",
-    "wink-nlp": "^1.14.3"
+    "string-strip-html": "^13.4.6",
+    "wink-nlp": "^1.14.3",
+    "wikipedia": "^2.1.2"
   },
   "devDependencies": {
-    "@swc/cli": "^0.3.12",
-    "@swc/core": "^1.4.13",
+    "@swc/cli": "^0.3.9",
+    "@swc/core": "^1.4.2",
     "concurrently": "^8.2.2",
-    "glob": "^10.3.12",
+    "glob": "^10.3.10",
     "madge": "^6.1.0",
-    "typescript": "^5.4.4"
+    "typescript": "^5.3.3"
   },
   "engines": {
     "node": ">=18.0.0"

packages/core/src/ChatHistory.ts

@@ -1,9 +1,7 @@
-import { globalsHelper } from "./GlobalsHelper.js";
+import { OpenAI } from "./llm/LLM.js";
+import type { ChatMessage, LLM, MessageType } from "./llm/types.js";
 import type { SummaryPrompt } from "./Prompt.js";
 import { defaultSummaryPrompt, messagesToHistoryStr } from "./Prompt.js";
-import { OpenAI } from "./llm/open_ai.js";
-import type { ChatMessage, LLM, MessageType } from "./llm/types.js";
-import { extractText } from "./llm/utils.js";
 
 /**
  * A ChatHistory is used to keep the state of back and forth chat messages
@@ -64,12 +62,6 @@ export class SimpleChatHistory extends ChatHistory {
 }
 
 export class SummaryChatHistory extends ChatHistory {
-  /**
-   * Tokenizer function that converts text to tokens,
-   *  this is used to calculate the number of tokens in a message.
-   */
-  tokenizer: (text: string) => Uint32Array =
-    globalsHelper.defaultTokenizer.encode;
   tokensToSummarize: number;
   messages: ChatMessage[];
   summaryPrompt: SummaryPrompt;
@@ -112,9 +104,7 @@ export class SummaryChatHistory extends ChatHistory {
       ];
       // remove oldest message until the chat history is short enough for the context window
       messagesToSummarize.shift();
-    } while (
-      this.tokenizer(promptMessages[0].content).length > this.tokensToSummarize
-    );
+    } while (this.llm.tokens(promptMessages) > this.tokensToSummarize);
 
     const response = await this.llm.chat({ messages: promptMessages });
     return { content: response.message.content, role: "memory" };
@@ -188,11 +178,7 @@ export class SummaryChatHistory extends ChatHistory {
     const requestMessages = this.calcCurrentRequestMessages(transientMessages);
 
     // get tokens of current request messages and the transient messages
-    const tokens = requestMessages.reduce(
-      (count, message) =>
-        count + this.tokenizer(extractText(message.content)).length,
-      0,
-    );
+    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();

packages/core/src/GlobalsHelper.ts

@@ -1,5 +1,12 @@
 import { encodingForModel } from "js-tiktoken";
 
+import { randomUUID } from "@llamaindex/env";
+import type {
+  Event,
+  EventTag,
+  EventType,
+} from "./callbacks/CallbackManager.js";
+
 export enum Tokenizers {
   CL100K_BASE = "cl100k_base",
 }
@@ -11,9 +18,9 @@ class GlobalsHelper {
   defaultTokenizer: {
     encode: (text: string) => Uint32Array;
     decode: (tokens: Uint32Array) => string;
-  };
+  } | null = null;
 
-  constructor() {
+  private initDefaultTokenizer() {
     const encoding = encodingForModel("text-embedding-ada-002"); // cl100k_base
 
     this.defaultTokenizer = {
@@ -33,6 +40,9 @@ class GlobalsHelper {
     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);
   }
@@ -41,9 +51,45 @@ class GlobalsHelper {
     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);
   }
+
+  /**
+   * @deprecated createEvent will be removed in the future,
+   *  please use `new CustomEvent(eventType, { detail: payload })` instead.
+   *
+   *  Also, `parentEvent` will not be used in the future,
+   *    use `AsyncLocalStorage` to track parent events instead.
+   *    @example - Usage of `AsyncLocalStorage`:
+   *    let id = 0;
+   *    const asyncLocalStorage = new AsyncLocalStorage<number>();
+   *    asyncLocalStorage.run(++id, async () => {
+   *      setTimeout(() => {
+   *        console.log('parent event id:', asyncLocalStorage.getStore()); // 1
+   *      }, 1000)
+   *    });
+   */
+  createEvent({
+    parentEvent,
+    type,
+    tags,
+  }: {
+    parentEvent?: Event;
+    type: EventType;
+    tags?: EventTag[];
+  }): Event {
+    return {
+      id: randomUUID(),
+      type,
+      // inherit parent tags if tags not set
+      tags: tags || parentEvent?.tags,
+      parentId: parentEvent?.id,
+    };
+  }
 }
 
 export const globalsHelper = new GlobalsHelper();

packages/core/src/QuestionGenerator.ts

@@ -5,7 +5,7 @@ import type {
   BaseQuestionGenerator,
   SubQuestion,
 } from "./engines/query/types.js";
-import { OpenAI } from "./llm/open_ai.js";
+import { OpenAI } from "./llm/LLM.js";
 import type { LLM } from "./llm/types.js";
 import { PromptMixin } from "./prompts/index.js";
 import type {

packages/core/src/Retriever.ts

@@ -1,8 +1,13 @@
+import type { Event } from "./callbacks/CallbackManager.js";
 import type { NodeWithScore } from "./Node.js";
 import type { ServiceContext } from "./ServiceContext.js";
 
 export type RetrieveParams = {
   query: string;
+  /**
+   * @deprecated will be removed in the next major version
+   */
+  parentEvent?: Event;
   preFilters?: unknown;
 };
 

packages/core/src/ServiceContext.ts

@@ -1,7 +1,7 @@
 import { PromptHelper } from "./PromptHelper.js";
 import { OpenAIEmbedding } from "./embeddings/OpenAIEmbedding.js";
 import type { BaseEmbedding } from "./embeddings/types.js";
-import { OpenAI } from "./llm/open_ai.js";
+import { OpenAI } from "./llm/LLM.js";
 import type { LLM } from "./llm/types.js";
 import { SimpleNodeParser } from "./nodeParsers/SimpleNodeParser.js";
 import type { NodeParser } from "./nodeParsers/types.js";

packages/core/src/Settings.ts

@@ -1,11 +1,11 @@
 import { CallbackManager } from "./callbacks/CallbackManager.js";
 import { OpenAIEmbedding } from "./embeddings/OpenAIEmbedding.js";
-import { OpenAI } from "./llm/open_ai.js";
+import { OpenAI } from "./llm/LLM.js";
 
 import { PromptHelper } from "./PromptHelper.js";
 import { SimpleNodeParser } from "./nodeParsers/SimpleNodeParser.js";
 
-import { AsyncLocalStorage, getEnv } from "@llamaindex/env";
+import { AsyncLocalStorage } from "@llamaindex/env";
 import type { ServiceContext } from "./ServiceContext.js";
 import type { BaseEmbedding } from "./embeddings/types.js";
 import {
@@ -52,15 +52,6 @@ class GlobalSettings implements Config {
   #chunkOverlapAsyncLocalStorage = new AsyncLocalStorage<number>();
   #promptAsyncLocalStorage = new AsyncLocalStorage<PromptConfig>();
 
-  get debug() {
-    const debug = getEnv("DEBUG");
-    return (
-      getEnv("NODE_ENV") === "development" &&
-      Boolean(debug) &&
-      debug?.includes("llamaindex")
-    );
-  }
-
   get llm(): LLM {
     if (this.#llm === null) {
       this.#llm = new OpenAI();

packages/core/src/TextSplitter.ts

@@ -66,9 +66,8 @@ export const defaultParagraphSeparator = EOL + EOL + EOL;
  * One of the advantages of SentenceSplitter is that even in the fixed length chunks it will try to keep sentences together.
  */
 export class SentenceSplitter {
-  public chunkSize: number;
-  public chunkOverlap: number;
-
+  private chunkSize: number;
+  private chunkOverlap: number;
   private tokenizer: any;
   private tokenizerDecoder: any;
   private paragraphSeparator: string;

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

@@ -1,4 +1,3 @@
-import { Settings } from "../../Settings.js";
 import type { ChatMessage } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
 import type { ObjectRetriever } from "../../objects/base.js";
@@ -11,6 +10,7 @@ type OpenAIAgentParams = {
   llm?: OpenAI;
   memory?: any;
   prefixMessages?: ChatMessage[];
+  verbose?: boolean;
   maxFunctionCalls?: number;
   defaultToolChoice?: string;
   toolRetriever?: ObjectRetriever;
@@ -28,19 +28,13 @@ export class OpenAIAgent extends AgentRunner {
     llm,
     memory,
     prefixMessages,
+    verbose,
     maxFunctionCalls = 5,
     defaultToolChoice = "auto",
     toolRetriever,
     systemPrompt,
   }: OpenAIAgentParams) {
-    if (!llm) {
-      if (Settings.llm instanceof OpenAI) {
-        llm = Settings.llm;
-      } else {
-        console.warn("No OpenAI model provided, creating a new one");
-        llm = new OpenAI({ model: "gpt-3.5-turbo-0613" });
-      }
-    }
+    llm = llm ?? new OpenAI({ model: "gpt-3.5-turbo-0613" });
 
     if (systemPrompt) {
       if (prefixMessages) {
@@ -65,11 +59,11 @@ export class OpenAIAgent extends AgentRunner {
       prefixMessages,
       maxFunctionCalls,
       toolRetriever,
+      verbose,
     });
 
     super({
       agentWorker: stepEngine,
-      llm,
       memory,
       defaultToolChoice,
       chatHistory: prefixMessages,

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

@@ -0,0 +1,27 @@
+import type { ToolMetadata } from "../../types.js";
+
+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

@@ -1,25 +1,17 @@
 import { randomUUID } from "@llamaindex/env";
-import type { ChatCompletionToolChoiceOption } from "openai/resources/chat/completions";
 import { Response } from "../../Response.js";
-import { Settings } from "../../Settings.js";
 import {
   AgentChatResponse,
   ChatResponseMode,
   StreamingAgentChatResponse,
 } from "../../engines/chat/types.js";
-import {
-  OpenAI,
-  isFunctionCallingModel,
-  type ChatMessage,
-  type ChatResponseChunk,
-  type LLMChatParamsBase,
-  type OpenAIAdditionalChatOptions,
+import type {
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
 } from "../../llm/index.js";
-import {
-  extractText,
-  streamConverter,
-  streamReducer,
-} from "../../llm/utils.js";
+import { OpenAI } from "../../llm/index.js";
+import { streamConverter, streamReducer } from "../../llm/utils.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
 import type { ObjectRetriever } from "../../objects/base.js";
 import type { ToolOutput } from "../../tools/types.js";
@@ -29,17 +21,28 @@ import type { AgentWorker, Task } from "../types.js";
 import { TaskStep, TaskStepOutput } from "../types.js";
 import { addUserStepToMemory, getFunctionByName } from "../utils.js";
 import type { OpenAIToolCall } from "./types/chat.js";
+import { toOpenAiTool } from "./utils.js";
 
+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 (Settings.debug) {
+  if (verbose) {
     console.log("=== Calling Function ===");
     console.log(`Calling function: ${name} with args: ${argumentsStr}`);
   }
@@ -51,7 +54,7 @@ async function callFunction(
   // Use default error message
   const output = await callToolWithErrorHandling(tool, argumentDict, null);
 
-  if (Settings.debug) {
+  if (verbose) {
     console.log(`Got output ${output}`);
     console.log("==========================");
   }
@@ -73,6 +76,7 @@ type OpenAIAgentWorkerParams = {
   tools?: BaseTool[];
   llm?: OpenAI;
   prefixMessages?: ChatMessage[];
+  verbose?: boolean;
   maxFunctionCalls?: number;
   toolRetriever?: ObjectRetriever;
 };
@@ -82,40 +86,40 @@ type CallFunctionOutput = {
   toolOutput: ToolOutput;
 };
 
-export class OpenAIAgentWorker
-  implements AgentWorker<LLMChatParamsBase<OpenAIAdditionalChatOptions>>
-{
+/**
+ * OpenAI agent worker.
+ * This class is responsible for running the agent.
+ */
+export class OpenAIAgentWorker implements AgentWorker {
   private llm: OpenAI;
-  private maxFunctionCalls: number = 5;
+  private verbose: boolean;
+  private maxFunctionCalls: number;
 
   public prefixMessages: ChatMessage[];
 
   private _getTools: (input: string) => Promise<BaseTool[]>;
 
+  /**
+   * Initialize.
+   */
   constructor({
     tools = [],
     llm,
     prefixMessages,
-    maxFunctionCalls,
+    verbose,
+    maxFunctionCalls = DEFAULT_MAX_FUNCTION_CALLS,
     toolRetriever,
   }: OpenAIAgentWorkerParams) {
-    this.llm =
-      llm ?? isFunctionCallingModel(Settings.llm)
-        ? (Settings.llm as OpenAI)
-        : new OpenAI({
-            model: "gpt-3.5-turbo-0613",
-          });
-    if (maxFunctionCalls) {
-      this.maxFunctionCalls = maxFunctionCalls;
-    }
+    this.llm = llm ?? new OpenAI({ model: "gpt-3.5-turbo-0613" });
+    this.verbose = verbose || false;
+    this.maxFunctionCalls = maxFunctionCalls;
     this.prefixMessages = prefixMessages || [];
 
-    if (Array.isArray(tools) && tools.length > 0 && toolRetriever) {
+    if (tools.length > 0 && toolRetriever) {
       throw new Error("Cannot specify both tools and tool_retriever");
-    } else if (Array.isArray(tools)) {
+    } else if (tools.length > 0) {
       this._getTools = async () => tools;
     } else if (toolRetriever) {
-      // fixme: this won't work, type mismatch
       this._getTools = async (message: string) =>
         toolRetriever.retrieve(message);
     } else {
@@ -123,6 +127,11 @@ export class OpenAIAgentWorker
     }
   }
 
+  /**
+   * Get all messages.
+   * @param task: task
+   * @returns: messages
+   */
   public getAllMessages(task: Task): ChatMessage[] {
     return [
       ...this.prefixMessages,
@@ -131,6 +140,11 @@ export class OpenAIAgentWorker
     ];
   }
 
+  /**
+   * Get latest tool calls.
+   * @param task: task
+   * @returns: tool calls
+   */
   public getLatestToolCalls(task: Task): OpenAIToolCall[] | null {
     const chatHistory: ChatMessage[] = task.extraState.newMemory.getAll();
 
@@ -141,67 +155,59 @@ export class OpenAIAgentWorker
     return chatHistory[chatHistory.length - 1].additionalKwargs?.toolCalls;
   }
 
-  private _getLlmChatParams(
+  /**
+   *
+   * @param task
+   * @param openaiTools
+   * @param toolChoice
+   * @returns
+   */
+  private _getLlmChatKwargs(
     task: Task,
-    openaiTools: BaseTool[],
-    toolChoice: ChatCompletionToolChoiceOption = "auto",
-  ): LLMChatParamsBase<OpenAIAdditionalChatOptions> {
-    const llmChatParams = {
+    openaiTools: { [key: string]: any }[],
+    toolChoice: string | { [key: string]: any } = "auto",
+  ): { [key: string]: any } {
+    const llmChatKwargs: { [key: string]: any } = {
       messages: this.getAllMessages(task),
-      tools: [] as BaseTool[],
-      additionalChatOptions: {} as OpenAIAdditionalChatOptions,
-    } satisfies LLMChatParamsBase<OpenAIAdditionalChatOptions>;
+    };
 
     if (openaiTools.length > 0) {
-      llmChatParams.tools = openaiTools;
-      llmChatParams.additionalChatOptions.tool_choice = toolChoice;
+      llmChatKwargs.tools = openaiTools;
+      llmChatKwargs.toolChoice = toolChoice;
     }
 
-    return llmChatParams;
+    return llmChatKwargs;
   }
 
+  /**
+   * Process message.
+   * @param task: task
+   * @param chatResponse: chat response
+   * @returns: agent chat response
+   */
   private _processMessage(
     task: Task,
-    aiMessage: ChatMessage,
+    chatResponse: ChatResponse,
   ): AgentChatResponse {
+    const aiMessage = chatResponse.message;
     task.extraState.newMemory.put(aiMessage);
 
-    return new AgentChatResponse(
-      extractText(aiMessage.content),
-      task.extraState.sources,
-    );
+    return new AgentChatResponse(aiMessage.content, task.extraState.sources);
   }
 
   private async _getStreamAiResponse(
     task: Task,
-    llmChatParams: LLMChatParamsBase<OpenAIAdditionalChatOptions>,
-  ): Promise<StreamingAgentChatResponse | AgentChatResponse> {
+    llmChatKwargs: any,
+  ): Promise<StreamingAgentChatResponse> {
     const stream = await this.llm.chat({
       stream: true,
-      ...llmChatParams,
+      ...llmChatKwargs,
     });
-    // read first chunk from stream to find out if we need to call tools
-    const iterator = stream[Symbol.asyncIterator]();
-    let { value } = await iterator.next();
-    let content = value.delta;
-    const hasToolCalls = value.additionalKwargs?.toolCalls.length > 0;
 
-    if (hasToolCalls) {
-      // consume stream until we have all the tool calls and return a non-streamed response
-      for await (value of stream) {
-        content += value.delta;
-      }
-      return this._processMessage(task, {
-        content,
-        role: "assistant",
-        additionalKwargs: value.additionalKwargs,
-      });
-    }
-
-    const newStream = streamConverter.bind(this)(
+    const iterator = streamConverter(
       streamReducer({
         stream,
-        initialValue: content,
+        initialValue: "",
         reducer: (accumulator, part) => (accumulator += part.delta),
         finished: (accumulator) => {
           task.extraState.newMemory.put({
@@ -213,31 +219,43 @@ export class OpenAIAgentWorker
       (r: ChatResponseChunk) => new Response(r.delta),
     );
 
-    return new StreamingAgentChatResponse(newStream, task.extraState.sources);
+    return new StreamingAgentChatResponse(iterator, 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,
-    llmChatParams: LLMChatParamsBase<OpenAIAdditionalChatOptions>,
+    llmChatKwargs: any,
   ): Promise<AgentChatResponse | StreamingAgentChatResponse> {
     if (mode === ChatResponseMode.WAIT) {
-      const chatResponse = await this.llm.chat({
+      const chatResponse = (await this.llm.chat({
         stream: false,
-        ...llmChatParams,
-      });
+        ...llmChatKwargs,
+      })) as unknown as ChatResponse;
 
-      return this._processMessage(
-        task,
-        chatResponse.message,
-      ) as AgentChatResponse;
+      return this._processMessage(task, chatResponse) as AgentChatResponse;
     } else if (mode === ChatResponseMode.STREAM) {
-      return this._getStreamAiResponse(task, llmChatParams);
+      return this._getStreamAiResponse(task, llmChatKwargs);
     }
 
     throw new Error("Invalid mode");
   }
 
+  /**
+   * Call function.
+   * @param tools: tools
+   * @param toolCall: tool call
+   * @param memory: memory
+   * @param sources: sources
+   * @returns: void
+   */
   async callFunction(
     tools: BaseTool[],
     toolCall: OpenAIToolCall,
@@ -248,7 +266,7 @@ export class OpenAIAgentWorker
       throw new Error("Invalid tool_call object");
     }
 
-    const functionMessage = await callFunction(tools, toolCall);
+    const functionMessage = await callFunction(tools, toolCall, this.verbose);
 
     const message = functionMessage[0];
     const toolOutput = functionMessage[1];
@@ -259,12 +277,16 @@ export class OpenAIAgentWorker
     };
   }
 
-  initializeStep(task: Task): TaskStep {
+  /**
+   * Initialize step.
+   * @param task: task
+   * @param kwargs: kwargs
+   * @returns: task step
+   */
+  initializeStep(task: Task, kwargs?: any): TaskStep {
     const sources: ToolOutput[] = [];
 
-    const newMemory = new ChatMemoryBuffer({
-      tokenLimit: task.memory.tokenLimit,
-    });
+    const newMemory = new ChatMemoryBuffer();
 
     const taskState = {
       sources,
@@ -280,6 +302,12 @@ export class OpenAIAgentWorker
     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,
@@ -295,6 +323,11 @@ export class OpenAIAgentWorker
     return true;
   }
 
+  /**
+   * Get tools.
+   * @param input: input
+   * @returns: tools
+   */
   async getTools(input: string): Promise<BaseTool[]> {
     return this._getTools(input);
   }
@@ -303,20 +336,28 @@ export class OpenAIAgentWorker
     step: TaskStep,
     task: Task,
     mode: ChatResponseMode = ChatResponseMode.WAIT,
-    toolChoice: ChatCompletionToolChoiceOption = "auto",
+    toolChoice: string | { [key: string]: any } = "auto",
   ): Promise<TaskStepOutput> {
     const tools = await this.getTools(task.input);
 
     if (step.input) {
-      addUserStepToMemory(step, task.extraState.newMemory);
+      addUserStepToMemory(step, task.extraState.newMemory, this.verbose);
     }
 
-    const llmChatParams = this._getLlmChatParams(task, tools, toolChoice);
+    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,
-      llmChatParams,
+      llmChatKwargs,
     );
 
     const latestToolCalls = this.getLatestToolCalls(task) || [];
@@ -349,25 +390,45 @@ export class OpenAIAgentWorker
     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,
-    chatParams: LLMChatParamsBase<OpenAIAdditionalChatOptions>,
+    kwargs?: any,
   ): Promise<TaskStepOutput> {
-    const toolChoice = chatParams?.additionalChatOptions?.tool_choice ?? "auto";
+    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,
-    chatParams: LLMChatParamsBase<OpenAIAdditionalChatOptions>,
+    kwargs?: any,
   ): Promise<TaskStepOutput> {
-    const toolChoice = chatParams?.additionalChatOptions?.tool_choice ?? "auto";
+    const toolChoice = kwargs?.toolChoice || "auto";
     return this._runStep(step, task, ChatResponseMode.STREAM, toolChoice);
   }
 
-  finalizeTask(task: Task): void {
+  /**
+   * 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/react/base.ts

@@ -9,6 +9,7 @@ type ReActAgentParams = {
   llm?: LLM;
   memory?: any;
   prefixMessages?: ChatMessage[];
+  verbose?: boolean;
   maxInteractions?: number;
   defaultToolChoice?: string;
   toolRetriever?: ObjectRetriever;
@@ -25,6 +26,7 @@ export class ReActAgent extends AgentRunner {
     llm,
     memory,
     prefixMessages,
+    verbose,
     maxInteractions = 10,
     defaultToolChoice = "auto",
     toolRetriever,
@@ -34,6 +36,7 @@ export class ReActAgent extends AgentRunner {
       llm,
       maxInteractions,
       toolRetriever,
+      verbose,
     });
 
     super({

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

@@ -1,5 +1,4 @@
 import type { ChatMessage } from "../../llm/index.js";
-import { extractText } from "../../llm/utils.js";
 
 export interface BaseReasoningStep {
   getContent(): string;
@@ -52,12 +51,10 @@ export abstract class BaseOutputParser {
   formatMessages(messages: ChatMessage[]): ChatMessage[] {
     if (messages) {
       if (messages[0].role === "system") {
-        messages[0].content = this.format(
-          extractText(messages[0].content) || "",
-        );
+        messages[0].content = this.format(messages[0].content || "");
       } else {
         messages[messages.length - 1].content = this.format(
-          extractText(messages[messages.length - 1].content) || "",
+          messages[messages.length - 1].content || "",
         );
       }
     }

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

@@ -1,9 +1,7 @@
 import { randomUUID } from "@llamaindex/env";
-import type { ChatMessage } from "cohere-ai/api";
-import { Settings } from "../../Settings.js";
 import { AgentChatResponse } from "../../engines/chat/index.js";
-import { type ChatResponse, type LLM } from "../../llm/index.js";
-import { extractText } from "../../llm/utils.js";
+import type { ChatResponse, LLM } from "../../llm/index.js";
+import { OpenAI } from "../../llm/index.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
 import type { ObjectRetriever } from "../../objects/base.js";
 import { ToolOutput } from "../../tools/index.js";
@@ -18,24 +16,32 @@ import {
   ObservationReasoningStep,
   ResponseReasoningStep,
 } from "./types.js";
-
 type ReActAgentWorkerParams = {
   tools: BaseTool[];
   llm?: LLM;
   maxInteractions?: number;
   reactChatFormatter?: ReActChatFormatter | undefined;
   outputParser?: ReActOutputParser | undefined;
+  verbose?: boolean | undefined;
   toolRetriever?: ObjectRetriever | undefined;
 };
 
+/**
+ *
+ * @param step
+ * @param memory
+ * @param currentReasoning
+ * @param verbose
+ */
 function addUserStepToReasoning(
   step: TaskStep,
   memory: ChatMemoryBuffer,
   currentReasoning: BaseReasoningStep[],
+  verbose: boolean = false,
 ): void {
   if (step.stepState.isFirst) {
     memory.put({
-      content: step.input ?? "",
+      content: step.input,
       role: "user",
     });
     step.stepState.isFirst = false;
@@ -44,22 +50,18 @@ function addUserStepToReasoning(
       observation: step.input ?? undefined,
     });
     currentReasoning.push(reasoningStep);
-    if (Settings.debug) {
+    if (verbose) {
       console.log(`Added user message to memory: ${step.input}`);
     }
   }
 }
 
-type ChatParams = {
-  messages: ChatMessage[];
-  tools?: BaseTool[];
-};
-
 /**
  * ReAct agent worker.
  */
-export class ReActAgentWorker implements AgentWorker<ChatParams> {
+export class ReActAgentWorker implements AgentWorker {
   llm: LLM;
+  verbose: boolean;
 
   maxInteractions: number = 10;
   reactChatFormatter: ReActChatFormatter;
@@ -73,13 +75,15 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
     maxInteractions,
     reactChatFormatter,
     outputParser,
+    verbose,
     toolRetriever,
   }: ReActAgentWorkerParams) {
-    this.llm = llm ?? Settings.llm;
+    this.llm = llm ?? new OpenAI({ model: "gpt-3.5-turbo-0613" });
 
     this.maxInteractions = maxInteractions ?? 10;
     this.reactChatFormatter = reactChatFormatter ?? new ReActChatFormatter();
     this.outputParser = outputParser ?? new ReActOutputParser();
+    this.verbose = verbose || false;
 
     if (tools.length > 0 && toolRetriever) {
       throw new Error("Cannot specify both tools and tool_retriever");
@@ -93,12 +97,16 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
     }
   }
 
-  initializeStep(task: Task): TaskStep {
+  /**
+   * Initialize a task step.
+   * @param task - task
+   * @param kwargs - keyword arguments
+   * @returns - task step
+   */
+  initializeStep(task: Task, kwargs?: any): TaskStep {
     const sources: ToolOutput[] = [];
     const currentReasoning: BaseReasoningStep[] = [];
-    const newMemory = new ChatMemoryBuffer({
-      tokenLimit: task.memory.tokenLimit,
-    });
+    const newMemory = new ChatMemoryBuffer();
 
     const taskState = {
       sources,
@@ -116,6 +124,12 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
     });
   }
 
+  /**
+   * Extract reasoning step from chat response.
+   * @param output - chat response
+   * @param isStreaming - whether the chat response is streaming
+   * @returns - [message content, reasoning steps, is done]
+   */
   extractReasoningStep(
     output: ChatResponse,
     isStreaming: boolean,
@@ -131,21 +145,21 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
 
     try {
       reasoningStep = this.outputParser.parse(
-        extractText(messageContent),
+        messageContent,
         isStreaming,
       ) as ActionReasoningStep;
     } catch (e) {
       throw new Error(`Could not parse output: ${e}`);
     }
 
-    if (Settings.debug) {
+    if (this.verbose) {
       console.log(`${reasoningStep.getContent()}\n`);
     }
 
     currentReasoning.push(reasoningStep);
 
     if (reasoningStep.isDone()) {
-      return [extractText(messageContent), currentReasoning, true];
+      return [messageContent, currentReasoning, true];
     }
 
     const actionReasoningStep = new ActionReasoningStep({
@@ -158,9 +172,17 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
       throw new Error(`Expected ActionReasoningStep, got ${reasoningStep}`);
     }
 
-    return [extractText(messageContent), currentReasoning, false];
+    return [messageContent, currentReasoning, false];
   }
 
+  /**
+   * Process actions.
+   * @param task - task
+   * @param tools - tools
+   * @param output - chat response
+   * @param isStreaming - whether the chat response is streaming
+   * @returns - [reasoning steps, is done]
+   */
   async _processActions(
     task: Task,
     tools: BaseTool[],
@@ -211,13 +233,19 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
 
     currentReasoning.push(observationStep);
 
-    if (Settings.debug) {
+    if (this.verbose) {
       console.log(`${observationStep.getContent()}`);
     }
 
     return [currentReasoning, false];
   }
 
+  /**
+   * Get response.
+   * @param currentReasoning - current reasoning steps
+   * @param sources - tool outputs
+   * @returns - agent chat response
+   */
   _getResponse(
     currentReasoning: BaseReasoningStep[],
     sources: ToolOutput[],
@@ -241,6 +269,13 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
     return new AgentChatResponse(responseStr, sources);
   }
 
+  /**
+   * Get task step response.
+   * @param agentResponse - agent chat response
+   * @param step - task step
+   * @param isDone - whether the task is done
+   * @returns - task step output
+   */
   _getTaskStepResponse(
     agentResponse: AgentChatResponse,
     step: TaskStep,
@@ -257,12 +292,24 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
     return new TaskStepOutput(agentResponse, step, newSteps, isDone);
   }
 
-  async _runStep(step: TaskStep, task: Task): Promise<TaskStepOutput> {
+  /**
+   * Run a task step.
+   * @param step - task step
+   * @param task - task
+   * @param kwargs - keyword arguments
+   * @returns - task step output
+   */
+  async _runStep(
+    step: TaskStep,
+    task: Task,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
     if (step.input) {
       addUserStepToReasoning(
         step,
         task.extraState.newMemory,
         task.extraState.currentReasoning,
+        this.verbose,
       );
     }
 
@@ -301,15 +348,42 @@ export class ReActAgentWorker implements AgentWorker<ChatParams> {
     return this._getTaskStepResponse(agentResponse, step, isDone);
   }
 
-  async runStep(step: TaskStep, task: Task): Promise<TaskStepOutput> {
+  /**
+   * Run a task step.
+   * @param step - task step
+   * @param task - task
+   * @param kwargs - keyword arguments
+   * @returns - task step output
+   */
+  async runStep(
+    step: TaskStep,
+    task: Task,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
     return await this._runStep(step, task);
   }
 
-  streamStep(): Promise<TaskStepOutput> {
+  /**
+   * Run a task step.
+   * @param step - task step
+   * @param task - task
+   * @param kwargs - keyword arguments
+   * @returns - task step output
+   */
+  streamStep(
+    step: TaskStep,
+    task: Task,
+    kwargs?: any,
+  ): Promise<TaskStepOutput> {
     throw new Error("Method not implemented.");
   }
 
-  finalizeTask(task: Task): void {
+  /**
+   * Finalize a task.
+   * @param task - task
+   * @param kwargs - keyword arguments
+   */
+  finalizeTask(task: Task, kwargs?: any): void {
     task.memory.set(task.memory.get() + task.extraState.newMemory.get());
     task.extraState.newMemory.reset();
   }

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

@@ -58,7 +58,6 @@ export class AgentRunner extends BaseAgentRunner {
     this.memory =
       params.memory ??
       new ChatMemoryBuffer({
-        llm: params.llm,
         chatHistory: params.chatHistory,
       });
     this.initTaskStateKwargs = params.initTaskStateKwargs ?? {};

packages/core/src/agent/types.ts

@@ -6,19 +6,11 @@ import type {
 
 import type { QueryEngineParamsNonStreaming } from "../types.js";
 
-export interface AgentWorker<ExtraParams extends object = object> {
-  initializeStep(task: Task, params?: ExtraParams): TaskStep;
-  runStep(
-    step: TaskStep,
-    task: Task,
-    params?: ExtraParams,
-  ): Promise<TaskStepOutput>;
-  streamStep(
-    step: TaskStep,
-    task: Task,
-    params?: ExtraParams,
-  ): Promise<TaskStepOutput>;
-  finalizeTask(task: Task, params?: ExtraParams): void;
+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 {
@@ -178,13 +170,13 @@ export class TaskStep implements ITaskStep {
  * @param isLast: isLast
  */
 export class TaskStepOutput {
-  output: AgentChatResponse | StreamingAgentChatResponse;
+  output: any;
   taskStep: TaskStep;
   nextSteps: TaskStep[];
   isLast: boolean;
 
   constructor(
-    output: AgentChatResponse | StreamingAgentChatResponse,
+    output: any,
     taskStep: TaskStep,
     nextSteps: TaskStep[],
     isLast: boolean = false,

packages/core/src/agent/utils.ts

@@ -1,12 +1,19 @@
-import { Settings } from "../Settings.js";
 import type { ChatMessage } from "../llm/index.js";
 import type { ChatMemoryBuffer } from "../memory/ChatMemoryBuffer.js";
 import type { BaseTool } from "../types.js";
 import type { TaskStep } from "./types.js";
 
+/**
+ * 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;
@@ -19,17 +26,26 @@ export function addUserStepToMemory(
 
   memory.put(userMessage);
 
-  if (Settings.debug) {
+  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 exist = tools.find((tool) => tool.metadata.name === name);
+  const nameToTool: { [key: string]: BaseTool } = {};
+  tools.forEach((tool) => {
+    nameToTool[tool.metadata.name] = tool;
+  });
 
-  if (!exist) {
+  if (!(name in nameToTool)) {
     throw new Error(`Tool with name ${name} not found`);
   }
 
-  return exist;
+  return nameToTool[name];
 }

packages/core/src/callbacks/CallbackManager.ts

@@ -1,52 +1,46 @@
 import type { Anthropic } from "@anthropic-ai/sdk";
 import { CustomEvent } from "@llamaindex/env";
 import type { NodeWithScore } from "../Node.js";
-import {
-  EventCaller,
-  getEventCaller,
-} from "../internal/context/EventCaller.js";
-import type { LLMEndEvent, LLMStartEvent } from "../llm/types.js";
-
-export class LlamaIndexCustomEvent<T = any> extends CustomEvent<T> {
-  reason: EventCaller | null;
-  private constructor(
-    event: string,
-    options?: CustomEventInit & {
-      reason?: EventCaller | null;
-    },
-  ) {
-    super(event, options);
-    this.reason = options?.reason ?? null;
-  }
-
-  static fromEvent<Type extends keyof LlamaIndexEventMaps>(
-    type: Type,
-    detail: LlamaIndexEventMaps[Type]["detail"],
-  ) {
-    return new LlamaIndexCustomEvent(type, {
-      detail: detail,
-      reason: getEventCaller(),
-    });
-  }
-}
 
 /**
- * This type is used to define the event maps.
+ * This type is used to define the event maps for the Llamaindex package.
  */
-export interface LlamaIndexEventMaps {
-  /**
-   * @deprecated
-   */
-  retrieve: CustomEvent<RetrievalCallbackResponse>;
-  /**
-   * @deprecated
-   */
-  stream: CustomEvent<StreamCallbackResponse>;
-  "llm-start": LLMStartEvent;
-  "llm-end": LLMEndEvent;
+export interface LlamaIndexEventMaps {}
+
+declare module "llamaindex" {
+  interface LlamaIndexEventMaps {
+    /**
+     * @deprecated
+     */
+    retrieve: CustomEvent<RetrievalCallbackResponse>;
+    /**
+     * @deprecated
+     */
+    stream: CustomEvent<StreamCallbackResponse>;
+  }
 }
 
 //#region @deprecated remove in the next major version
+/*
+  An event is a wrapper that groups related operations.
+  For example, during retrieve and synthesize,
+  a parent event wraps both operations, and each operation has it's own
+  event. In this case, both sub-events will share a parentId.
+*/
+
+export type EventTag = "intermediate" | "final";
+export type EventType = "retrieve" | "llmPredict" | "wrapper";
+export interface Event {
+  id: string;
+  type: EventType;
+  tags?: EventTag[];
+  parentId?: string;
+}
+
+interface BaseCallbackResponse {
+  event: Event;
+}
+
 //Specify StreamToken per mainstream LLM
 export interface DefaultStreamToken {
   id: string;
@@ -74,13 +68,13 @@ export type AnthropicStreamToken = Anthropic.Completion;
 
 //StreamCallbackResponse should let practitioners implement callbacks out of the box...
 //When custom streaming LLMs are involved, people are expected to write their own StreamCallbackResponses
-export interface StreamCallbackResponse {
+export interface StreamCallbackResponse extends BaseCallbackResponse {
   index: number;
   isDone?: boolean;
   token?: DefaultStreamToken;
 }
 
-export interface RetrievalCallbackResponse {
+export interface RetrievalCallbackResponse extends BaseCallbackResponse {
   query: string;
   nodes: NodeWithScore[];
 }
@@ -104,11 +98,7 @@ interface CallbackManagerMethods {
 
 const noop: (...args: any[]) => any = () => void 0;
 
-type EventHandler<Event extends CustomEvent> = (
-  event: Event & {
-    reason: EventCaller | null;
-  },
-) => void;
+type EventHandler<Event extends CustomEvent> = (event: Event) => void;
 
 export class CallbackManager implements CallbackManagerMethods {
   /**
@@ -120,7 +110,7 @@ export class CallbackManager implements CallbackManagerMethods {
         this.#handlers
           .get("stream")!
           .map((handler) =>
-            handler(LlamaIndexCustomEvent.fromEvent("stream", response)),
+            handler(new CustomEvent("stream", { detail: response })),
           ),
       );
     };
@@ -135,7 +125,7 @@ export class CallbackManager implements CallbackManagerMethods {
         this.#handlers
           .get("retrieve")!
           .map((handler) =>
-            handler(LlamaIndexCustomEvent.fromEvent("retrieve", response)),
+            handler(new CustomEvent("retrieve", { detail: response })),
           ),
       );
     };
@@ -198,8 +188,6 @@ export class CallbackManager implements CallbackManagerMethods {
     if (!handlers) {
       return;
     }
-    handlers.forEach((handler) =>
-      handler(LlamaIndexCustomEvent.fromEvent(event, detail)),
-    );
+    handlers.forEach((handler) => handler(new CustomEvent(event, { detail })));
   }
 }

packages/core/src/cloud/LlamaCloudIndex.ts

@@ -1,20 +1,11 @@
-import { PlatformApi } from "@llamaindex/cloud";
-import type { Document } from "../Node.js";
 import type { BaseRetriever } from "../Retriever.js";
 import { RetrieverQueryEngine } from "../engines/query/RetrieverQueryEngine.js";
-import type { TransformComponent } from "../ingestion/types.js";
 import type { BaseNodePostprocessor } from "../postprocessors/types.js";
 import type { BaseSynthesizer } from "../synthesizers/types.js";
 import type { BaseQueryEngine } from "../types.js";
 import type { CloudRetrieveParams } from "./LlamaCloudRetriever.js";
 import { LlamaCloudRetriever } from "./LlamaCloudRetriever.js";
-import { getPipelineCreate } from "./config.js";
 import type { CloudConstructorParams } from "./types.js";
-import { getAppBaseUrl, getClient } from "./utils.js";
-
-import { getEnv } from "@llamaindex/env";
-import { OpenAIEmbedding } from "../embeddings/OpenAIEmbedding.js";
-import { SimpleNodeParser } from "../nodeParsers/SimpleNodeParser.js";
 
 export class LlamaCloudIndex {
   params: CloudConstructorParams;
@@ -23,151 +14,6 @@ export class LlamaCloudIndex {
     this.params = params;
   }
 
-  static async fromDocuments(
-    params: {
-      documents: Document[];
-      transformations?: TransformComponent[];
-      verbose?: boolean;
-    } & CloudConstructorParams,
-  ): Promise<LlamaCloudIndex> {
-    const defaultTransformations: TransformComponent[] = [
-      new OpenAIEmbedding({
-        apiKey: getEnv("OPENAI_API_KEY"),
-      }),
-      new SimpleNodeParser(),
-    ];
-
-    const appUrl = getAppBaseUrl(params.baseUrl);
-
-    const client = await getClient({ ...params, baseUrl: appUrl });
-
-    const pipelineCreateParams = await getPipelineCreate({
-      pipelineName: params.name,
-      pipelineType: "MANAGED",
-      inputNodes: params.documents,
-      transformations: params.transformations ?? defaultTransformations,
-    });
-
-    const project = await client.project.upsertProject({
-      name: params.projectName ?? "default",
-    });
-
-    if (!project.id) {
-      throw new Error("Project ID should be defined");
-    }
-
-    const pipeline = await client.project.upsertPipelineForProject(
-      project.id,
-      pipelineCreateParams,
-    );
-
-    if (!pipeline.id) {
-      throw new Error("Pipeline ID must be defined");
-    }
-
-    if (params.verbose) {
-      console.log(`Created pipeline ${pipeline.id} with name ${params.name}`);
-    }
-
-    const executionsIds: {
-      exectionId: string;
-      dataSourceId: string;
-    }[] = [];
-
-    for (const dataSource of pipeline.dataSources) {
-      const dataSourceExection =
-        await client.dataSource.createDataSourceExecution(dataSource.id);
-
-      if (!dataSourceExection.id) {
-        throw new Error("Data Source Execution ID must be defined");
-      }
-
-      executionsIds.push({
-        exectionId: dataSourceExection.id,
-        dataSourceId: dataSource.id,
-      });
-    }
-
-    let isDone = false;
-
-    while (!isDone) {
-      const statuses = [];
-
-      for await (const execution of executionsIds) {
-        const dataSourceExecution =
-          await client.dataSource.getDataSourceExecution(
-            execution.dataSourceId,
-            execution.exectionId,
-          );
-
-        statuses.push(dataSourceExecution.status);
-
-        if (
-          statuses.every((status) => status === PlatformApi.StatusEnum.Success)
-        ) {
-          isDone = true;
-          if (params.verbose) {
-            console.info("Data Source Execution completed");
-          }
-          break;
-        } else if (
-          statuses.some((status) => status === PlatformApi.StatusEnum.Error)
-        ) {
-          throw new Error("Data Source Execution failed");
-        } else {
-          await new Promise((resolve) => setTimeout(resolve, 1000));
-          if (params.verbose) {
-            process.stdout.write(".");
-          }
-        }
-      }
-    }
-
-    isDone = false;
-
-    const execution = await client.pipeline.runManagedPipelineIngestion(
-      pipeline.id,
-    );
-
-    const ingestionId = execution.id;
-
-    if (!ingestionId) {
-      throw new Error("Ingestion ID must be defined");
-    }
-
-    while (!isDone) {
-      const pipelineStatus = await client.pipeline.getManagedIngestionExecution(
-        pipeline.id,
-        ingestionId,
-      );
-
-      if (pipelineStatus.status === PlatformApi.StatusEnum.Success) {
-        isDone = true;
-
-        if (params.verbose) {
-          console.info("Ingestion completed");
-        }
-
-        break;
-      } else if (pipelineStatus.status === PlatformApi.StatusEnum.Error) {
-        throw new Error("Ingestion failed");
-      } else {
-        await new Promise((resolve) => setTimeout(resolve, 1000));
-        if (params.verbose) {
-          process.stdout.write(".");
-        }
-      }
-    }
-
-    if (params.verbose) {
-      console.info(
-        `Ingestion completed, find your index at ${appUrl}/project/${project.id}/deploy/${pipeline.id}`,
-      );
-    }
-
-    return new LlamaCloudIndex({ ...params });
-  }
-
   asRetriever(params: CloudRetrieveParams = {}): BaseRetriever {
     return new LlamaCloudRetriever({ ...this.params, ...params });
   }

packages/core/src/cloud/LlamaCloudRetriever.ts

@@ -1,12 +1,13 @@
 import type { PlatformApi, PlatformApiClient } from "@llamaindex/cloud";
+import { globalsHelper } from "../GlobalsHelper.js";
 import type { NodeWithScore } from "../Node.js";
 import { ObjectType, jsonToNode } from "../Node.js";
 import type { BaseRetriever, RetrieveParams } from "../Retriever.js";
 import { Settings } from "../Settings.js";
-import { wrapEventCaller } from "../internal/context/EventCaller.js";
 import type { ClientParams, CloudConstructorParams } from "./types.js";
 import { DEFAULT_PROJECT_NAME } from "./types.js";
 import { getClient } from "./utils.js";
+
 export type CloudRetrieveParams = Omit<
   PlatformApi.RetrievalParams,
   "query" | "searchFilters" | "pipelineId" | "className"
@@ -50,9 +51,9 @@ export class LlamaCloudRetriever implements BaseRetriever {
     return this.client;
   }
 
-  @wrapEventCaller
   async retrieve({
     query,
+    parentEvent,
     preFilters,
   }: RetrieveParams): Promise<NodeWithScore[]> {
     const pipelines = await (
@@ -76,9 +77,13 @@ export class LlamaCloudRetriever implements BaseRetriever {
 
     const nodes = this.resultNodesToNodeWithScore(results.retrievalNodes);
 
-    Settings.callbackManager.dispatchEvent("retrieve", {
+    Settings.callbackManager.onRetrieve({
       query,
       nodes,
+      event: globalsHelper.createEvent({
+        parentEvent,
+        type: "retrieve",
+      }),
     });
 
     return nodes;

packages/core/src/cloud/config.ts

@@ -18,11 +18,11 @@ function getTransformationConfig(
     return {
       configurableTransformationType: "SENTENCE_AWARE_NODE_PARSER",
       component: {
-        // TODO: API doesnt accept camelCase
-        chunk_size: transformation.textSplitter.chunkSize, // TODO: set to public in SentenceSplitter
-        chunk_overlap: transformation.textSplitter.chunkOverlap, // TODO: set to public in SentenceSplitter
-        include_metadata: transformation.includeMetadata,
-        include_prev_next_rel: transformation.includePrevNextRel,
+        // TODO: API returns 422 if these parameters are included
+        // chunkSize: transformation.textSplitter.chunkSize, // TODO: set to public in SentenceSplitter
+        // chunkOverlap: transformation.textSplitter.chunkOverlap, // TODO: set to public in SentenceSplitter
+        // includeMetadata: transformation.includeMetadata,
+        // includePrevNextRel: transformation.includePrevNextRel,
       },
     };
   }
@@ -30,10 +30,9 @@ function getTransformationConfig(
     return {
       configurableTransformationType: "OPENAI_EMBEDDING",
       component: {
-        // TODO: API doesnt accept camelCase
-        model: transformation.model,
-        api_key: transformation.apiKey,
-        embed_batch_size: transformation.embedBatchSize,
+        modelName: transformation.model,
+        apiKey: transformation.apiKey,
+        embedBatchSize: transformation.embedBatchSize,
         dimensions: transformation.dimensions,
       },
     };
@@ -72,12 +71,10 @@ export async function getPipelineCreate(
     inputNodes = [],
   } = params;
 
-  const dataSources = inputNodes.map(getDataSourceConfig);
-
   return {
     name: pipelineName,
     configuredTransformations: transformations.map(getTransformationConfig),
-    dataSources,
+    dataSources: inputNodes.map(getDataSourceConfig),
     dataSinks: [],
     pipelineType,
   };

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

@@ -8,7 +8,6 @@ import {
 import type { Response } from "../../Response.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { llmFromSettingsOrContext } from "../../Settings.js";
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { ChatMessage, LLM } from "../../llm/index.js";
 import { extractText, streamReducer } from "../../llm/utils.js";
 import { PromptMixin } from "../../prompts/index.js";
@@ -18,6 +17,7 @@ import type {
   ChatEngineParamsNonStreaming,
   ChatEngineParamsStreaming,
 } from "./types.js";
+
 /**
  * CondenseQuestionChatEngine is used in conjunction with a Index (for example VectorStoreIndex).
  * It does two steps on taking a user's chat message: first, it condenses the chat message
@@ -82,7 +82,6 @@ export class CondenseQuestionChatEngine
 
   chat(params: ChatEngineParamsStreaming): Promise<AsyncIterable<Response>>;
   chat(params: ChatEngineParamsNonStreaming): Promise<Response>;
-  @wrapEventCaller
   async chat(
     params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
   ): Promise<Response | AsyncIterable<Response>> {

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

@@ -1,9 +1,10 @@
+import { randomUUID } from "@llamaindex/env";
 import type { ChatHistory } from "../../ChatHistory.js";
 import { getHistory } from "../../ChatHistory.js";
 import type { ContextSystemPrompt } from "../../Prompt.js";
 import { Response } from "../../Response.js";
 import type { BaseRetriever } from "../../Retriever.js";
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
+import type { Event } from "../../callbacks/CallbackManager.js";
 import type { ChatMessage, ChatResponseChunk, LLM } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
 import type { MessageContent } from "../../llm/types.js";
@@ -59,7 +60,6 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
 
   chat(params: ChatEngineParamsStreaming): Promise<AsyncIterable<Response>>;
   chat(params: ChatEngineParamsNonStreaming): Promise<Response>;
-  @wrapEventCaller
   async chat(
     params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
   ): Promise<Response | AsyncIterable<Response>> {
@@ -67,14 +67,21 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     const chatHistory = params.chatHistory
       ? getHistory(params.chatHistory)
       : this.chatHistory;
+    const parentEvent: Event = {
+      id: randomUUID(),
+      type: "wrapper",
+      tags: ["final"],
+    };
     const requestMessages = await this.prepareRequestMessages(
       message,
       chatHistory,
+      parentEvent,
     );
 
     if (stream) {
       const stream = await this.chatModel.chat({
         messages: requestMessages.messages,
+        parentEvent,
         stream: true,
       });
       return streamConverter(
@@ -91,12 +98,10 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     }
     const response = await this.chatModel.chat({
       messages: requestMessages.messages,
+      parentEvent,
     });
     chatHistory.addMessage(response.message);
-    return new Response(
-      extractText(response.message.content),
-      requestMessages.nodes,
-    );
+    return new Response(response.message.content, requestMessages.nodes);
   }
 
   reset() {
@@ -106,13 +111,14 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
   private async prepareRequestMessages(
     message: MessageContent,
     chatHistory: ChatHistory,
+    parentEvent?: Event,
   ) {
     chatHistory.addMessage({
       content: message,
       role: "user",
     });
     const textOnly = extractText(message);
-    const context = await this.contextGenerator.generate(textOnly);
+    const context = await this.contextGenerator.generate(textOnly, parentEvent);
     const nodes = context.nodes.map((r) => r.node);
     const messages = await chatHistory.requestMessages(
       context ? [context.message] : undefined,

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

@@ -1,7 +1,9 @@
+import { randomUUID } from "@llamaindex/env";
 import type { NodeWithScore, TextNode } from "../../Node.js";
 import type { ContextSystemPrompt } from "../../Prompt.js";
 import { defaultContextSystemPrompt } from "../../Prompt.js";
 import type { BaseRetriever } from "../../Retriever.js";
+import type { Event } from "../../callbacks/CallbackManager.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import { PromptMixin } from "../../prompts/index.js";
 import type { Context, ContextGenerator } from "./types.js";
@@ -54,9 +56,17 @@ export class DefaultContextGenerator
     return nodesWithScore;
   }
 
-  async generate(message: string): Promise<Context> {
+  async generate(message: string, parentEvent?: Event): Promise<Context> {
+    if (!parentEvent) {
+      parentEvent = {
+        id: randomUUID(),
+        type: "wrapper",
+        tags: ["final"],
+      };
+    }
     const sourceNodesWithScore = await this.retriever.retrieve({
       query: message,
+      parentEvent,
     });
 
     const nodes = await this.applyNodePostprocessors(

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

@@ -1,14 +1,9 @@
 import type { ChatHistory } from "../../ChatHistory.js";
 import { getHistory } from "../../ChatHistory.js";
 import { Response } from "../../Response.js";
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { ChatResponseChunk, LLM } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
-import {
-  extractText,
-  streamConverter,
-  streamReducer,
-} from "../../llm/utils.js";
+import { streamConverter, streamReducer } from "../../llm/utils.js";
 import type {
   ChatEngine,
   ChatEngineParamsNonStreaming,
@@ -30,7 +25,6 @@ export class SimpleChatEngine implements ChatEngine {
 
   chat(params: ChatEngineParamsStreaming): Promise<AsyncIterable<Response>>;
   chat(params: ChatEngineParamsNonStreaming): Promise<Response>;
-  @wrapEventCaller
   async chat(
     params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
   ): Promise<Response | AsyncIterable<Response>> {
@@ -50,7 +44,7 @@ export class SimpleChatEngine implements ChatEngine {
         streamReducer({
           stream,
           initialValue: "",
-          reducer: (accumulator, part) => accumulator + part.delta,
+          reducer: (accumulator, part) => (accumulator += part.delta),
           finished: (accumulator) => {
             chatHistory.addMessage({ content: accumulator, role: "assistant" });
           },
@@ -63,7 +57,7 @@ export class SimpleChatEngine implements ChatEngine {
       messages: await chatHistory.requestMessages(),
     });
     chatHistory.addMessage(response.message);
-    return new Response(extractText(response.message.content));
+    return new Response(response.message.content);
   }
 
   reset() {

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

@@ -1,6 +1,7 @@
 import type { ChatHistory } from "../../ChatHistory.js";
 import type { BaseNode, NodeWithScore } from "../../Node.js";
 import type { Response } from "../../Response.js";
+import type { Event } from "../../callbacks/CallbackManager.js";
 import type { ChatMessage } from "../../llm/index.js";
 import type { MessageContent } from "../../llm/types.js";
 import type { ToolOutput } from "../../tools/types.js";
@@ -55,7 +56,7 @@ export interface Context {
  * A ContextGenerator is used to generate a context based on a message's text content
  */
 export interface ContextGenerator {
-  generate(message: string): Promise<Context>;
+  generate(message: string, parentEvent?: Event): Promise<Context>;
 }
 
 export enum ChatResponseMode {

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

@@ -1,7 +1,8 @@
+import { randomUUID } from "@llamaindex/env";
 import type { NodeWithScore } from "../../Node.js";
 import type { Response } from "../../Response.js";
 import type { BaseRetriever } from "../../Retriever.js";
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
+import type { Event } from "../../callbacks/CallbackManager.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import { PromptMixin } from "../../prompts/Mixin.js";
 import type { BaseSynthesizer } from "../../synthesizers/index.js";
@@ -61,9 +62,10 @@ export class RetrieverQueryEngine
     return nodesWithScore;
   }
 
-  private async retrieve(query: string) {
+  private async retrieve(query: string, parentEvent: Event) {
     const nodes = await this.retriever.retrieve({
       query,
+      parentEvent,
       preFilters: this.preFilters,
     });
 
@@ -72,22 +74,28 @@ export class RetrieverQueryEngine
 
   query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
   query(params: QueryEngineParamsNonStreaming): Promise<Response>;
-  @wrapEventCaller
   async query(
     params: QueryEngineParamsStreaming | QueryEngineParamsNonStreaming,
   ): Promise<Response | AsyncIterable<Response>> {
     const { query, stream } = params;
-    const nodesWithScore = await this.retrieve(query);
+    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/SubQuestionQueryEngine.ts

@@ -1,8 +1,10 @@
+import { randomUUID } from "@llamaindex/env";
 import type { NodeWithScore } from "../../Node.js";
 import { TextNode } from "../../Node.js";
 import { LLMQuestionGenerator } from "../../QuestionGenerator.js";
 import type { Response } from "../../Response.js";
 import type { ServiceContext } from "../../ServiceContext.js";
+import type { Event } from "../../callbacks/CallbackManager.js";
 import { PromptMixin } from "../../prompts/Mixin.js";
 import type { BaseSynthesizer } from "../../synthesizers/index.js";
 import {
@@ -18,7 +20,6 @@ import type {
   ToolMetadata,
 } from "../../types.js";
 
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { BaseQuestionGenerator, SubQuestion } from "./types.js";
 
 /**
@@ -79,15 +80,29 @@ export class SubQuestionQueryEngine
 
   query(params: QueryEngineParamsStreaming): Promise<AsyncIterable<Response>>;
   query(params: QueryEngineParamsNonStreaming): Promise<Response>;
-  @wrapEventCaller
   async query(
     params: QueryEngineParamsStreaming | QueryEngineParamsNonStreaming,
   ): Promise<Response | AsyncIterable<Response>> {
     const { query, stream } = params;
     const subQuestions = await this.questionGen.generate(this.metadatas, query);
 
+    // groups final retrieval+synthesis operation
+    const parentEvent: Event = params.parentEvent || {
+      id: randomUUID(),
+      type: "wrapper",
+      tags: ["final"],
+    };
+
+    // groups all sub-queries
+    const subQueryParentEvent: Event = {
+      id: randomUUID(),
+      parentId: parentEvent.id,
+      type: "wrapper",
+      tags: ["intermediate"],
+    };
+
     const subQNodes = await Promise.all(
-      subQuestions.map((subQ) => this.querySubQ(subQ)),
+      subQuestions.map((subQ) => this.querySubQ(subQ, subQueryParentEvent)),
     );
 
     const nodesWithScore = subQNodes
@@ -97,16 +112,21 @@ export class SubQuestionQueryEngine
       return this.responseSynthesizer.synthesize({
         query,
         nodesWithScore,
+        parentEvent,
         stream: true,
       });
     }
     return this.responseSynthesizer.synthesize({
       query,
       nodesWithScore,
+      parentEvent,
     });
   }
 
-  private async querySubQ(subQ: SubQuestion): Promise<NodeWithScore | null> {
+  private async querySubQ(
+    subQ: SubQuestion,
+    parentEvent?: Event,
+  ): Promise<NodeWithScore | null> {
     try {
       const question = subQ.subQuestion;
 
@@ -120,6 +140,7 @@ export class SubQuestionQueryEngine
 
       const responseText = await queryEngine?.call?.({
         query: question,
+        parentEvent,
       });
 
       if (!responseText) {

packages/core/src/evaluation/Correctness.ts

@@ -2,7 +2,6 @@ import { MetadataMode } from "../Node.js";
 import type { ServiceContext } from "../ServiceContext.js";
 import { llmFromSettingsOrContext } from "../Settings.js";
 import type { ChatMessage, LLM } from "../llm/types.js";
-import { extractText } from "../llm/utils.js";
 import { PromptMixin } from "../prompts/Mixin.js";
 import type { CorrectnessSystemPrompt } from "./prompts.js";
 import {
@@ -86,7 +85,7 @@ export class CorrectnessEvaluator extends PromptMixin implements BaseEvaluator {
     });
 
     const [score, reasoning] = this.parserFunction(
-      extractText(evalResponse.message.content),
+      evalResponse.message.content,
     );
 
     return {

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

@@ -278,7 +278,7 @@ export class KeywordTableIndex extends BaseIndex<KeywordTable> {
     serviceContext = serviceContext ?? serviceContextFromDefaults({});
     const docStore = storageContext.docStore;
 
-    await docStore.addDocuments(documents, true);
+    docStore.addDocuments(documents, true);
     for (const doc of documents) {
       docStore.setDocumentHash(doc.id_, doc.hash);
     }

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

@@ -1,4 +1,5 @@
 import _ from "lodash";
+import { globalsHelper } from "../../GlobalsHelper.js";
 import type { BaseNode, Document, NodeWithScore } from "../../Node.js";
 import type { ChoiceSelectPrompt } from "../../Prompt.js";
 import { defaultChoiceSelectPrompt } from "../../Prompt.js";
@@ -10,7 +11,6 @@ import {
   nodeParserFromSettingsOrContext,
 } from "../../Settings.js";
 import { RetrieverQueryEngine } from "../../engines/query/index.js";
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import type { StorageContext } from "../../storage/StorageContext.js";
 import { storageContextFromDefaults } from "../../storage/StorageContext.js";
@@ -135,7 +135,7 @@ export class SummaryIndex extends BaseIndex<IndexList> {
     serviceContext = serviceContext;
     const docStore = storageContext.docStore;
 
-    await docStore.addDocuments(documents, true);
+    docStore.addDocuments(documents, true);
     for (const doc of documents) {
       docStore.setDocumentHash(doc.id_, doc.hash);
     }
@@ -287,8 +287,10 @@ export class SummaryIndexRetriever implements BaseRetriever {
     this.index = index;
   }
 
-  @wrapEventCaller
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async retrieve({
+    query,
+    parentEvent,
+  }: RetrieveParams): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const nodes = await this.index.docStore.getNodes(nodeIds);
     const result = nodes.map((node) => ({
@@ -296,9 +298,13 @@ export class SummaryIndexRetriever implements BaseRetriever {
       score: 1,
     }));
 
-    Settings.callbackManager.dispatchEvent("retrieve", {
+    Settings.callbackManager.onRetrieve({
       query,
       nodes: result,
+      event: globalsHelper.createEvent({
+        parentEvent,
+        type: "retrieve",
+      }),
     });
 
     return result;
@@ -334,7 +340,10 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
     this.serviceContext = serviceContext || index.serviceContext;
   }
 
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async retrieve({
+    query,
+    parentEvent,
+  }: RetrieveParams): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const results: NodeWithScore[] = [];
 
@@ -371,9 +380,13 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
       results.push(...nodeWithScores);
     }
 
-    Settings.callbackManager.dispatchEvent("retrieve", {
+    Settings.callbackManager.onRetrieve({
       query,
       nodes: results,
+      event: globalsHelper.createEvent({
+        parentEvent,
+        type: "retrieve",
+      }),
     });
 
     return results;

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

@@ -1,3 +1,4 @@
+import { globalsHelper } from "../../GlobalsHelper.js";
 import type {
   BaseNode,
   Document,
@@ -17,6 +18,7 @@ import {
   embedModelFromSettingsOrContext,
   nodeParserFromSettingsOrContext,
 } from "../../Settings.js";
+import { type Event } from "../../callbacks/CallbackManager.js";
 import { DEFAULT_SIMILARITY_TOP_K } from "../../constants.js";
 import type {
   BaseEmbedding,
@@ -29,7 +31,6 @@ import {
   DocStoreStrategy,
   createDocStoreStrategy,
 } from "../../ingestion/strategies/index.js";
-import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/types.js";
 import type { StorageContext } from "../../storage/StorageContext.js";
 import { storageContextFromDefaults } from "../../storage/StorageContext.js";
@@ -364,7 +365,7 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     vectorStore: VectorStore,
     refDocId: string,
   ): Promise<void> {
-    await vectorStore.delete(refDocId);
+    vectorStore.delete(refDocId);
 
     if (!vectorStore.storesText) {
       const refDocInfo = await this.docStore.getRefDocInfo(refDocId);
@@ -372,7 +373,7 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
       if (refDocInfo) {
         for (const nodeId of refDocInfo.nodeIds) {
           this.indexStruct.delete(nodeId);
-          await vectorStore.delete(nodeId);
+          vectorStore.delete(nodeId);
         }
       }
       await this.indexStore.addIndexStruct(this.indexStruct);
@@ -439,6 +440,7 @@ export class VectorIndexRetriever implements BaseRetriever {
 
   async retrieve({
     query,
+    parentEvent,
     preFilters,
   }: RetrieveParams): Promise<NodeWithScore[]> {
     let nodesWithScores = await this.textRetrieve(
@@ -448,7 +450,7 @@ export class VectorIndexRetriever implements BaseRetriever {
     nodesWithScores = nodesWithScores.concat(
       await this.textToImageRetrieve(query, preFilters as MetadataFilters),
     );
-    this.sendEvent(query, nodesWithScores);
+    this.sendEvent(query, nodesWithScores, parentEvent);
     return nodesWithScores;
   }
 
@@ -485,14 +487,18 @@ export class VectorIndexRetriever implements BaseRetriever {
     return this.buildNodeListFromQueryResult(result);
   }
 
-  @wrapEventCaller
   protected sendEvent(
     query: string,
     nodesWithScores: NodeWithScore<Metadata>[],
+    parentEvent: Event | undefined,
   ) {
-    Settings.callbackManager.dispatchEvent("retrieve", {
+    Settings.callbackManager.onRetrieve({
       query,
       nodes: nodesWithScores,
+      event: globalsHelper.createEvent({
+        parentEvent,
+        type: "retrieve",
+      }),
     });
   }
 

packages/core/src/ingestion/strategies/DuplicatesStrategy.ts

@@ -25,7 +25,7 @@ export class DuplicatesStrategy implements TransformComponent {
       }
     }
 
-    await this.docStore.addDocuments(nodesToRun, true);
+    this.docStore.addDocuments(nodesToRun, true);
 
     return nodesToRun;
   }

packages/core/src/ingestion/strategies/UpsertsStrategy.ts

@@ -26,7 +26,7 @@ export class UpsertsStrategy implements TransformComponent {
       }
     }
     // add non-duplicate docs
-    await this.docStore.addDocuments(dedupedNodes, true);
+    this.docStore.addDocuments(dedupedNodes, true);
     return dedupedNodes;
   }
 }

packages/core/src/ingestion/strategies/index.ts

@@ -5,16 +5,9 @@ import { DuplicatesStrategy } from "./DuplicatesStrategy.js";
 import { UpsertsAndDeleteStrategy } from "./UpsertsAndDeleteStrategy.js";
 import { UpsertsStrategy } from "./UpsertsStrategy.js";
 
-/**
- * Document de-deduplication strategies work by comparing the hashes or ids stored in the document store.
- * They require a document store to be set which must be persisted across pipeline runs.
- */
 export enum DocStoreStrategy {
-  // Use upserts to handle duplicates. Checks if the a document is already in the doc store based on its id. If it is not, or if the hash of the document is updated, it will update the document in the doc store and run the transformations.
   UPSERTS = "upserts",
-  // Only handle duplicates. Checks if the hash of a document is already in the doc store. Only then it will add the document to the doc store and run the transformations
   DUPLICATES_ONLY = "duplicates_only",
-  // Use upserts and delete to handle duplicates. Like the upsert strategy but it will also delete non-existing documents from the doc store
   UPSERTS_AND_DELETE = "upserts_and_delete",
   NONE = "none", // no-op strategy
 }

packages/core/src/internal/context/EventCaller.ts

@@ -1,99 +0,0 @@
-import { AsyncLocalStorage, randomUUID } from "@llamaindex/env";
-import { isAsyncGenerator, isGenerator } from "../utils.js";
-
-const eventReasonAsyncLocalStorage = new AsyncLocalStorage<EventCaller>();
-
-/**
- * EventCaller is used to track the caller of an event.
- */
-export class EventCaller {
-  public readonly id = randomUUID();
-
-  private constructor(
-    public readonly caller: unknown,
-    public readonly parent: EventCaller | null,
-  ) {}
-
-  #computedCallers: unknown[] | null = null;
-
-  public get computedCallers(): unknown[] {
-    if (this.#computedCallers != null) {
-      return this.#computedCallers;
-    }
-    const callers = [this.caller];
-    let parent = this.parent;
-    while (parent != null) {
-      callers.push(parent.caller);
-      parent = parent.parent;
-    }
-    this.#computedCallers = callers;
-    return callers;
-  }
-
-  public static create(
-    caller: unknown,
-    parent: EventCaller | null,
-  ): EventCaller {
-    return new EventCaller(caller, parent);
-  }
-}
-
-export function getEventCaller(): EventCaller | null {
-  return eventReasonAsyncLocalStorage.getStore() ?? null;
-}
-
-/**
- * @param caller who is calling this function, pass in `this` if it's a class method
- * @param fn
- */
-function withEventCaller<T>(caller: unknown, fn: () => T) {
-  // create a chain of event callers
-  const parentCaller = getEventCaller();
-  return eventReasonAsyncLocalStorage.run(
-    EventCaller.create(caller, parentCaller),
-    fn,
-  );
-}
-
-export function wrapEventCaller<This, Result, Args extends unknown[]>(
-  originalMethod: (this: This, ...args: Args) => Result,
-  context: ClassMethodDecoratorContext<object>,
-) {
-  const name = context.name;
-  context.addInitializer(function () {
-    // @ts-expect-error
-    const fn = this[name].bind(this);
-    // @ts-expect-error
-    this[name] = (...args: unknown[]) => {
-      return withEventCaller(this, () => fn(...args));
-    };
-  });
-  return function (this: This, ...args: Args): Result {
-    const result = originalMethod.call(this, ...args);
-    // patch for iterators because AsyncLocalStorage doesn't work with them
-    if (isAsyncGenerator(result)) {
-      const snapshot = AsyncLocalStorage.snapshot();
-      return (async function* asyncGeneratorWrapper() {
-        while (true) {
-          const { value, done } = await snapshot(() => result.next());
-          if (done) {
-            break;
-          }
-          yield value;
-        }
-      })() as Result;
-    } else if (isGenerator(result)) {
-      const snapshot = AsyncLocalStorage.snapshot();
-      return (function* generatorWrapper() {
-        while (true) {
-          const { value, done } = snapshot(() => result.next());
-          if (done) {
-            break;
-          }
-          yield value;
-        }
-      })() as Result;
-    }
-    return result;
-  };
-}