Add streaming for OpenAI agents (#693)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>
Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/long-cobras-play.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Support streaming for OpenAI agent

.changeset/orange-lions-remember.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Support streaming for OpenAI tool calls

.changeset/warm-ligers-hammer.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: llamacloud index from documents

.github/workflows/test.yml

@@ -4,6 +4,11 @@ on: [push, pull_request]
 
 jobs:
   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:
@@ -12,7 +17,7 @@ jobs:
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
-          node-version-file: ".nvmrc"
+          node-version: ${{ matrix.node-version }}
           cache: "pnpm"
       - name: Install dependencies
         run: pnpm install

examples/agent/multi_document_agent.ts

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

examples/agent/openai.ts

@@ -71,6 +71,6 @@ async function main() {
   console.log(String(response));
 }
 
-main().then(() => {
+void main().then(() => {
   console.log("Done");
 });

examples/agent/query_openai_agent.ts

@@ -41,6 +41,6 @@ async function main() {
   console.log(String(response));
 }
 
-main().then(() => {
+void main().then(() => {
   console.log("Done");
 });

examples/agent/react_agent.ts

@@ -77,6 +77,6 @@ async function main() {
   console.log(String(response));
 }
 
-main().then(() => {
+void main().then(() => {
   console.log("Done");
 });

examples/agent/step_wise_openai.ts

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

examples/agent/step_wise_query_tool.ts

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

examples/agent/step_wise_react.ts

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

examples/agent/stream_openai_agent.ts

@@ -72,6 +72,6 @@ async function main() {
   }
 }
 
-main().then(() => {
+void main().then(() => {
   console.log("\nDone");
 });

examples/agent/wiki.ts

@@ -0,0 +1,28 @@
+import { OpenAI, OpenAIAgent, WikipediaTool } from "llamaindex";
+
+async function main() {
+  const llm = new OpenAI({ model: "gpt-4-turbo-preview" });
+  const wikiTool = new WikipediaTool();
+
+  // Create an OpenAIAgent with the Wikipedia tool
+  const agent = new OpenAIAgent({
+    llm,
+    tools: [wikiTool],
+    verbose: true,
+  });
+
+  // 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

@@ -1,23 +0,0 @@
-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() {
   }
 }
 
-main();
+void main();

examples/astradb/load.ts

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

examples/astradb/query.ts

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

examples/chatHistory.ts

@@ -1,7 +1,18 @@
 import { stdin as input, stdout as output } from "node:process";
 import readline from "node:readline/promises";
 
-import { OpenAI, SimpleChatEngine, SummaryChatHistory } from "llamaindex";
+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);
+  });
+}
 
 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() {
   }
 }
 
-main();
+void main();

examples/chromadb/test.ts

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

examples/cloud/from_documents.ts

@@ -0,0 +1,44 @@
+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);
 }
 
-main();
+void main();

examples/evaluation/faithfulness.ts

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

examples/evaluation/relevancy.ts

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

examples/milvus/load.ts

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

examples/milvus/query.ts

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

examples/mongo.ts

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

examples/mongodb/1_import.ts

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

examples/mongodb/3_query.ts

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

examples/multimodal/clip.ts

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

examples/nodeParser/MarkdownNodeParser.ts

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

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}`);
-  var fileName = "";
+  const 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);
-    pgvs.clearCollection();
+    await pgvs.clearCollection();
 
     const ctx = await storageContextFromDefaults({ vectorStore: pgvs });
 
@@ -65,4 +65,4 @@ async function main(args: any) {
   process.exit(0);
 }
 
-main(process.argv).catch((err) => console.error(err));
+void 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}`);
-  var fileName = "";
+  const 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);
 }
 
-main(process.argv).catch((err) => console.error(err));
+void main(process.argv).catch((err) => console.error(err));

examples/prompts/promptMixin.ts

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

examples/qdrantdb/preFilters.ts

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

examples/readers/src/docx.ts

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

examples/readers/src/markdown.ts

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

examples/recipes/cost-analysis.ts

@@ -31,7 +31,7 @@ Settings.callbackManager.on("llm-end", (event) => {
 
 const question = "Hello, how are you?";
 console.log("Question:", question);
-llm
+void llm
   .chat({
     stream: true,
     messages: [

examples/routerQueryEngine.ts

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

examples/split.ts

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

examples/toolsStream.ts

@@ -0,0 +1,46 @@
+import { ChatResponseChunk, LLMChatParamsBase, OpenAI } from "llamaindex";
+
+async function main() {
+  const llm = new OpenAI({ model: "gpt-4-turbo-preview" });
+
+  const args: LLMChatParamsBase = {
+    messages: [
+      {
+        content: "Who was Goethe?",
+        role: "user",
+      },
+    ],
+    tools: [
+      {
+        type: "function",
+        function: {
+          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"],
+          },
+        },
+      },
+    ],
+    toolChoice: "auto",
+  };
+
+  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");
+})();

packages/core/package.json

@@ -9,7 +9,7 @@
     "@aws-crypto/sha256-js": "^5.2.0",
     "@datastax/astra-db-ts": "^0.1.4",
     "@grpc/grpc-js": "^1.10.2",
-    "@llamaindex/cloud": "0.0.4",
+    "@llamaindex/cloud": "0.0.5",
     "@llamaindex/env": "workspace:*",
     "@mistralai/mistralai": "^0.0.10",
     "@notionhq/client": "^2.2.14",

packages/core/src/ChatHistory.ts

@@ -1,7 +1,7 @@
 import { globalsHelper } from "./GlobalsHelper.js";
 import type { SummaryPrompt } from "./Prompt.js";
 import { defaultSummaryPrompt, messagesToHistoryStr } from "./Prompt.js";
-import { OpenAI } from "./llm/LLM.js";
+import { OpenAI } from "./llm/open_ai.js";
 import type { ChatMessage, LLM, MessageType } from "./llm/types.js";
 
 /**

packages/core/src/GlobalsHelper.ts

@@ -1,12 +1,5 @@
 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",
 }
@@ -51,39 +44,6 @@ class GlobalsHelper {
 
     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/LLM.js";
+import { OpenAI } from "./llm/open_ai.js";
 import type { LLM } from "./llm/types.js";
 import { PromptMixin } from "./prompts/index.js";
 import type {

packages/core/src/Retriever.ts

@@ -1,13 +1,8 @@
-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/LLM.js";
+import { OpenAI } from "./llm/open_ai.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,6 +1,6 @@
 import { CallbackManager } from "./callbacks/CallbackManager.js";
 import { OpenAIEmbedding } from "./embeddings/OpenAIEmbedding.js";
-import { OpenAI } from "./llm/LLM.js";
+import { OpenAI } from "./llm/open_ai.js";
 
 import { PromptHelper } from "./PromptHelper.js";
 import { SimpleNodeParser } from "./nodeParsers/SimpleNodeParser.js";

packages/core/src/TextSplitter.ts

@@ -66,8 +66,9 @@ 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 {
-  private chunkSize: number;
-  private chunkOverlap: number;
+  public chunkSize: number;
+  public chunkOverlap: number;
+
   private tokenizer: any;
   private tokenizerDecoder: any;
   private paragraphSeparator: string;

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

@@ -64,6 +64,7 @@ export class OpenAIAgent extends AgentRunner {
 
     super({
       agentWorker: stepEngine,
+      llm,
       memory,
       defaultToolChoice,
       chatHistory: prefixMessages,

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

@@ -9,6 +9,7 @@ import type {
   ChatMessage,
   ChatResponse,
   ChatResponseChunk,
+  LLMChatParamsBase,
 } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
 import { streamConverter, streamReducer } from "../../llm/utils.js";
@@ -166,8 +167,8 @@ export class OpenAIAgentWorker implements AgentWorker {
     task: Task,
     openaiTools: { [key: string]: any }[],
     toolChoice: string | { [key: string]: any } = "auto",
-  ): { [key: string]: any } {
-    const llmChatKwargs: { [key: string]: any } = {
+  ): LLMChatParamsBase {
+    const llmChatKwargs: LLMChatParamsBase = {
       messages: this.getAllMessages(task),
     };
 
@@ -179,17 +180,10 @@ export class OpenAIAgentWorker implements AgentWorker {
     return llmChatKwargs;
   }
 
-  /**
-   * Process message.
-   * @param task: task
-   * @param chatResponse: chat response
-   * @returns: agent chat response
-   */
   private _processMessage(
     task: Task,
-    chatResponse: ChatResponse,
+    aiMessage: ChatMessage,
   ): AgentChatResponse {
-    const aiMessage = chatResponse.message;
     task.extraState.newMemory.put(aiMessage);
 
     return new AgentChatResponse(aiMessage.content, task.extraState.sources);
@@ -198,16 +192,33 @@ export class OpenAIAgentWorker implements AgentWorker {
   private async _getStreamAiResponse(
     task: Task,
     llmChatKwargs: any,
-  ): Promise<StreamingAgentChatResponse> {
+  ): Promise<StreamingAgentChatResponse | AgentChatResponse> {
     const stream = await this.llm.chat({
       stream: true,
       ...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;
 
-    const iterator = streamConverter(
+    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)(
       streamReducer({
         stream,
-        initialValue: "",
+        initialValue: content,
         reducer: (accumulator, part) => (accumulator += part.delta),
         finished: (accumulator) => {
           task.extraState.newMemory.put({
@@ -219,7 +230,7 @@ export class OpenAIAgentWorker implements AgentWorker {
       (r: ChatResponseChunk) => new Response(r.delta),
     );
 
-    return new StreamingAgentChatResponse(iterator, task.extraState.sources);
+    return new StreamingAgentChatResponse(newStream, task.extraState.sources);
   }
 
   /**
@@ -240,7 +251,10 @@ export class OpenAIAgentWorker implements AgentWorker {
         ...llmChatKwargs,
       })) as unknown as ChatResponse;
 
-      return this._processMessage(task, chatResponse) as AgentChatResponse;
+      return this._processMessage(
+        task,
+        chatResponse.message,
+      ) as AgentChatResponse;
     } else if (mode === ChatResponseMode.STREAM) {
       return this._getStreamAiResponse(task, llmChatKwargs);
     }
@@ -286,7 +300,9 @@ export class OpenAIAgentWorker implements AgentWorker {
   initializeStep(task: Task, kwargs?: any): TaskStep {
     const sources: ToolOutput[] = [];
 
-    const newMemory = new ChatMemoryBuffer();
+    const newMemory = new ChatMemoryBuffer({
+      tokenLimit: task.memory.tokenLimit,
+    });
 
     const taskState = {
       sources,

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

@@ -106,7 +106,9 @@ export class ReActAgentWorker implements AgentWorker {
   initializeStep(task: Task, kwargs?: any): TaskStep {
     const sources: ToolOutput[] = [];
     const currentReasoning: BaseReasoningStep[] = [];
-    const newMemory = new ChatMemoryBuffer();
+    const newMemory = new ChatMemoryBuffer({
+      tokenLimit: task.memory.tokenLimit,
+    });
 
     const taskState = {
       sources,

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

@@ -58,6 +58,7 @@ 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

@@ -170,13 +170,13 @@ export class TaskStep implements ITaskStep {
  * @param isLast: isLast
  */
 export class TaskStepOutput {
-  output: any;
+  output: AgentChatResponse | StreamingAgentChatResponse;
   taskStep: TaskStep;
   nextSteps: TaskStep[];
   isLast: boolean;
 
   constructor(
-    output: any,
+    output: AgentChatResponse | StreamingAgentChatResponse,
     taskStep: TaskStep,
     nextSteps: TaskStep[],
     isLast: boolean = false,

packages/core/src/callbacks/CallbackManager.ts

@@ -1,6 +1,33 @@
 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";
+
+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 for the Llamaindex package.
@@ -21,26 +48,6 @@ declare module "llamaindex" {
 }
 
 //#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;
@@ -68,13 +75,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 extends BaseCallbackResponse {
+export interface StreamCallbackResponse {
   index: number;
   isDone?: boolean;
   token?: DefaultStreamToken;
 }
 
-export interface RetrievalCallbackResponse extends BaseCallbackResponse {
+export interface RetrievalCallbackResponse {
   query: string;
   nodes: NodeWithScore[];
 }
@@ -98,7 +105,11 @@ interface CallbackManagerMethods {
 
 const noop: (...args: any[]) => any = () => void 0;
 
-type EventHandler<Event extends CustomEvent> = (event: Event) => void;
+type EventHandler<Event extends CustomEvent> = (
+  event: Event & {
+    reason: EventCaller | null;
+  },
+) => void;
 
 export class CallbackManager implements CallbackManagerMethods {
   /**
@@ -110,7 +121,7 @@ export class CallbackManager implements CallbackManagerMethods {
         this.#handlers
           .get("stream")!
           .map((handler) =>
-            handler(new CustomEvent("stream", { detail: response })),
+            handler(LlamaIndexCustomEvent.fromEvent("stream", response)),
           ),
       );
     };
@@ -125,7 +136,7 @@ export class CallbackManager implements CallbackManagerMethods {
         this.#handlers
           .get("retrieve")!
           .map((handler) =>
-            handler(new CustomEvent("retrieve", { detail: response })),
+            handler(LlamaIndexCustomEvent.fromEvent("retrieve", response)),
           ),
       );
     };
@@ -188,6 +199,8 @@ export class CallbackManager implements CallbackManagerMethods {
     if (!handlers) {
       return;
     }
-    handlers.forEach((handler) => handler(new CustomEvent(event, { detail })));
+    handlers.forEach((handler) =>
+      handler(LlamaIndexCustomEvent.fromEvent(event, detail)),
+    );
   }
 }

packages/core/src/cloud/LlamaCloudIndex.ts

@@ -1,11 +1,20 @@
+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;
@@ -14,6 +23,151 @@ 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,13 +1,12 @@
 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"
@@ -51,9 +50,9 @@ export class LlamaCloudRetriever implements BaseRetriever {
     return this.client;
   }
 
+  @wrapEventCaller
   async retrieve({
     query,
-    parentEvent,
     preFilters,
   }: RetrieveParams): Promise<NodeWithScore[]> {
     const pipelines = await (
@@ -77,13 +76,9 @@ export class LlamaCloudRetriever implements BaseRetriever {
 
     const nodes = this.resultNodesToNodeWithScore(results.retrievalNodes);
 
-    Settings.callbackManager.onRetrieve({
+    Settings.callbackManager.dispatchEvent("retrieve", {
       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 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,
+        // 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,
       },
     };
   }
@@ -30,9 +30,10 @@ function getTransformationConfig(
     return {
       configurableTransformationType: "OPENAI_EMBEDDING",
       component: {
-        modelName: transformation.model,
-        apiKey: transformation.apiKey,
-        embedBatchSize: transformation.embedBatchSize,
+        // TODO: API doesnt accept camelCase
+        model: transformation.model,
+        api_key: transformation.apiKey,
+        embed_batch_size: transformation.embedBatchSize,
         dimensions: transformation.dimensions,
       },
     };
@@ -71,10 +72,12 @@ export async function getPipelineCreate(
     inputNodes = [],
   } = params;
 
+  const dataSources = inputNodes.map(getDataSourceConfig);
+
   return {
     name: pipelineName,
     configuredTransformations: transformations.map(getTransformationConfig),
-    dataSources: inputNodes.map(getDataSourceConfig),
+    dataSources,
     dataSinks: [],
     pipelineType,
   };

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

@@ -8,6 +8,7 @@ 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";
@@ -17,7 +18,6 @@ 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,6 +82,7 @@ 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,10 +1,9 @@
-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 type { Event } from "../../callbacks/CallbackManager.js";
+import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { ChatMessage, ChatResponseChunk, LLM } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
 import type { MessageContent } from "../../llm/types.js";
@@ -60,6 +59,7 @@ 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,21 +67,14 @@ 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(
@@ -98,7 +91,6 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     }
     const response = await this.chatModel.chat({
       messages: requestMessages.messages,
-      parentEvent,
     });
     chatHistory.addMessage(response.message);
     return new Response(response.message.content, requestMessages.nodes);
@@ -111,14 +103,13 @@ 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, parentEvent);
+    const context = await this.contextGenerator.generate(textOnly);
     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,9 +1,7 @@
-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";
@@ -56,17 +54,9 @@ export class DefaultContextGenerator
     return nodesWithScore;
   }
 
-  async generate(message: string, parentEvent?: Event): Promise<Context> {
-    if (!parentEvent) {
-      parentEvent = {
-        id: randomUUID(),
-        type: "wrapper",
-        tags: ["final"],
-      };
-    }
+  async generate(message: string): Promise<Context> {
     const sourceNodesWithScore = await this.retriever.retrieve({
       query: message,
-      parentEvent,
     });
 
     const nodes = await this.applyNodePostprocessors(

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

@@ -1,6 +1,7 @@
 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 { streamConverter, streamReducer } from "../../llm/utils.js";
@@ -25,6 +26,7 @@ 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>> {

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

@@ -1,7 +1,6 @@
 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";
@@ -56,7 +55,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, parentEvent?: Event): Promise<Context>;
+  generate(message: string): Promise<Context>;
 }
 
 export enum ChatResponseMode {

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

@@ -1,8 +1,7 @@
-import { randomUUID } from "@llamaindex/env";
 import type { NodeWithScore } from "../../Node.js";
 import type { Response } from "../../Response.js";
 import type { BaseRetriever } from "../../Retriever.js";
-import type { Event } from "../../callbacks/CallbackManager.js";
+import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import { PromptMixin } from "../../prompts/Mixin.js";
 import type { BaseSynthesizer } from "../../synthesizers/index.js";
@@ -62,10 +61,9 @@ export class RetrieverQueryEngine
     return nodesWithScore;
   }
 
-  private async retrieve(query: string, parentEvent: Event) {
+  private async retrieve(query: string) {
     const nodes = await this.retriever.retrieve({
       query,
-      parentEvent,
       preFilters: this.preFilters,
     });
 
@@ -74,28 +72,22 @@ 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 parentEvent: Event = params.parentEvent || {
-      id: randomUUID(),
-      type: "wrapper",
-      tags: ["final"],
-    };
-    const nodesWithScore = await this.retrieve(query, parentEvent);
+    const nodesWithScore = await this.retrieve(query);
     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,10 +1,8 @@
-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 {
@@ -20,6 +18,7 @@ import type {
   ToolMetadata,
 } from "../../types.js";
 
+import { wrapEventCaller } from "../../internal/context/EventCaller.js";
 import type { BaseQuestionGenerator, SubQuestion } from "./types.js";
 
 /**
@@ -80,29 +79,15 @@ 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, subQueryParentEvent)),
+      subQuestions.map((subQ) => this.querySubQ(subQ)),
     );
 
     const nodesWithScore = subQNodes
@@ -112,21 +97,16 @@ export class SubQuestionQueryEngine
       return this.responseSynthesizer.synthesize({
         query,
         nodesWithScore,
-        parentEvent,
         stream: true,
       });
     }
     return this.responseSynthesizer.synthesize({
       query,
       nodesWithScore,
-      parentEvent,
     });
   }
 
-  private async querySubQ(
-    subQ: SubQuestion,
-    parentEvent?: Event,
-  ): Promise<NodeWithScore | null> {
+  private async querySubQ(subQ: SubQuestion): Promise<NodeWithScore | null> {
     try {
       const question = subQ.subQuestion;
 
@@ -140,7 +120,6 @@ export class SubQuestionQueryEngine
 
       const responseText = await queryEngine?.call?.({
         query: question,
-        parentEvent,
       });
 
       if (!responseText) {

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

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

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

@@ -1,5 +1,4 @@
 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";
@@ -11,6 +10,7 @@ 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;
 
-    docStore.addDocuments(documents, true);
+    await docStore.addDocuments(documents, true);
     for (const doc of documents) {
       docStore.setDocumentHash(doc.id_, doc.hash);
     }
@@ -287,10 +287,8 @@ export class SummaryIndexRetriever implements BaseRetriever {
     this.index = index;
   }
 
-  async retrieve({
-    query,
-    parentEvent,
-  }: RetrieveParams): Promise<NodeWithScore[]> {
+  @wrapEventCaller
+  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const nodes = await this.index.docStore.getNodes(nodeIds);
     const result = nodes.map((node) => ({
@@ -298,13 +296,9 @@ export class SummaryIndexRetriever implements BaseRetriever {
       score: 1,
     }));
 
-    Settings.callbackManager.onRetrieve({
+    Settings.callbackManager.dispatchEvent("retrieve", {
       query,
       nodes: result,
-      event: globalsHelper.createEvent({
-        parentEvent,
-        type: "retrieve",
-      }),
     });
 
     return result;
@@ -340,10 +334,7 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
     this.serviceContext = serviceContext || index.serviceContext;
   }
 
-  async retrieve({
-    query,
-    parentEvent,
-  }: RetrieveParams): Promise<NodeWithScore[]> {
+  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const results: NodeWithScore[] = [];
 
@@ -380,13 +371,9 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
       results.push(...nodeWithScores);
     }
 
-    Settings.callbackManager.onRetrieve({
+    Settings.callbackManager.dispatchEvent("retrieve", {
       query,
       nodes: results,
-      event: globalsHelper.createEvent({
-        parentEvent,
-        type: "retrieve",
-      }),
     });
 
     return results;

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

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

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

@@ -5,9 +5,16 @@ 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

@@ -0,0 +1,99 @@
+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;
+  };
+}

packages/core/src/internal/utils.ts

@@ -0,0 +1,7 @@
+export const isAsyncGenerator = (obj: unknown): obj is AsyncGenerator => {
+  return obj != null && typeof obj === "object" && Symbol.asyncIterator in obj;
+};
+
+export const isGenerator = (obj: unknown): obj is Generator => {
+  return obj != null && typeof obj === "object" && Symbol.iterator in obj;
+};

packages/core/src/llm/LLM.ts

@@ -1,28 +1,10 @@
-import type OpenAILLM from "openai";
-import type { ClientOptions as OpenAIClientOptions } from "openai";
-import {
-  type Event,
-  type EventType,
-  type OpenAIStreamToken,
-  type StreamCallbackResponse,
-} from "../callbacks/CallbackManager.js";
+import { type StreamCallbackResponse } from "../callbacks/CallbackManager.js";
 
-import type { ChatCompletionMessageParam } from "openai/resources/index.js";
 import type { LLMOptions } from "portkey-ai";
-import { Tokenizers } from "../GlobalsHelper.js";
 import { getCallbackManager } from "../internal/settings/CallbackManager.js";
 import type { AnthropicSession } from "./anthropic.js";
 import { getAnthropicSession } from "./anthropic.js";
-import type { AzureOpenAIConfig } from "./azure.js";
-import {
-  getAzureBaseUrl,
-  getAzureConfigFromEnv,
-  getAzureModel,
-  shouldUseAzure,
-} from "./azure.js";
 import { BaseLLM } from "./base.js";
-import type { OpenAISession } from "./open_ai.js";
-import { getOpenAISession } from "./open_ai.js";
 import type { PortkeySession } from "./portkey.js";
 import { getPortkeySession } from "./portkey.js";
 import { ReplicateSession } from "./replicate_ai.js";
@@ -35,290 +17,7 @@ import type {
   LLMMetadata,
   MessageType,
 } from "./types.js";
-import { llmEvent } from "./utils.js";
-
-export const GPT4_MODELS = {
-  "gpt-4": { contextWindow: 8192 },
-  "gpt-4-32k": { contextWindow: 32768 },
-  "gpt-4-32k-0613": { contextWindow: 32768 },
-  "gpt-4-turbo-preview": { contextWindow: 128000 },
-  "gpt-4-1106-preview": { contextWindow: 128000 },
-  "gpt-4-0125-preview": { contextWindow: 128000 },
-  "gpt-4-vision-preview": { contextWindow: 128000 },
-};
-
-// NOTE we don't currently support gpt-3.5-turbo-instruct and don't plan to in the near future
-export const GPT35_MODELS = {
-  "gpt-3.5-turbo": { contextWindow: 4096 },
-  "gpt-3.5-turbo-0613": { contextWindow: 4096 },
-  "gpt-3.5-turbo-16k": { contextWindow: 16384 },
-  "gpt-3.5-turbo-16k-0613": { contextWindow: 16384 },
-  "gpt-3.5-turbo-1106": { contextWindow: 16384 },
-  "gpt-3.5-turbo-0125": { contextWindow: 16384 },
-};
-
-/**
- * We currently support GPT-3.5 and GPT-4 models
- */
-export const ALL_AVAILABLE_OPENAI_MODELS = {
-  ...GPT4_MODELS,
-  ...GPT35_MODELS,
-};
-
-export const isFunctionCallingModel = (model: string): boolean => {
-  const isChatModel = Object.keys(ALL_AVAILABLE_OPENAI_MODELS).includes(model);
-  const isOld = model.includes("0314") || model.includes("0301");
-  return isChatModel && !isOld;
-};
-
-/**
- * OpenAI LLM implementation
- */
-export class OpenAI extends BaseLLM {
-  // Per completion OpenAI params
-  model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS | string;
-  temperature: number;
-  topP: number;
-  maxTokens?: number;
-  additionalChatOptions?: Omit<
-    Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
-    | "max_tokens"
-    | "messages"
-    | "model"
-    | "temperature"
-    | "top_p"
-    | "stream"
-    | "tools"
-    | "toolChoice"
-  >;
-
-  // OpenAI session params
-  apiKey?: string = undefined;
-  maxRetries: number;
-  timeout?: number;
-  session: OpenAISession;
-  additionalSessionOptions?: Omit<
-    Partial<OpenAIClientOptions>,
-    "apiKey" | "maxRetries" | "timeout"
-  >;
-
-  constructor(
-    init?: Partial<OpenAI> & {
-      azure?: AzureOpenAIConfig;
-    },
-  ) {
-    super();
-    this.model = init?.model ?? "gpt-3.5-turbo";
-    this.temperature = init?.temperature ?? 0.1;
-    this.topP = init?.topP ?? 1;
-    this.maxTokens = init?.maxTokens ?? undefined;
-
-    this.maxRetries = init?.maxRetries ?? 10;
-    this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
-    this.additionalChatOptions = init?.additionalChatOptions;
-    this.additionalSessionOptions = init?.additionalSessionOptions;
-
-    if (init?.azure || shouldUseAzure()) {
-      const azureConfig = getAzureConfigFromEnv({
-        ...init?.azure,
-        model: getAzureModel(this.model),
-      });
-
-      if (!azureConfig.apiKey) {
-        throw new Error(
-          "Azure API key is required for OpenAI Azure models. Please set the AZURE_OPENAI_KEY environment variable.",
-        );
-      }
-
-      this.apiKey = azureConfig.apiKey;
-      this.session =
-        init?.session ??
-        getOpenAISession({
-          azure: true,
-          apiKey: this.apiKey,
-          baseURL: getAzureBaseUrl(azureConfig),
-          maxRetries: this.maxRetries,
-          timeout: this.timeout,
-          defaultQuery: { "api-version": azureConfig.apiVersion },
-          ...this.additionalSessionOptions,
-        });
-    } else {
-      this.apiKey = init?.apiKey ?? undefined;
-      this.session =
-        init?.session ??
-        getOpenAISession({
-          apiKey: this.apiKey,
-          maxRetries: this.maxRetries,
-          timeout: this.timeout,
-          ...this.additionalSessionOptions,
-        });
-    }
-  }
-
-  get metadata() {
-    const contextWindow =
-      ALL_AVAILABLE_OPENAI_MODELS[
-        this.model as keyof typeof ALL_AVAILABLE_OPENAI_MODELS
-      ]?.contextWindow ?? 1024;
-    return {
-      model: this.model,
-      temperature: this.temperature,
-      topP: this.topP,
-      maxTokens: this.maxTokens,
-      contextWindow,
-      tokenizer: Tokenizers.CL100K_BASE,
-      isFunctionCallingModel: isFunctionCallingModel(this.model),
-    };
-  }
-
-  mapMessageType(
-    messageType: MessageType,
-  ): "user" | "assistant" | "system" | "function" | "tool" {
-    switch (messageType) {
-      case "user":
-        return "user";
-      case "assistant":
-        return "assistant";
-      case "system":
-        return "system";
-      case "function":
-        return "function";
-      case "tool":
-        return "tool";
-      default:
-        return "user";
-    }
-  }
-
-  toOpenAIMessage(messages: ChatMessage[]) {
-    return messages.map((message) => {
-      const additionalKwargs = message.additionalKwargs ?? {};
-
-      if (message.additionalKwargs?.toolCalls) {
-        additionalKwargs.tool_calls = message.additionalKwargs.toolCalls;
-        delete additionalKwargs.toolCalls;
-      }
-
-      return {
-        role: this.mapMessageType(message.role),
-        content: message.content,
-        ...additionalKwargs,
-      };
-    });
-  }
-
-  chat(
-    params: LLMChatParamsStreaming,
-  ): Promise<AsyncIterable<ChatResponseChunk>>;
-  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  @llmEvent
-  async chat(
-    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
-  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, parentEvent, stream, tools, toolChoice } = params;
-    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
-      model: this.model,
-      temperature: this.temperature,
-      max_tokens: this.maxTokens,
-      tools: tools,
-      tool_choice: toolChoice,
-      messages: this.toOpenAIMessage(messages) as ChatCompletionMessageParam[],
-      top_p: this.topP,
-      ...this.additionalChatOptions,
-    };
-
-    // Streaming
-    if (stream) {
-      return this.streamChat(params);
-    }
-
-    // Non-streaming
-    const response = await this.session.openai.chat.completions.create({
-      ...baseRequestParams,
-      stream: false,
-    });
-
-    const content = response.choices[0].message?.content ?? null;
-
-    const kwargsOutput: Record<string, any> = {};
-
-    if (response.choices[0].message?.tool_calls) {
-      kwargsOutput.toolCalls = response.choices[0].message.tool_calls;
-    }
-
-    return {
-      message: {
-        content,
-        role: response.choices[0].message.role,
-        additionalKwargs: kwargsOutput,
-      },
-    };
-  }
-
-  protected async *streamChat({
-    messages,
-    parentEvent,
-  }: LLMChatParamsStreaming): AsyncIterable<ChatResponseChunk> {
-    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
-      model: this.model,
-      temperature: this.temperature,
-      max_tokens: this.maxTokens,
-      messages: messages.map(
-        (message) =>
-          ({
-            role: this.mapMessageType(message.role),
-            content: message.content,
-          }) as ChatCompletionMessageParam,
-      ),
-      top_p: this.topP,
-      ...this.additionalChatOptions,
-    };
-
-    //Now let's wrap our stream in a callback
-    const onLLMStream = getCallbackManager().onLLMStream;
-
-    const chunk_stream: AsyncIterable<OpenAIStreamToken> =
-      await this.session.openai.chat.completions.create({
-        ...baseRequestParams,
-        stream: true,
-      });
-
-    const event: Event = parentEvent
-      ? parentEvent
-      : {
-          id: "unspecified",
-          type: "llmPredict" as EventType,
-        };
-
-    // TODO: add callback to streamConverter and use streamConverter here
-    //Indices
-    let idx_counter: number = 0;
-    for await (const part of chunk_stream) {
-      if (!part.choices.length) continue;
-
-      //Increment
-      part.choices[0].index = idx_counter;
-      const is_done: boolean =
-        part.choices[0].finish_reason === "stop" ? true : false;
-      //onLLMStream Callback
-
-      const stream_callback: StreamCallbackResponse = {
-        event: event,
-        index: idx_counter,
-        isDone: is_done,
-        token: part,
-      };
-      onLLMStream(stream_callback);
-
-      idx_counter++;
-
-      yield {
-        delta: part.choices[0].delta.content ?? "",
-      };
-    }
-    return;
-  }
-}
+import { wrapLLMEvent } from "./utils.js";
 
 export const ALL_AVAILABLE_LLAMADEUCE_MODELS = {
   "Llama-2-70b-chat-old": {
@@ -541,11 +240,11 @@ If a question does not make any sense, or is not factually coherent, explain why
     params: LLMChatParamsStreaming,
   ): Promise<AsyncIterable<ChatResponseChunk>>;
   chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  @llmEvent
+  @wrapLLMEvent
   async chat(
     params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
   ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, parentEvent, stream } = params;
+    const { messages, stream } = params;
     const api = ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model]
       .replicateApi as `${string}/${string}:${string}`;
 
@@ -683,13 +382,13 @@ export class Anthropic extends BaseLLM {
     params: LLMChatParamsStreaming,
   ): Promise<AsyncIterable<ChatResponseChunk>>;
   chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  @llmEvent
+  @wrapLLMEvent
   async chat(
     params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
   ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
     let { messages } = params;
 
-    const { parentEvent, stream } = params;
+    const { stream } = params;
 
     let systemPrompt: string | null = null;
 
@@ -706,7 +405,7 @@ export class Anthropic extends BaseLLM {
 
     //Streaming
     if (stream) {
-      return this.streamChat(messages, parentEvent, systemPrompt);
+      return this.streamChat(messages, systemPrompt);
     }
 
     //Non-streaming
@@ -726,7 +425,6 @@ export class Anthropic extends BaseLLM {
 
   protected async *streamChat(
     messages: ChatMessage[],
-    parentEvent?: Event | undefined,
     systemPrompt?: string | null,
   ): AsyncIterable<ChatResponseChunk> {
     const stream = await this.session.anthropic.messages.create({
@@ -782,13 +480,13 @@ export class Portkey extends BaseLLM {
     params: LLMChatParamsStreaming,
   ): Promise<AsyncIterable<ChatResponseChunk>>;
   chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
-  @llmEvent
+  @wrapLLMEvent
   async chat(
     params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
   ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, parentEvent, stream, extraParams } = params;
+    const { messages, stream, extraParams } = params;
     if (stream) {
-      return this.streamChat(messages, parentEvent, extraParams);
+      return this.streamChat(messages, extraParams);
     } else {
       const bodyParams = extraParams || {};
       const response = await this.session.portkey.chatCompletions.create({
@@ -804,25 +502,14 @@ export class Portkey extends BaseLLM {
 
   async *streamChat(
     messages: ChatMessage[],
-    parentEvent?: Event,
     params?: Record<string, any>,
   ): AsyncIterable<ChatResponseChunk> {
-    // Wrapping the stream in a callback.
-    const onLLMStream = getCallbackManager().onLLMStream;
-
     const chunkStream = await this.session.portkey.chatCompletions.create({
       messages,
       ...params,
       stream: true,
     });
 
-    const event: Event = parentEvent
-      ? parentEvent
-      : {
-          id: "unspecified",
-          type: "llmPredict" as EventType,
-        };
-
     //Indices
     let idx_counter: number = 0;
     for await (const part of chunkStream) {
@@ -833,12 +520,11 @@ export class Portkey extends BaseLLM {
       //onLLMStream Callback
 
       const stream_callback: StreamCallbackResponse = {
-        event: event,
         index: idx_counter,
         isDone: is_done,
         // token: part,
       };
-      onLLMStream(stream_callback);
+      getCallbackManager().dispatchEvent("stream", stream_callback);
 
       idx_counter++;
 

packages/core/src/llm/base.ts

@@ -23,11 +23,10 @@ export abstract class BaseLLM implements LLM {
   async complete(
     params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
   ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
-    const { prompt, parentEvent, stream } = params;
+    const { prompt, stream } = params;
     if (stream) {
       const stream = await this.chat({
         messages: [{ content: prompt, role: "user" }],
-        parentEvent,
         stream: true,
       });
       return streamConverter(stream, (chunk) => {
@@ -38,7 +37,6 @@ export abstract class BaseLLM implements LLM {
     }
     const chatResponse = await this.chat({
       messages: [{ content: prompt, role: "user" }],
-      parentEvent,
     });
     return { text: chatResponse.message.content as string };
   }

packages/core/src/llm/fireworks.ts

@@ -1,5 +1,5 @@
 import { getEnv } from "@llamaindex/env";
-import { OpenAI } from "./LLM.js";
+import { OpenAI } from "./open_ai.js";
 
 export class FireworksLLM extends OpenAI {
   constructor(init?: Partial<OpenAI>) {

packages/core/src/llm/groq.ts

@@ -1,5 +1,5 @@
 import { getEnv } from "@llamaindex/env";
-import { OpenAI } from "./LLM.js";
+import { OpenAI } from "./open_ai.js";
 
 export class Groq extends OpenAI {
   constructor(init?: Partial<OpenAI>) {

packages/core/src/llm/mistral.ts

@@ -1,10 +1,6 @@
 import { getEnv } from "@llamaindex/env";
 import { Settings } from "../Settings.js";
-import {
-  type Event,
-  type EventType,
-  type StreamCallbackResponse,
-} from "../callbacks/CallbackManager.js";
+import { type StreamCallbackResponse } from "../callbacks/CallbackManager.js";
 import { BaseLLM } from "./base.js";
 import type {
   ChatMessage,
@@ -116,21 +112,10 @@ export class MistralAI extends BaseLLM {
 
   protected async *streamChat({
     messages,
-    parentEvent,
   }: LLMChatParamsStreaming): AsyncIterable<ChatResponseChunk> {
-    //Now let's wrap our stream in a callback
-    const onLLMStream = Settings.callbackManager.onLLMStream;
-
     const client = await this.session.getClient();
     const chunkStream = await client.chatStream(this.buildParams(messages));
 
-    const event: Event = parentEvent
-      ? parentEvent
-      : {
-          id: "unspecified",
-          type: "llmPredict" as EventType,
-        };
-
     //Indices
     let idx_counter: number = 0;
     for await (const part of chunkStream) {
@@ -141,12 +126,12 @@ export class MistralAI extends BaseLLM {
         part.choices[0].finish_reason === "stop" ? true : false;
 
       const stream_callback: StreamCallbackResponse = {
-        event: event,
         index: idx_counter,
         isDone: isDone,
         token: part,
       };
-      onLLMStream(stream_callback);
+
+      Settings.callbackManager.dispatchEvent("stream", stream_callback);
 
       idx_counter++;
 

packages/core/src/llm/ollama.ts

@@ -1,5 +1,4 @@
 import { ok } from "@llamaindex/env";
-import type { Event } from "../callbacks/CallbackManager.js";
 import { BaseEmbedding } from "../embeddings/types.js";
 import type {
   ChatResponse,
@@ -69,7 +68,7 @@ export class Ollama extends BaseEmbedding implements LLM {
   async chat(
     params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
   ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
-    const { messages, parentEvent, stream } = params;
+    const { messages, stream } = params;
     const payload = {
       model: this.model,
       messages: messages.map((message) => ({
@@ -106,14 +105,13 @@ export class Ollama extends BaseEmbedding implements LLM {
       const stream = response.body;
       ok(stream, "stream is null");
       ok(stream instanceof ReadableStream, "stream is not readable");
-      return this.streamChat(stream, messageAccessor, parentEvent);
+      return this.streamChat(stream, messageAccessor);
     }
   }
 
   private async *streamChat<T>(
     stream: ReadableStream<Uint8Array>,
     accessor: (data: any) => T,
-    parentEvent?: Event,
   ): AsyncIterable<T> {
     const reader = stream.getReader();
     while (true) {
@@ -147,7 +145,7 @@ export class Ollama extends BaseEmbedding implements LLM {
   async complete(
     params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
   ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
-    const { prompt, parentEvent, stream } = params;
+    const { prompt, stream } = params;
     const payload = {
       model: this.model,
       prompt: prompt,
@@ -177,7 +175,7 @@ export class Ollama extends BaseEmbedding implements LLM {
       const stream = response.body;
       ok(stream, "stream is null");
       ok(stream instanceof ReadableStream, "stream is not readable");
-      return this.streamChat(stream, completionAccessor, parentEvent);
+      return this.streamChat(stream, completionAccessor);
     }
   }
 

packages/core/src/llm/open_ai.ts

@@ -1,16 +1,43 @@
 import { getEnv } from "@llamaindex/env";
 import _ from "lodash";
-import type { ClientOptions } from "openai";
-import OpenAI from "openai";
+import type OpenAILLM from "openai";
+import type {
+  ClientOptions,
+  ClientOptions as OpenAIClientOptions,
+} from "openai";
+import { OpenAI as OrigOpenAI } from "openai";
 
-export class AzureOpenAI extends OpenAI {
+import type { ChatCompletionMessageParam } from "openai/resources/index.js";
+import { Tokenizers } from "../GlobalsHelper.js";
+import { wrapEventCaller } from "../internal/context/EventCaller.js";
+import { getCallbackManager } from "../internal/settings/CallbackManager.js";
+import type { AzureOpenAIConfig } from "./azure.js";
+import {
+  getAzureBaseUrl,
+  getAzureConfigFromEnv,
+  getAzureModel,
+  shouldUseAzure,
+} from "./azure.js";
+import { BaseLLM } from "./base.js";
+import type {
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+  MessageToolCall,
+  MessageType,
+} from "./types.js";
+import { wrapLLMEvent } from "./utils.js";
+
+export class AzureOpenAI extends OrigOpenAI {
   protected override authHeaders() {
     return { "api-key": this.apiKey };
   }
 }
 
 export class OpenAISession {
-  openai: OpenAI;
+  openai: OrigOpenAI;
 
   constructor(options: ClientOptions & { azure?: boolean } = {}) {
     if (!options.apiKey) {
@@ -24,7 +51,7 @@ export class OpenAISession {
     if (options.azure) {
       this.openai = new AzureOpenAI(options);
     } else {
-      this.openai = new OpenAI({
+      this.openai = new OrigOpenAI({
         ...options,
         // defaultHeaders: { "OpenAI-Beta": "assistants=v1" },
       });
@@ -60,3 +87,288 @@ export function getOpenAISession(
 
   return session;
 }
+
+export const GPT4_MODELS = {
+  "gpt-4": { contextWindow: 8192 },
+  "gpt-4-32k": { contextWindow: 32768 },
+  "gpt-4-32k-0613": { contextWindow: 32768 },
+  "gpt-4-turbo-preview": { contextWindow: 128000 },
+  "gpt-4-1106-preview": { contextWindow: 128000 },
+  "gpt-4-0125-preview": { contextWindow: 128000 },
+  "gpt-4-vision-preview": { contextWindow: 128000 },
+};
+
+// NOTE we don't currently support gpt-3.5-turbo-instruct and don't plan to in the near future
+export const GPT35_MODELS = {
+  "gpt-3.5-turbo": { contextWindow: 4096 },
+  "gpt-3.5-turbo-0613": { contextWindow: 4096 },
+  "gpt-3.5-turbo-16k": { contextWindow: 16384 },
+  "gpt-3.5-turbo-16k-0613": { contextWindow: 16384 },
+  "gpt-3.5-turbo-1106": { contextWindow: 16384 },
+  "gpt-3.5-turbo-0125": { contextWindow: 16384 },
+};
+
+/**
+ * We currently support GPT-3.5 and GPT-4 models
+ */
+export const ALL_AVAILABLE_OPENAI_MODELS = {
+  ...GPT4_MODELS,
+  ...GPT35_MODELS,
+};
+
+export const isFunctionCallingModel = (model: string): boolean => {
+  const isChatModel = Object.keys(ALL_AVAILABLE_OPENAI_MODELS).includes(model);
+  const isOld = model.includes("0314") || model.includes("0301");
+  return isChatModel && !isOld;
+};
+
+/**
+ * OpenAI LLM implementation
+ */
+export class OpenAI extends BaseLLM {
+  // Per completion OpenAI params
+  model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS | string;
+  temperature: number;
+  topP: number;
+  maxTokens?: number;
+  additionalChatOptions?: Omit<
+    Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
+    | "max_tokens"
+    | "messages"
+    | "model"
+    | "temperature"
+    | "top_p"
+    | "stream"
+    | "tools"
+    | "toolChoice"
+  >;
+
+  // OpenAI session params
+  apiKey?: string = undefined;
+  maxRetries: number;
+  timeout?: number;
+  session: OpenAISession;
+  additionalSessionOptions?: Omit<
+    Partial<OpenAIClientOptions>,
+    "apiKey" | "maxRetries" | "timeout"
+  >;
+
+  constructor(
+    init?: Partial<OpenAI> & {
+      azure?: AzureOpenAIConfig;
+    },
+  ) {
+    super();
+    this.model = init?.model ?? "gpt-3.5-turbo";
+    this.temperature = init?.temperature ?? 0.1;
+    this.topP = init?.topP ?? 1;
+    this.maxTokens = init?.maxTokens ?? undefined;
+
+    this.maxRetries = init?.maxRetries ?? 10;
+    this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
+    this.additionalChatOptions = init?.additionalChatOptions;
+    this.additionalSessionOptions = init?.additionalSessionOptions;
+
+    if (init?.azure || shouldUseAzure()) {
+      const azureConfig = getAzureConfigFromEnv({
+        ...init?.azure,
+        model: getAzureModel(this.model),
+      });
+
+      if (!azureConfig.apiKey) {
+        throw new Error(
+          "Azure API key is required for OpenAI Azure models. Please set the AZURE_OPENAI_KEY environment variable.",
+        );
+      }
+
+      this.apiKey = azureConfig.apiKey;
+      this.session =
+        init?.session ??
+        getOpenAISession({
+          azure: true,
+          apiKey: this.apiKey,
+          baseURL: getAzureBaseUrl(azureConfig),
+          maxRetries: this.maxRetries,
+          timeout: this.timeout,
+          defaultQuery: { "api-version": azureConfig.apiVersion },
+          ...this.additionalSessionOptions,
+        });
+    } else {
+      this.apiKey = init?.apiKey ?? undefined;
+      this.session =
+        init?.session ??
+        getOpenAISession({
+          apiKey: this.apiKey,
+          maxRetries: this.maxRetries,
+          timeout: this.timeout,
+          ...this.additionalSessionOptions,
+        });
+    }
+  }
+
+  get metadata() {
+    const contextWindow =
+      ALL_AVAILABLE_OPENAI_MODELS[
+        this.model as keyof typeof ALL_AVAILABLE_OPENAI_MODELS
+      ]?.contextWindow ?? 1024;
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      contextWindow,
+      tokenizer: Tokenizers.CL100K_BASE,
+      isFunctionCallingModel: isFunctionCallingModel(this.model),
+    };
+  }
+
+  mapMessageType(
+    messageType: MessageType,
+  ): "user" | "assistant" | "system" | "function" | "tool" {
+    switch (messageType) {
+      case "user":
+        return "user";
+      case "assistant":
+        return "assistant";
+      case "system":
+        return "system";
+      case "function":
+        return "function";
+      case "tool":
+        return "tool";
+      default:
+        return "user";
+    }
+  }
+
+  toOpenAIMessage(messages: ChatMessage[]) {
+    return messages.map((message) => {
+      const additionalKwargs = message.additionalKwargs ?? {};
+
+      if (message.additionalKwargs?.toolCalls) {
+        additionalKwargs.tool_calls = message.additionalKwargs.toolCalls;
+        delete additionalKwargs.toolCalls;
+      }
+
+      return {
+        role: this.mapMessageType(message.role),
+        content: message.content,
+        ...additionalKwargs,
+      };
+    });
+  }
+
+  chat(
+    params: LLMChatParamsStreaming,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  @wrapEventCaller
+  @wrapLLMEvent
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, stream, tools, toolChoice } = params;
+    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
+      model: this.model,
+      temperature: this.temperature,
+      max_tokens: this.maxTokens,
+      tools: tools,
+      tool_choice: toolChoice,
+      messages: this.toOpenAIMessage(messages) as ChatCompletionMessageParam[],
+      top_p: this.topP,
+      ...this.additionalChatOptions,
+    };
+
+    // Streaming
+    if (stream) {
+      return this.streamChat(baseRequestParams);
+    }
+
+    // Non-streaming
+    const response = await this.session.openai.chat.completions.create({
+      ...baseRequestParams,
+      stream: false,
+    });
+
+    const content = response.choices[0].message?.content ?? null;
+
+    const kwargsOutput: Record<string, any> = {};
+
+    if (response.choices[0].message?.tool_calls) {
+      kwargsOutput.toolCalls = response.choices[0].message.tool_calls;
+    }
+
+    return {
+      message: {
+        content,
+        role: response.choices[0].message.role,
+        additionalKwargs: kwargsOutput,
+      },
+    };
+  }
+
+  @wrapEventCaller
+  protected async *streamChat(
+    baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams,
+  ): AsyncIterable<ChatResponseChunk> {
+    const stream: AsyncIterable<OpenAILLM.Chat.ChatCompletionChunk> =
+      await this.session.openai.chat.completions.create({
+        ...baseRequestParams,
+        stream: true,
+      });
+
+    // TODO: add callback to streamConverter and use streamConverter here
+    //Indices
+    let idxCounter: number = 0;
+    const toolCalls: MessageToolCall[] = [];
+    for await (const part of stream) {
+      if (!part.choices.length) continue;
+      const choice = part.choices[0];
+      updateToolCalls(toolCalls, choice.delta.tool_calls);
+
+      const isDone: boolean = choice.finish_reason !== null;
+
+      getCallbackManager().dispatchEvent("stream", {
+        index: idxCounter++,
+        isDone: isDone,
+        token: part,
+      });
+
+      yield {
+        // add tool calls to final chunk
+        additionalKwargs:
+          toolCalls.length > 0 ? { toolCalls: toolCalls } : undefined,
+        delta: choice.delta.content ?? "",
+      };
+    }
+    return;
+  }
+}
+
+function updateToolCalls(
+  toolCalls: MessageToolCall[],
+  toolCallDeltas?: OpenAILLM.Chat.Completions.ChatCompletionChunk.Choice.Delta.ToolCall[],
+) {
+  function augmentToolCall(
+    toolCall?: MessageToolCall,
+    toolCallDelta?: OpenAILLM.Chat.Completions.ChatCompletionChunk.Choice.Delta.ToolCall,
+  ) {
+    toolCall =
+      toolCall ??
+      ({ function: { name: "", arguments: "" } } as MessageToolCall);
+    toolCall.id = toolCall.id ?? toolCallDelta?.id;
+    toolCall.type = toolCall.type ?? toolCallDelta?.type;
+    if (toolCallDelta?.function?.arguments) {
+      toolCall.function.arguments += toolCallDelta.function.arguments;
+    }
+    if (toolCallDelta?.function?.name) {
+      toolCall.function.name += toolCallDelta.function.name;
+    }
+    return toolCall;
+  }
+  if (toolCallDeltas) {
+    toolCallDeltas?.forEach((toolCall, i) => {
+      toolCalls[i] = augmentToolCall(toolCalls[i], toolCall);
+    });
+  }
+}

packages/core/src/llm/together.ts

@@ -1,5 +1,5 @@
 import { getEnv } from "@llamaindex/env";
-import { OpenAI } from "./LLM.js";
+import { OpenAI } from "./open_ai.js";
 
 export class TogetherLLM extends OpenAI {
   constructor(init?: Partial<OpenAI>) {

packages/core/src/llm/types.ts

@@ -1,5 +1,4 @@
 import type { Tokenizers } from "../GlobalsHelper.js";
-import { type Event } from "../callbacks/CallbackManager.js";
 
 type LLMBaseEvent<
   Type extends string,
@@ -85,6 +84,7 @@ export interface ChatResponse {
 
 export interface ChatResponseChunk {
   delta: string;
+  additionalKwargs?: Record<string, any>;
 }
 
 export interface CompletionResponse {
@@ -103,7 +103,6 @@ export interface LLMMetadata {
 
 export interface LLMChatParamsBase {
   messages: ChatMessage[];
-  parentEvent?: Event;
   extraParams?: Record<string, any>;
   tools?: any;
   toolChoice?: any;
@@ -120,7 +119,6 @@ export interface LLMChatParamsNonStreaming extends LLMChatParamsBase {
 
 export interface LLMCompletionParamsBase {
   prompt: any;
-  parentEvent?: Event;
 }
 
 export interface LLMCompletionParamsStreaming extends LLMCompletionParamsBase {
@@ -142,3 +140,14 @@ export interface MessageContentDetail {
  * Extended type for the content of a message that allows for multi-modal messages.
  */
 export type MessageContent = string | MessageContentDetail[];
+
+interface Function {
+  arguments: string;
+  name: string;
+}
+
+export interface MessageToolCall {
+  id: string;
+  function: Function;
+  type: "function";
+}

packages/core/src/llm/utils.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from "@llamaindex/env";
 import { getCallbackManager } from "../internal/settings/CallbackManager.js";
 import type { ChatResponse, LLM, LLMChat, MessageContent } from "./types.js";
 
@@ -47,7 +48,7 @@ export function extractText(message: MessageContent): string {
 /**
  * @internal
  */
-export function llmEvent(
+export function wrapLLMEvent(
   originalMethod: LLMChat["chat"],
   _context: ClassMethodDecoratorContext,
 ) {
@@ -62,6 +63,8 @@ export function llmEvent(
     });
     const response = await originalMethod.call(this, ...params);
     if (Symbol.asyncIterator in response) {
+      // save snapshot to restore it after the response is done
+      const snapshot = AsyncLocalStorage.snapshot();
       const originalAsyncIterator = {
         [Symbol.asyncIterator]: response[Symbol.asyncIterator].bind(response),
       };
@@ -82,10 +85,12 @@ export function llmEvent(
           }
           yield chunk;
         }
-        getCallbackManager().dispatchEvent("llm-end", {
-          payload: {
-            response: finalResponse,
-          },
+        snapshot(() => {
+          getCallbackManager().dispatchEvent("llm-end", {
+            payload: {
+              response: finalResponse,
+            },
+          });
         });
       };
     } else {

packages/core/src/memory/ChatMemoryBuffer.ts

@@ -1,13 +1,17 @@
-import type { ChatMessage } from "../llm/index.js";
+import type { ChatMessage, LLM } from "../llm/index.js";
 import { SimpleChatStore } from "../storage/chatStore/SimpleChatStore.js";
 import type { BaseChatStore } from "../storage/chatStore/types.js";
 import type { BaseMemory } from "./types.js";
 
+const DEFAULT_TOKEN_LIMIT_RATIO = 0.75;
+const DEFAULT_TOKEN_LIMIT = 3000;
+
 type ChatMemoryBufferParams = {
   tokenLimit?: number;
   chatStore?: BaseChatStore;
   chatStoreKey?: string;
   chatHistory?: ChatMessage[];
+  llm?: LLM;
 };
 
 /**
@@ -23,9 +27,16 @@ export class ChatMemoryBuffer implements BaseMemory {
    * Initialize.
    */
   constructor(init?: Partial<ChatMemoryBufferParams>) {
-    this.tokenLimit = init?.tokenLimit ?? 3000;
     this.chatStore = init?.chatStore ?? new SimpleChatStore();
     this.chatStoreKey = init?.chatStoreKey ?? "chat_history";
+    if (init?.llm) {
+      const contextWindow = init.llm.metadata.contextWindow;
+      this.tokenLimit =
+        init?.tokenLimit ??
+        Math.ceil(contextWindow * DEFAULT_TOKEN_LIMIT_RATIO);
+    } else {
+      this.tokenLimit = init?.tokenLimit ?? DEFAULT_TOKEN_LIMIT;
+    }
 
     if (init?.chatHistory) {
       this.chatStore.setMessages(this.chatStoreKey, init.chatHistory);
@@ -49,7 +60,7 @@ export class ChatMemoryBuffer implements BaseMemory {
 
     while (tokenCount > this.tokenLimit && messageCount > 1) {
       messageCount -= 1;
-      if (chatHistory[-messageCount].role === "assistant") {
+      if (chatHistory.at(-messageCount)?.role === "assistant") {
         // we cannot have an assistant message at the start of the chat history
         // if after removal of the first, we have an assistant message,
         // we need to remove the assistant message too

packages/core/src/objects/base.ts

@@ -170,10 +170,10 @@ export class ObjectIndex {
     return new ObjectIndex(index, objectMapping);
   }
 
-  insertObject(obj: any): void {
+  async insertObject(obj: any): Promise<void> {
     this._objectNodeMapping.addObj(obj);
     const node = this._objectNodeMapping.toNode(obj);
-    this._index.insertNodes([node]);
+    await this._index.insertNodes([node]);
   }
 
   get tools(): Record<string, BaseTool> {

packages/core/src/storage/docStore/KVDocumentStore.ts

@@ -70,7 +70,7 @@ export class KVDocumentStore extends BaseDocumentStore {
         metadata.refDocId = doc.sourceNode.nodeId!;
       }
 
-      this.kvstore.put(nodeKey, metadata, this.metadataCollection);
+      await this.kvstore.put(nodeKey, metadata, this.metadataCollection);
     }
   }
 
@@ -126,9 +126,9 @@ export class KVDocumentStore extends BaseDocumentStore {
       !_.pull(refDocInfo.nodeIds, docId);
 
       if (refDocInfo.nodeIds.length > 0) {
-        this.kvstore.put(refDocId, refDocInfo, this.refDocCollection);
+        await this.kvstore.put(refDocId, refDocInfo, this.refDocCollection);
       }
-      this.kvstore.delete(refDocId, this.metadataCollection);
+      await this.kvstore.delete(refDocId, this.metadataCollection);
     }
   }
 

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

@@ -167,7 +167,7 @@ export class MilvusVectorStore implements VectorStore {
     refDocId: string,
     deleteOptions?: Omit<DeleteReq, "ids">,
   ): Promise<void> {
-    this.ensureCollection();
+    await this.ensureCollection();
 
     await this.milvusClient.delete({
       ids: [refDocId],

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

@@ -82,7 +82,9 @@ export class PGVectorStore implements VectorStore {
   private async getDb(): Promise<pg.Client> {
     if (!this.db) {
       try {
-        const { Client } = await import("pg");
+        const pg = await import("pg");
+        const { Client } = pg.default ? pg.default : pg;
+
         const { registerType } = await import("pgvector/pg");
         // Create DB connection
         // Read connection params from env - see comment block above
@@ -92,7 +94,7 @@ export class PGVectorStore implements VectorStore {
         await db.connect();
 
         // Check vector extension
-        db.query("CREATE EXTENSION IF NOT EXISTS vector");
+        await db.query("CREATE EXTENSION IF NOT EXISTS vector");
         await registerType(db);
 
         // Check schema, table(s), index(es)

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

@@ -142,11 +142,11 @@ export class PineconeVectorStore implements VectorStore {
    */
   async query(
     query: VectorStoreQuery,
-    options?: any,
+    _options?: any,
   ): Promise<VectorStoreQueryResult> {
     const filter = this.toPineconeFilter(query.filters);
 
-    var options: any = {
+    const defaultOptions: any = {
       vector: query.queryEmbedding,
       topK: query.similarityTopK,
       includeValues: true,
@@ -155,7 +155,7 @@ export class PineconeVectorStore implements VectorStore {
     };
 
     const idx = await this.index();
-    const results = await idx.query(options);
+    const results = await idx.query(defaultOptions);
 
     const idList = results.matches.map((row) => row.id);
     const records: FetchResponse<any> = await idx.fetch(idList);

packages/core/src/synthesizers/MultiModalResponseSynthesizer.ts

@@ -55,7 +55,6 @@ export class MultiModalResponseSynthesizer
   async synthesize({
     query,
     nodesWithScore,
-    parentEvent,
     stream,
   }: SynthesizeParamsStreaming | SynthesizeParamsNonStreaming): Promise<
     AsyncIterable<Response> | Response
@@ -90,7 +89,6 @@ export class MultiModalResponseSynthesizer
 
     const response = await llm.complete({
       prompt,
-      parentEvent,
     });
 
     return new Response(response.text, nodes);

packages/core/src/synthesizers/ResponseSynthesizer.ts

@@ -62,7 +62,6 @@ export class ResponseSynthesizer
   async synthesize({
     query,
     nodesWithScore,
-    parentEvent,
     stream,
   }: SynthesizeParamsStreaming | SynthesizeParamsNonStreaming): Promise<
     AsyncIterable<Response> | Response
@@ -75,7 +74,6 @@ export class ResponseSynthesizer
       const response = await this.responseBuilder.getResponse({
         query,
         textChunks,
-        parentEvent,
         stream,
       });
       return streamConverter(response, (chunk) => new Response(chunk, nodes));
@@ -83,7 +81,6 @@ export class ResponseSynthesizer
     const response = await this.responseBuilder.getResponse({
       query,
       textChunks,
-      parentEvent,
     });
     return new Response(response, nodes);
   }

packages/core/src/synthesizers/builders.ts

@@ -1,4 +1,3 @@
-import type { Event } from "../callbacks/CallbackManager.js";
 import type { LLM } from "../llm/index.js";
 import { streamConverter } from "../llm/utils.js";
 import type {
@@ -55,7 +54,6 @@ export class SimpleResponseBuilder implements ResponseBuilder {
   async getResponse({
     query,
     textChunks,
-    parentEvent,
     stream,
   }:
     | ResponseBuilderParamsStreaming
@@ -69,10 +67,10 @@ export class SimpleResponseBuilder implements ResponseBuilder {
 
     const prompt = this.textQATemplate(input);
     if (stream) {
-      const response = await this.llm.complete({ prompt, parentEvent, stream });
+      const response = await this.llm.complete({ prompt, stream });
       return streamConverter(response, (chunk) => chunk.text);
     } else {
-      const response = await this.llm.complete({ prompt, parentEvent, stream });
+      const response = await this.llm.complete({ prompt, stream });
       return response.text;
     }
   }
@@ -130,7 +128,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
   async getResponse({
     query,
     textChunks,
-    parentEvent,
     prevResponse,
     stream,
   }:
@@ -148,7 +145,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
           query,
           chunk,
           !!stream && lastChunk,
-          parentEvent,
         );
       } else {
         response = await this.refineResponseSingle(
@@ -156,7 +152,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
           query,
           chunk,
           !!stream && lastChunk,
-          parentEvent,
         );
       }
     }
@@ -168,7 +163,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
     queryStr: string,
     textChunk: string,
     stream: boolean,
-    parentEvent?: Event,
   ) {
     const textQATemplate: SimplePrompt = (input) =>
       this.textQATemplate({ ...input, query: queryStr });
@@ -184,7 +178,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
           prompt: textQATemplate({
             context: chunk,
           }),
-          parentEvent,
           stream: stream && lastChunk,
         });
       } else {
@@ -193,7 +186,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
           queryStr,
           chunk,
           stream && lastChunk,
-          parentEvent,
         );
       }
     }
@@ -207,7 +199,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
     queryStr: string,
     textChunk: string,
     stream: boolean,
-    parentEvent?: Event,
   ) {
     const refineTemplate: SimplePrompt = (input) =>
       this.refineTemplate({ ...input, query: queryStr });
@@ -224,7 +215,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
           context: chunk,
           existingAnswer: response as string,
         }),
-        parentEvent,
         stream: stream && lastChunk,
       });
     }
@@ -234,7 +224,6 @@ export class Refine extends PromptMixin implements ResponseBuilder {
   async complete(params: {
     prompt: string;
     stream: boolean;
-    parentEvent?: Event;
   }): Promise<AsyncIterable<string> | string> {
     if (params.stream) {
       const response = await this.llm.complete({ ...params, stream: true });
@@ -257,7 +246,6 @@ export class CompactAndRefine extends Refine {
   async getResponse({
     query,
     textChunks,
-    parentEvent,
     prevResponse,
     stream,
   }:
@@ -275,7 +263,6 @@ export class CompactAndRefine extends Refine {
     const params = {
       query,
       textChunks: newTexts,
-      parentEvent,
       prevResponse,
     };
     if (stream) {
@@ -328,7 +315,6 @@ export class TreeSummarize extends PromptMixin implements ResponseBuilder {
   async getResponse({
     query,
     textChunks,
-    parentEvent,
     stream,
   }:
     | ResponseBuilderParamsStreaming
@@ -351,7 +337,6 @@ export class TreeSummarize extends PromptMixin implements ResponseBuilder {
           context: packedTextChunks[0],
           query,
         }),
-        parentEvent,
       };
       if (stream) {
         const response = await this.llm.complete({ ...params, stream });
@@ -366,7 +351,6 @@ export class TreeSummarize extends PromptMixin implements ResponseBuilder {
               context: chunk,
               query,
             }),
-            parentEvent,
           }),
         ),
       );

packages/core/src/synthesizers/types.ts

@@ -1,4 +1,3 @@
-import type { Event } from "../callbacks/CallbackManager.js";
 import type { NodeWithScore } from "../Node.js";
 import type { PromptMixin } from "../prompts/Mixin.js";
 import type { Response } from "../Response.js";
@@ -6,7 +5,6 @@ import type { Response } from "../Response.js";
 export interface SynthesizeParamsBase {
   query: string;
   nodesWithScore: NodeWithScore[];
-  parentEvent?: Event;
 }
 
 export interface SynthesizeParamsStreaming extends SynthesizeParamsBase {
@@ -30,7 +28,6 @@ export interface BaseSynthesizer {
 export interface ResponseBuilderParamsBase {
   query: string;
   textChunks: string[];
-  parentEvent?: Event;
   prevResponse?: string;
 }
 

packages/core/src/types.ts

@@ -1,7 +1,6 @@
 /**
  * Top level types to avoid circular dependencies
  */
-import type { Event } from "./callbacks/CallbackManager.js";
 import type { Response } from "./Response.js";
 
 /**
@@ -9,7 +8,6 @@ import type { Response } from "./Response.js";
  */
 export interface QueryEngineParamsBase {
   query: string;
-  parentEvent?: Event;
 }
 
 export interface QueryEngineParamsStreaming extends QueryEngineParamsBase {

packages/core/tests/CallbackManager.test.ts

@@ -20,20 +20,13 @@ import { CallbackManager } from "llamaindex/callbacks/CallbackManager";
 import { OpenAIEmbedding } from "llamaindex/embeddings/index";
 import { SummaryIndex } from "llamaindex/indices/summary/index";
 import { VectorStoreIndex } from "llamaindex/indices/vectorStore/index";
-import { OpenAI } from "llamaindex/llm/LLM";
+import { OpenAI } from "llamaindex/llm/open_ai";
 import {
   ResponseSynthesizer,
   SimpleResponseBuilder,
 } from "llamaindex/synthesizers/index";
 import { mockEmbeddingModel, mockLlmGeneration } from "./utility/mockOpenAI.js";
 
-// Mock the OpenAI getOpenAISession function during testing
-vi.mock("llamaindex/llm/open_ai", () => {
-  return {
-    getOpenAISession: vi.fn().mockImplementation(() => null),
-  };
-});
-
 describe("CallbackManager: onLLMStream and onRetrieve", () => {
   let serviceContext: ServiceContext;
   let streamCallbackData: StreamCallbackResponse[] = [];
@@ -88,12 +81,6 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
     expect(response.toString()).toBe("MOCK_TOKEN_1-MOCK_TOKEN_2");
     expect(streamCallbackData).toEqual([
       {
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "llmPredict",
-          tags: ["final"],
-        },
         index: 0,
         token: {
           id: "id",
@@ -104,12 +91,6 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
         },
       },
       {
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "llmPredict",
-          tags: ["final"],
-        },
         index: 1,
         token: {
           id: "id",
@@ -120,12 +101,6 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
         },
       },
       {
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "llmPredict",
-          tags: ["final"],
-        },
         index: 2,
         isDone: true,
       },
@@ -134,19 +109,8 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
       {
         query: query,
         nodes: expect.any(Array),
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "retrieve",
-          tags: ["final"],
-        },
       },
     ]);
-    // both retrieval and streaming should have
-    // the same parent event
-    expect(streamCallbackData[0].event.parentId).toBe(
-      retrieveCallbackData[0].event.parentId,
-    );
   });
 
   test("For SummaryIndex w/ a SummaryIndexRetriever", async () => {
@@ -169,12 +133,6 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
     expect(response.toString()).toBe("MOCK_TOKEN_1-MOCK_TOKEN_2");
     expect(streamCallbackData).toEqual([
       {
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "llmPredict",
-          tags: ["final"],
-        },
         index: 0,
         token: {
           id: "id",
@@ -185,12 +143,6 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
         },
       },
       {
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "llmPredict",
-          tags: ["final"],
-        },
         index: 1,
         token: {
           id: "id",
@@ -201,12 +153,6 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
         },
       },
       {
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "llmPredict",
-          tags: ["final"],
-        },
         index: 2,
         isDone: true,
       },
@@ -215,18 +161,7 @@ describe("CallbackManager: onLLMStream and onRetrieve", () => {
       {
         query: query,
         nodes: expect.any(Array),
-        event: {
-          id: expect.any(String),
-          parentId: expect.any(String),
-          type: "retrieve",
-          tags: ["final"],
-        },
       },
     ]);
-    // both retrieval and streaming should have
-    // the same parent event
-    expect(streamCallbackData[0].event.parentId).toBe(
-      retrieveCallbackData[0].event.parentId,
-    );
   });
 });

packages/core/tests/Embedding.test.ts

@@ -3,16 +3,9 @@ import {
   SimilarityType,
   similarity,
 } from "llamaindex/embeddings/index";
-import { beforeAll, describe, expect, test, vi } from "vitest";
+import { beforeAll, describe, expect, test } from "vitest";
 import { mockEmbeddingModel } from "./utility/mockOpenAI.js";
 
-// Mock the OpenAI getOpenAISession function during testing
-vi.mock("llamaindex/llm/open_ai", () => {
-  return {
-    getOpenAISession: vi.fn().mockImplementation(() => null),
-  };
-});
-
 describe("similarity", () => {
   test("throws error on mismatched lengths", () => {
     const embedding1 = [1, 2, 3];

packages/core/tests/MetadataExtractors.test.ts

@@ -8,7 +8,7 @@ import {
   SummaryExtractor,
   TitleExtractor,
 } from "llamaindex/extractors/index";
-import { OpenAI } from "llamaindex/llm/LLM";
+import { OpenAI } from "llamaindex/llm/open_ai";
 import { SimpleNodeParser } from "llamaindex/nodeParsers/index";
 import { afterAll, beforeAll, describe, expect, test, vi } from "vitest";
 import {
@@ -17,13 +17,6 @@ import {
   mockLlmGeneration,
 } from "./utility/mockOpenAI.js";
 
-// Mock the OpenAI getOpenAISession function during testing
-vi.mock("llamaindex/llm/open_ai", () => {
-  return {
-    getOpenAISession: vi.fn().mockImplementation(() => null),
-  };
-});
-
 describe("[MetadataExtractor]: Extractors should populate the metadata", () => {
   let serviceContext: ServiceContext;
 

packages/core/tests/Selectors.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, test, vi } from "vitest";
+import { describe, expect, test } from "vitest";
 // from unittest.mock import patch
 
 import { serviceContextFromDefaults } from "llamaindex/ServiceContext";
@@ -6,12 +6,6 @@ import { OpenAI } from "llamaindex/llm/index";
 import { LLMSingleSelector } from "llamaindex/selectors/index";
 import { mocStructuredkLlmGeneration } from "./utility/mockOpenAI.js";
 
-vi.mock("llamaindex/llm/open_ai", () => {
-  return {
-    getOpenAISession: vi.fn().mockImplementation(() => null),
-  };
-});
-
 describe("LLMSelector", () => {
   test("should be able to output a selection with a reason", async () => {
     const serviceContext = serviceContextFromDefaults({});

packages/core/tests/agent/OpenAIAgent.test.ts

@@ -1,7 +1,7 @@
 import { OpenAIAgent } from "llamaindex/agent/index";
 import { OpenAI } from "llamaindex/llm/index";
 import { FunctionTool } from "llamaindex/tools/index";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it } from "vitest";
 import { mockLlmToolCallGeneration } from "../utility/mockOpenAI.js";
 
 // Define a function to sum two numbers
@@ -24,12 +24,6 @@ const sumJSON = {
   required: ["a", "b"],
 };
 
-vi.mock("llamaindex/llm/open_ai", () => {
-  return {
-    getOpenAISession: vi.fn().mockImplementation(() => null),
-  };
-});
-
 describe("OpenAIAgent", () => {
   let openaiAgent: OpenAIAgent;
 

packages/core/tests/agent/runner/AgentRunner.test.ts

@@ -1,19 +1,13 @@
 import { OpenAIAgentWorker } from "llamaindex/agent/index";
 import { AgentRunner } from "llamaindex/agent/runner/base";
-import { OpenAI } from "llamaindex/llm/LLM";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { OpenAI } from "llamaindex/llm/open_ai";
+import { beforeEach, describe, expect, it } from "vitest";
 
 import {
   DEFAULT_LLM_TEXT_OUTPUT,
   mockLlmGeneration,
 } from "../../utility/mockOpenAI.js";
 
-vi.mock("llamaindex/llm/open_ai", () => {
-  return {
-    getOpenAISession: vi.fn().mockImplementation(() => null),
-  };
-});
-
 describe("Agent Runner", () => {
   let agentRunner: AgentRunner;