refactor: move chat engine & retriever into core (#1242)

.changeset/four-beers-kick.md

@@ -0,0 +1,12 @@
+---
+"@llamaindex/core": minor
+"llamaindex": minor
+---
+
+refactor: move chat engine & retriever into core.
+
+This is a breaking change since `stream` option has moved to second parameter.
+
+- `chat` API in BaseChatEngine has changed, Move `stream` option into second parameter.
+- `chatHistory` in BaseChatEngine now returns `ChatMessage[] | Promise<ChatMessage[]>`, instead of `BaseMemory`
+- update `retrieve-end` type

README.md

@@ -167,11 +167,13 @@ export async function chatWithAgent(
       // ... adding your tools here
     ],
   });
-  const responseStream = await agent.chat({
-    stream: true,
-    message: question,
-    chatHistory: prevMessages,
-  });
+  const responseStream = await agent.chat(
+    {
+      message: question,
+      chatHistory: prevMessages,
+    },
+    true,
+  );
   const uiStream = createStreamableUI(<div>loading...</div>);
   responseStream
     .pipeTo(

examples/agent/azure_dynamic_session.ts

@@ -42,7 +42,6 @@ async function main() {
   const response = await agent.chat({
     message:
       "plot a chart of 5 random numbers and save it to /mnt/data/chart.png",
-    stream: false,
   });
 
   // Print the response

examples/agent/stream_openai_agent.ts

@@ -61,10 +61,12 @@ async function main() {
     tools: [functionTool, functionTool2],
   });
 
-  const stream = await agent.chat({
-    message: "Divide 16 by 2 then add 20",
-    stream: true,
-  });
+  const stream = await agent.chat(
+    {
+      message: "Divide 16 by 2 then add 20",
+    },
+    true,
+  );
 
   console.log("Response:");
 

examples/agent/wiki.ts

@@ -11,10 +11,12 @@ async function main() {
   });
 
   // Chat with the agent
-  const response = await agent.chat({
-    message: "Who was Goethe?",
-    stream: true,
-  });
+  const response = await agent.chat(
+    {
+      message: "Who was Goethe?",
+    },
+    true,
+  );
 
   for await (const { delta } of response) {
     process.stdout.write(delta);

examples/anthropic/chat_interactive.ts

@@ -18,14 +18,14 @@ import readline from "node:readline/promises";
   });
   const chatEngine = new SimpleChatEngine({
     llm,
-    chatHistory,
+    memory: chatHistory,
   });
   const rl = readline.createInterface({ input, output });
 
   while (true) {
     const query = await rl.question("User: ");
     process.stdout.write("Assistant: ");
-    const stream = await chatEngine.chat({ message: query, stream: true });
+    const stream = await chatEngine.chat({ message: query }, true);
     for await (const chunk of stream) {
       process.stdout.write(chunk.response);
     }

examples/chatEngine.ts

@@ -24,7 +24,7 @@ async function main() {
 
   while (true) {
     const query = await rl.question("Query: ");
-    const stream = await chatEngine.chat({ message: query, stream: true });
+    const stream = await chatEngine.chat({ message: query }, true);
     console.log();
     for await (const chunk of stream) {
       process.stdout.write(chunk.response);

examples/chatHistory.ts

@@ -24,11 +24,13 @@ async function main() {
 
   while (true) {
     const query = await rl.question("Query: ");
-    const stream = await chatEngine.chat({
-      message: query,
-      chatHistory,
-      stream: true,
-    });
+    const stream = await chatEngine.chat(
+      {
+        message: query,
+        chatHistory,
+      },
+      true,
+    );
     if (chatHistory.getLastSummary()) {
       // Print the summary of the conversation so far that is produced by the SummaryChatHistory
       console.log(`Summary: ${chatHistory.getLastSummary()?.content}`);

examples/cloud/chat.ts

@@ -18,7 +18,7 @@ async function main() {
 
   while (true) {
     const query = await rl.question("User: ");
-    const stream = await chatEngine.chat({ message: query, stream: true });
+    const stream = await chatEngine.chat({ message: query }, true);
     for await (const chunk of stream) {
       process.stdout.write(chunk.response);
     }

examples/multimodal/context.ts

@@ -1,4 +1,5 @@
 // call pnpm tsx multimodal/load.ts first to init the storage
+import { extractText } from "@llamaindex/core/utils";
 import {
   ContextChatEngine,
   NodeWithScore,
@@ -25,8 +26,9 @@ Settings.callbackManager.on("retrieve-end", (event) => {
   const textNodes = nodes.filter(
     (node: NodeWithScore) => node.node.type === ObjectType.TEXT,
   );
+  const text = extractText(query);
   console.log(
-    `Retrieved ${textNodes.length} text nodes and ${imageNodes.length} image nodes for query: ${query}`,
+    `Retrieved ${textNodes.length} text nodes and ${imageNodes.length} image nodes for query: ${text}`,
   );
 });
 

examples/multimodal/rag.ts

@@ -1,3 +1,4 @@
+import { extractText } from "@llamaindex/core/utils";
 import {
   getResponseSynthesizer,
   OpenAI,
@@ -16,7 +17,8 @@ Settings.llm = new OpenAI({ model: "gpt-4-turbo", maxTokens: 512 });
 // Update callbackManager
 Settings.callbackManager.on("retrieve-end", (event) => {
   const { nodes, query } = event.detail;
-  console.log(`Retrieved ${nodes.length} nodes for query: ${query}`);
+  const text = extractText(query);
+  console.log(`Retrieved ${nodes.length} nodes for query: ${text}`);
 });
 
 async function main() {

packages/autotool/examples/02_nextjs/actions.ts

@@ -14,10 +14,12 @@ export async function chatWithAI(message: string): Promise<ReactNode> {
   const uiStream = createStreamableUI();
   runWithStreamableUI(uiStream, () =>
     agent
-      .chat({
-        stream: true,
-        message,
-      })
+      .chat(
+        {
+          message,
+        },
+        true,
+      )
       .then(async (responseStream) => {
         return responseStream.pipeTo(
           new WritableStream({

packages/core/package.json

@@ -199,6 +199,34 @@
         "types": "./dist/response-synthesizers/index.d.ts",
         "default": "./dist/response-synthesizers/index.js"
       }
+    },
+    "./chat-engine": {
+      "require": {
+        "types": "./dist/chat-engine/index.d.cts",
+        "default": "./dist/chat-engine/index.cjs"
+      },
+      "import": {
+        "types": "./dist/chat-engine/index.d.ts",
+        "default": "./dist/chat-engine/index.js"
+      },
+      "default": {
+        "types": "./dist/chat-engine/index.d.ts",
+        "default": "./dist/chat-engine/index.js"
+      }
+    },
+    "./retriever": {
+      "require": {
+        "types": "./dist/retriever/index.d.cts",
+        "default": "./dist/retriever/index.cjs"
+      },
+      "import": {
+        "types": "./dist/retriever/index.d.ts",
+        "default": "./dist/retriever/index.js"
+      },
+      "default": {
+        "types": "./dist/retriever/index.d.ts",
+        "default": "./dist/retriever/index.js"
+      }
     }
   },
   "files": [

packages/core/src/chat-engine/index.ts

@@ -0,0 +1,28 @@
+import type { ChatMessage, MessageContent } from "../llms";
+import type { BaseMemory } from "../memory";
+import { EngineResponse } from "../schema";
+
+export interface ChatEngineParams<
+  AdditionalMessageOptions extends object = object,
+> {
+  message: MessageContent;
+  /**
+   * Optional chat history if you want to customize the chat history.
+   */
+  chatHistory?:
+    | ChatMessage<AdditionalMessageOptions>[]
+    | BaseMemory<AdditionalMessageOptions>;
+}
+
+export abstract class BaseChatEngine {
+  abstract chat(
+    params: ChatEngineParams,
+    stream?: false,
+  ): Promise<EngineResponse>;
+  abstract chat(
+    params: ChatEngineParams,
+    stream: true,
+  ): Promise<AsyncIterable<EngineResponse>>;
+
+  abstract chatHistory: ChatMessage[] | Promise<ChatMessage[]>;
+}

packages/core/src/global/settings/callback-manager.ts

@@ -11,6 +11,7 @@ import type {
   SynthesizeEndEvent,
   SynthesizeStartEvent,
 } from "../../response-synthesizers";
+import type { RetrieveEndEvent, RetrieveStartEvent } from "../../retriever";
 import { TextNode } from "../../schema";
 import { EventCaller, getEventCaller } from "../../utils";
 import type { UUID } from "../type";
@@ -69,6 +70,8 @@ export interface LlamaIndexEventMaps {
   "query-end": QueryEndEvent;
   "synthesize-start": SynthesizeStartEvent;
   "synthesize-end": SynthesizeEndEvent;
+  "retrieve-start": RetrieveStartEvent;
+  "retrieve-end": RetrieveEndEvent;
 }
 
 export class LlamaIndexCustomEvent<T = any> extends CustomEvent<T> {

packages/core/src/query-engine/base.ts

@@ -2,7 +2,7 @@ import { randomUUID } from "@llamaindex/env";
 import { Settings } from "../global";
 import type { MessageContent } from "../llms";
 import { PromptMixin } from "../prompts";
-import { EngineResponse } from "../schema";
+import { EngineResponse, type NodeWithScore } from "../schema";
 import { wrapEventCaller } from "../utils";
 
 /**
@@ -28,6 +28,12 @@ export abstract class BaseQueryEngine extends PromptMixin {
     super();
   }
 
+  async retrieve(params: QueryType): Promise<NodeWithScore[]> {
+    throw new Error(
+      "This query engine does not support retrieve, use query directly",
+    );
+  }
+
   query(
     strOrQueryBundle: QueryType,
     stream: true,

packages/core/src/retriever/index.ts

@@ -0,0 +1,112 @@
+import { randomUUID } from "@llamaindex/env";
+import { Settings } from "../global";
+import type { MessageContent } from "../llms";
+import { PromptMixin } from "../prompts";
+import type { QueryBundle, QueryType } from "../query-engine";
+import { BaseNode, IndexNode, type NodeWithScore, ObjectType } from "../schema";
+
+export type RetrieveParams = {
+  query: MessageContent;
+  preFilters?: unknown;
+};
+
+export type RetrieveStartEvent = {
+  id: string;
+  query: QueryBundle;
+};
+
+export type RetrieveEndEvent = {
+  id: string;
+  query: QueryBundle;
+  nodes: NodeWithScore[];
+};
+
+export abstract class BaseRetriever extends PromptMixin {
+  objectMap: Map<string, unknown> = new Map();
+
+  protected _updatePrompts() {}
+  protected _getPrompts() {
+    return {};
+  }
+
+  protected _getPromptModules() {
+    return {};
+  }
+
+  protected constructor() {
+    super();
+  }
+
+  public async retrieve(params: QueryType): Promise<NodeWithScore[]> {
+    const cb = Settings.callbackManager;
+    const queryBundle = typeof params === "string" ? { query: params } : params;
+    const id = randomUUID();
+    cb.dispatchEvent("retrieve-start", { id, query: queryBundle });
+    let response = await this._retrieve(queryBundle);
+    response = await this._handleRecursiveRetrieval(queryBundle, response);
+    cb.dispatchEvent("retrieve-end", {
+      id,
+      query: queryBundle,
+      nodes: response,
+    });
+    return response;
+  }
+
+  abstract _retrieve(params: QueryBundle): Promise<NodeWithScore[]>;
+
+  async _handleRecursiveRetrieval(
+    params: QueryBundle,
+    nodes: NodeWithScore[],
+  ): Promise<NodeWithScore[]> {
+    const retrievedNodes = [];
+    for (const { node, score = 1.0 } of nodes) {
+      if (node.type === ObjectType.INDEX) {
+        const indexNode = node as IndexNode;
+        const object = this.objectMap.get(indexNode.indexId);
+        if (object !== undefined) {
+          retrievedNodes.push(
+            ...this._retrieveFromObject(object, params, score),
+          );
+        } else {
+          retrievedNodes.push({ node, score });
+        }
+      } else {
+        retrievedNodes.push({ node, score });
+      }
+    }
+    return nodes;
+  }
+
+  _retrieveFromObject(
+    object: unknown,
+    queryBundle: QueryBundle,
+    score: number,
+  ): NodeWithScore[] {
+    if (object == null) {
+      throw new TypeError("Object is not retrievable");
+    }
+    if (typeof object !== "object") {
+      throw new TypeError("Object is not retrievable");
+    }
+    if ("node" in object && object.node instanceof BaseNode) {
+      return [
+        {
+          node: object.node,
+          score:
+            "score" in object && typeof object.score === "number"
+              ? object.score
+              : score,
+        },
+      ];
+    }
+    if (object instanceof BaseNode) {
+      return [{ node: object, score }];
+    } else {
+      // todo: support other types
+      // BaseQueryEngine
+      // BaseRetriever
+      // QueryComponent
+      throw new TypeError("Object is not retrievable");
+    }
+  }
+}

packages/llamaindex/e2e/examples/cloudflare-worker-agent/src/index.ts

@@ -11,10 +11,12 @@ export default {
       tools: [],
     });
     console.log(1);
-    const responseStream = await agent.chat({
-      stream: true,
-      message: "Hello? What is the weather today?",
-    });
+    const responseStream = await agent.chat(
+      {
+        message: "Hello? What is the weather today?",
+      },
+      true,
+    );
     console.log(2);
     const textEncoder = new TextEncoder();
     const response = responseStream.pipeThrough<Uint8Array>(

packages/llamaindex/e2e/examples/nextjs-agent/src/actions/index.tsx

@@ -10,11 +10,13 @@ export async function chatWithAgent(
   const agent = new OpenAIAgent({
     tools: [],
   });
-  const responseStream = await agent.chat({
-    stream: true,
-    message: question,
-    chatHistory: prevMessages,
-  });
+  const responseStream = await agent.chat(
+    {
+      message: question,
+      chatHistory: prevMessages,
+    },
+    true,
+  );
   const uiStream = createStreamableUI(<div>loading...</div>);
   responseStream
     .pipeTo(

packages/llamaindex/e2e/node/openai.e2e.ts

@@ -322,10 +322,12 @@ await test("agent stream", async (t) => {
       tools: [sumNumbersTool, divideNumbersTool],
     });
 
-    const stream = await agent.chat({
-      message: "Divide 16 by 2 then add 20",
-      stream: true,
-    });
+    const stream = await agent.chat(
+      {
+        message: "Divide 16 by 2 then add 20",
+      },
+      true,
+    );
 
     let message = "";
 

packages/llamaindex/e2e/node/react.e2e.ts

@@ -20,7 +20,6 @@ await test("react agent", async (t) => {
       tools: [getWeatherTool],
     });
     const response = await agent.chat({
-      stream: false,
       message: "What is the weather like in San Francisco?",
     });
 
@@ -35,10 +34,12 @@ await test("react agent stream", async (t) => {
       tools: [getWeatherTool],
     });
 
-    const stream = await agent.chat({
-      stream: true,
-      message: "What is the weather like in San Francisco?",
-    });
+    const stream = await agent.chat(
+      {
+        message: "What is the weather like in San Francisco?",
+      },
+      true,
+    );
 
     let content = "";
     for await (const response of stream) {

packages/llamaindex/src/Retriever.ts

@@ -1,20 +0,0 @@
-import type { NodeWithScore } from "@llamaindex/core/schema";
-import type { ServiceContext } from "./ServiceContext.js";
-import type { MessageContent } from "./index.edge.js";
-
-export type RetrieveParams = {
-  query: MessageContent;
-  preFilters?: unknown;
-};
-
-/**
- * Retrievers retrieve the nodes that most closely match our query in similarity.
- */
-export interface BaseRetriever {
-  retrieve(params: RetrieveParams): Promise<NodeWithScore[]>;
-
-  /**
-   * @deprecated to be deprecated soon
-   */
-  serviceContext?: ServiceContext | undefined;
-}

packages/llamaindex/src/agent/anthropic.ts

@@ -1,9 +1,6 @@
+import type { ChatEngineParams } from "@llamaindex/core/chat-engine";
+import type { EngineResponse } from "@llamaindex/core/schema";
 import { Settings } from "../Settings.js";
-import type {
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-  EngineResponse,
-} from "../index.edge.js";
 import { Anthropic } from "../llm/anthropic.js";
 import { LLMAgent, LLMAgentWorker, type LLMAgentParams } from "./llm.js";
 
@@ -24,12 +21,10 @@ export class AnthropicAgent extends LLMAgent {
     });
   }
 
-  async chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
-  async chat(params: ChatEngineParamsStreaming): Promise<never>;
-  override async chat(
-    params: ChatEngineParamsNonStreaming | ChatEngineParamsStreaming,
-  ) {
-    if (params.stream) {
+  async chat(params: ChatEngineParams, stream?: false): Promise<EngineResponse>;
+  async chat(params: ChatEngineParams, stream: true): Promise<never>;
+  override async chat(params: ChatEngineParams, stream?: boolean) {
+    if (stream) {
       // Anthropic does support this, but looks like it's not supported in the LITS LLM
       throw new Error("Anthropic does not support streaming");
     }

packages/llamaindex/src/agent/base.ts

@@ -1,3 +1,7 @@
+import {
+  BaseChatEngine,
+  type ChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type {
   BaseToolWithCall,
   ChatMessage,
@@ -10,11 +14,6 @@ import { EngineResponse } from "@llamaindex/core/schema";
 import { wrapEventCaller } from "@llamaindex/core/utils";
 import { randomUUID } from "@llamaindex/env";
 import { Settings } from "../Settings.js";
-import {
-  type ChatEngine,
-  type ChatEngineParamsNonStreaming,
-  type ChatEngineParamsStreaming,
-} from "../engines/chat/index.js";
 import { consoleLogger, emptyLogger } from "../internal/logger.js";
 import { isReadableStream } from "../internal/utils.js";
 import { ObjectRetriever } from "../objects/index.js";
@@ -207,8 +206,7 @@ export abstract class AgentRunner<
   >
     ? AdditionalMessageOptions
     : never,
-> implements ChatEngine
-{
+> extends BaseChatEngine {
   readonly #llm: AI;
   readonly #tools:
     | BaseToolWithCall[]
@@ -259,6 +257,7 @@ export abstract class AgentRunner<
   protected constructor(
     params: AgentRunnerParams<AI, Store, AdditionalMessageOptions>,
   ) {
+    super();
     const { llm, chatHistory, systemPrompt, runner, tools, verbose } = params;
     this.#llm = llm;
     this.#chatHistory = chatHistory;
@@ -345,13 +344,15 @@ export abstract class AgentRunner<
     });
   }
 
-  async chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
+  async chat(params: ChatEngineParams, stream?: false): Promise<EngineResponse>;
   async chat(
-    params: ChatEngineParamsStreaming,
+    params: ChatEngineParams,
+    stream: true,
   ): Promise<ReadableStream<EngineResponse>>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsNonStreaming | ChatEngineParamsStreaming,
+    params: ChatEngineParams,
+    stream?: boolean,
   ): Promise<EngineResponse | ReadableStream<EngineResponse>> {
     let chatHistory: ChatMessage<AdditionalMessageOptions>[] = [];
 
@@ -363,12 +364,7 @@ export abstract class AgentRunner<
         params.chatHistory as ChatMessage<AdditionalMessageOptions>[];
     }
 
-    const task = this.createTask(
-      params.message,
-      !!params.stream,
-      false,
-      chatHistory,
-    );
+    const task = this.createTask(params.message, !!stream, false, chatHistory);
     for await (const stepOutput of task) {
       // update chat history for each round
       this.#chatHistory = [...stepOutput.taskStep.context.store.messages];

packages/llamaindex/src/cloud/LlamaCloudIndex.ts

@@ -1,7 +1,6 @@
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import type { Document, TransformComponent } from "@llamaindex/core/schema";
-import type { BaseRetriever } from "../Retriever.js";
 import { RetrieverQueryEngine } from "../engines/query/RetrieverQueryEngine.js";
 import type { BaseNodePostprocessor } from "../postprocessors/types.js";
 import type { CloudRetrieveParams } from "./LlamaCloudRetriever.js";
@@ -12,6 +11,7 @@ import { getAppBaseUrl, getProjectId, initService } from "./utils.js";
 
 import { PipelinesService, ProjectsService } from "@llamaindex/cloud/api";
 import { SentenceSplitter } from "@llamaindex/core/node-parser";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import { getEnv } from "@llamaindex/env";
 import { OpenAIEmbedding } from "@llamaindex/openai";
 import { Settings } from "../Settings.js";

packages/llamaindex/src/cloud/LlamaCloudRetriever.ts

@@ -4,11 +4,12 @@ import {
   type RetrievalParams,
   type TextNodeWithScore,
 } from "@llamaindex/cloud/api";
-import { DEFAULT_PROJECT_NAME, Settings } from "@llamaindex/core/global";
+import { DEFAULT_PROJECT_NAME } from "@llamaindex/core/global";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import type { NodeWithScore } from "@llamaindex/core/schema";
 import { jsonToNode, ObjectType } from "@llamaindex/core/schema";
-import { extractText, wrapEventCaller } from "@llamaindex/core/utils";
-import type { BaseRetriever, RetrieveParams } from "../Retriever.js";
+import { extractText } from "@llamaindex/core/utils";
 import type { ClientParams, CloudConstructorParams } from "./type.js";
 import { getProjectId, initService } from "./utils.js";
 
@@ -17,7 +18,7 @@ export type CloudRetrieveParams = Omit<
   "query" | "search_filters" | "dense_similarity_top_k"
 > & { similarityTopK?: number; filters?: MetadataFilters };
 
-export class LlamaCloudRetriever implements BaseRetriever {
+export class LlamaCloudRetriever extends BaseRetriever {
   clientParams: ClientParams;
   retrieveParams: CloudRetrieveParams;
   organizationId?: string;
@@ -42,6 +43,7 @@ export class LlamaCloudRetriever implements BaseRetriever {
   }
 
   constructor(params: CloudConstructorParams & CloudRetrieveParams) {
+    super();
     this.clientParams = { apiKey: params.apiKey, baseUrl: params.baseUrl };
     initService(this.clientParams);
     this.retrieveParams = params;
@@ -54,11 +56,7 @@ export class LlamaCloudRetriever implements BaseRetriever {
     }
   }
 
-  @wrapEventCaller
-  async retrieve({
-    query,
-    preFilters,
-  }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
     const { data: pipelines } =
       await PipelinesService.searchPipelinesApiV1PipelinesGet({
         query: {
@@ -97,19 +95,11 @@ export class LlamaCloudRetriever implements BaseRetriever {
         body: {
           ...this.retrieveParams,
           query: extractText(query),
-          search_filters:
-            this.retrieveParams.filters ?? (preFilters as MetadataFilters),
+          search_filters: this.retrieveParams.filters as MetadataFilters,
           dense_similarity_top_k: this.retrieveParams.similarityTopK!,
         },
       });
 
-    const nodesWithScores = this.resultNodesToNodeWithScore(
-      results.retrieval_nodes,
-    );
-    Settings.callbackManager.dispatchEvent("retrieve-end", {
-      query,
-      nodes: nodesWithScores,
-    });
-    return nodesWithScores;
+    return this.resultNodesToNodeWithScore(results.retrieval_nodes);
   }
 }

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

@@ -1,10 +1,13 @@
+import {
+  BaseChatEngine,
+  type ChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type { ChatMessage, LLM } from "@llamaindex/core/llms";
 import { BaseMemory, ChatMemoryBuffer } from "@llamaindex/core/memory";
 import {
   type CondenseQuestionPrompt,
   defaultCondenseQuestionPrompt,
   type ModuleRecord,
-  PromptMixin,
 } from "@llamaindex/core/prompts";
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { EngineResponse } from "@llamaindex/core/schema";
@@ -16,11 +19,6 @@ import {
 } from "@llamaindex/core/utils";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { llmFromSettingsOrContext } from "../../Settings.js";
-import type {
-  ChatEngine,
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-} from "./types.js";
 
 /**
  * CondenseQuestionChatEngine is used in conjunction with a Index (for example VectorStoreIndex).
@@ -32,16 +30,16 @@ import type {
  * underlying data. It performs less well when the chat messages are not questions about the
  * data, or are very referential to previous context.
  */
-
-export class CondenseQuestionChatEngine
-  extends PromptMixin
-  implements ChatEngine
-{
+export class CondenseQuestionChatEngine extends BaseChatEngine {
   queryEngine: BaseQueryEngine;
-  chatHistory: BaseMemory;
+  memory: BaseMemory;
   llm: LLM;
   condenseMessagePrompt: CondenseQuestionPrompt;
 
+  get chatHistory() {
+    return this.memory.getMessages();
+  }
+
   constructor(init: {
     queryEngine: BaseQueryEngine;
     chatHistory: ChatMessage[];
@@ -51,7 +49,7 @@ export class CondenseQuestionChatEngine
     super();
 
     this.queryEngine = init.queryEngine;
-    this.chatHistory = new ChatMemoryBuffer({
+    this.memory = new ChatMemoryBuffer({
       chatHistory: init?.chatHistory,
     });
     this.llm = llmFromSettingsOrContext(init?.serviceContext);
@@ -88,15 +86,17 @@ export class CondenseQuestionChatEngine
     });
   }
 
+  chat(params: ChatEngineParams, stream?: false): Promise<EngineResponse>;
   chat(
-    params: ChatEngineParamsStreaming,
+    params: ChatEngineParams,
+    stream: true,
   ): Promise<AsyncIterable<EngineResponse>>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
+    params: ChatEngineParams,
+    stream = false,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
-    const { message, stream } = params;
+    const { message } = params;
     const chatHistory = params.chatHistory
       ? new ChatMemoryBuffer({
           chatHistory:
@@ -104,7 +104,7 @@ export class CondenseQuestionChatEngine
               ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
-      : this.chatHistory;
+      : this.memory;
 
     const condensedQuestion = (
       await this.condenseQuestion(chatHistory, extractText(message))
@@ -140,6 +140,6 @@ export class CondenseQuestionChatEngine
   }
 
   reset() {
-    this.chatHistory.reset();
+    this.memory.reset();
   }
 }

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

@@ -1,3 +1,7 @@
+import type {
+  BaseChatEngine,
+  ChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type {
   ChatMessage,
   LLM,
@@ -11,6 +15,7 @@ import {
   PromptMixin,
   type PromptsRecord,
 } from "@llamaindex/core/prompts";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import { EngineResponse, MetadataMode } from "@llamaindex/core/schema";
 import {
   extractText,
@@ -18,27 +23,25 @@ import {
   streamReducer,
   wrapEventCaller,
 } from "@llamaindex/core/utils";
-import type { BaseRetriever } from "../../Retriever.js";
 import { Settings } from "../../Settings.js";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
 import { DefaultContextGenerator } from "./DefaultContextGenerator.js";
-import type {
-  ChatEngine,
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-  ContextGenerator,
-} from "./types.js";
+import type { ContextGenerator } from "./types.js";
 
 /**
  * ContextChatEngine uses the Index to get the appropriate context for each query.
  * The context is stored in the system prompt, and the chat history is chunk: ChatResponseChunk, nodes?: NodeWithScore<import("/Users/marcus/code/llamaindex/LlamaIndexTS/packages/core/src/Node").Metadata>[], nodes?: NodeWithScore<import("/Users/marcus/code/llamaindex/LlamaIndexTS/packages/core/src/Node").Metadata>[]lowing the appropriate context to be surfaced for each query.
  */
-export class ContextChatEngine extends PromptMixin implements ChatEngine {
+export class ContextChatEngine extends PromptMixin implements BaseChatEngine {
   chatModel: LLM;
-  chatHistory: BaseMemory;
+  memory: BaseMemory;
   contextGenerator: ContextGenerator & PromptMixin;
   systemPrompt?: string | undefined;
 
+  get chatHistory() {
+    return this.memory.getMessages();
+  }
+
   constructor(init: {
     retriever: BaseRetriever;
     chatModel?: LLM | undefined;
@@ -50,7 +53,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
   }) {
     super();
     this.chatModel = init.chatModel ?? Settings.llm;
-    this.chatHistory = new ChatMemoryBuffer({ chatHistory: init?.chatHistory });
+    this.memory = new ChatMemoryBuffer({ chatHistory: init?.chatHistory });
     this.contextGenerator = new DefaultContextGenerator({
       retriever: init.retriever,
       contextSystemPrompt: init?.contextSystemPrompt,
@@ -79,15 +82,17 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
     };
   }
 
+  chat(params: ChatEngineParams, stream?: false): Promise<EngineResponse>;
   chat(
-    params: ChatEngineParamsStreaming,
+    params: ChatEngineParams,
+    stream: true,
   ): Promise<AsyncIterable<EngineResponse>>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
+    params: ChatEngineParams,
+    stream = false,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
-    const { message, stream } = params;
+    const { message } = params;
     const chatHistory = params.chatHistory
       ? new ChatMemoryBuffer({
           chatHistory:
@@ -95,7 +100,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
               ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
-      : this.chatHistory;
+      : this.memory;
     const requestMessages = await this.prepareRequestMessages(
       message,
       chatHistory,
@@ -125,7 +130,7 @@ export class ContextChatEngine extends PromptMixin implements ChatEngine {
   }
 
   reset() {
-    this.chatHistory.reset();
+    this.memory.reset();
   }
 
   private async prepareRequestMessages(

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

@@ -5,10 +5,10 @@ import {
   type ModuleRecord,
   PromptMixin,
 } from "@llamaindex/core/prompts";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import { MetadataMode, type NodeWithScore } from "@llamaindex/core/schema";
 import { createMessageContent } from "@llamaindex/core/utils";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
-import type { BaseRetriever } from "../../Retriever.js";
 import type { Context, ContextGenerator } from "./types.js";
 
 export class DefaultContextGenerator

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

@@ -1,3 +1,7 @@
+import type {
+  BaseChatEngine,
+  ChatEngineParams,
+} from "@llamaindex/core/chat-engine";
 import type { LLM } from "@llamaindex/core/llms";
 import { BaseMemory, ChatMemoryBuffer } from "@llamaindex/core/memory";
 import { EngineResponse } from "@llamaindex/core/schema";
@@ -7,34 +11,35 @@ import {
   wrapEventCaller,
 } from "@llamaindex/core/utils";
 import { Settings } from "../../Settings.js";
-import type {
-  ChatEngine,
-  ChatEngineParamsNonStreaming,
-  ChatEngineParamsStreaming,
-} from "./types.js";
 
 /**
  * SimpleChatEngine is the simplest possible chat engine. Useful for using your own custom prompts.
  */
 
-export class SimpleChatEngine implements ChatEngine {
-  chatHistory: BaseMemory;
+export class SimpleChatEngine implements BaseChatEngine {
+  memory: BaseMemory;
   llm: LLM;
 
+  get chatHistory() {
+    return this.memory.getMessages();
+  }
+
   constructor(init?: Partial<SimpleChatEngine>) {
-    this.chatHistory = init?.chatHistory ?? new ChatMemoryBuffer();
+    this.memory = init?.memory ?? new ChatMemoryBuffer();
     this.llm = init?.llm ?? Settings.llm;
   }
 
+  chat(params: ChatEngineParams, stream?: false): Promise<EngineResponse>;
   chat(
-    params: ChatEngineParamsStreaming,
+    params: ChatEngineParams,
+    stream: true,
   ): Promise<AsyncIterable<EngineResponse>>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<EngineResponse>;
   @wrapEventCaller
   async chat(
-    params: ChatEngineParamsStreaming | ChatEngineParamsNonStreaming,
+    params: ChatEngineParams,
+    stream = false,
   ): Promise<EngineResponse | AsyncIterable<EngineResponse>> {
-    const { message, stream } = params;
+    const { message } = params;
 
     const chatHistory = params.chatHistory
       ? new ChatMemoryBuffer({
@@ -43,7 +48,7 @@ export class SimpleChatEngine implements ChatEngine {
               ? await params.chatHistory.getMessages()
               : params.chatHistory,
         })
-      : this.chatHistory;
+      : this.memory;
     chatHistory.put({ content: message, role: "user" });
 
     if (stream) {
@@ -73,6 +78,6 @@ export class SimpleChatEngine implements ChatEngine {
   }
 
   reset() {
-    this.chatHistory.reset();
+    this.memory.reset();
   }
 }

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

@@ -1,58 +1,10 @@
-import type { ChatMessage, MessageContent } from "@llamaindex/core/llms";
-import type { BaseMemory } from "@llamaindex/core/memory";
-import { EngineResponse, type NodeWithScore } from "@llamaindex/core/schema";
-
-/**
- * Represents the base parameters for ChatEngine.
- */
-export interface ChatEngineParamsBase {
-  message: MessageContent;
-  /**
-   * Optional chat history if you want to customize the chat history.
-   */
-  chatHistory?: ChatMessage[] | BaseMemory;
-  /**
-   * Optional flag to enable verbose mode.
-   * @default false
-   */
-  verbose?: boolean;
-}
-
-export interface ChatEngineParamsStreaming extends ChatEngineParamsBase {
-  stream: true;
-}
-
-export interface ChatEngineParamsNonStreaming extends ChatEngineParamsBase {
-  stream?: false | null;
-}
-
-/**
- * A ChatEngine is used to handle back and forth chats between the application and the LLM.
- */
-export interface ChatEngine<
-  // synchronous response
-  R = EngineResponse,
-  // asynchronous response
-  AR extends AsyncIterable<unknown> = AsyncIterable<R>,
-> {
-  /**
-   * Send message along with the class's current chat history to the LLM.
-   * @param params
-   */
-  chat(params: ChatEngineParamsStreaming): Promise<AR>;
-  chat(params: ChatEngineParamsNonStreaming): Promise<R>;
-
-  /**
-   * Resets the chat history so that it's empty.
-   */
-  reset(): void;
-}
+import type { ChatMessage } from "@llamaindex/core/llms";
+import type { NodeWithScore } from "@llamaindex/core/schema";
 
 export interface Context {
   message: ChatMessage;
   nodes: NodeWithScore[];
 }
-
 /**
  * A ContextGenerator is used to generate a context based on a message's text content
  */

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

@@ -1,10 +1,11 @@
-import { BaseQueryEngine } from "@llamaindex/core/query-engine";
+import type { MessageContent } from "@llamaindex/core/llms";
+import { BaseQueryEngine, type QueryType } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import { getResponseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import { type NodeWithScore } from "@llamaindex/core/schema";
 import { extractText } from "@llamaindex/core/utils";
 import type { BaseNodePostprocessor } from "../../postprocessors/index.js";
-import type { BaseRetriever } from "../../Retriever.js";
 
 /**
  * A query engine that uses a retriever to query an index and then synthesizes the response.
@@ -67,7 +68,10 @@ export class RetrieverQueryEngine extends BaseQueryEngine {
     };
   }
 
-  private async applyNodePostprocessors(nodes: NodeWithScore[], query: string) {
+  private async applyNodePostprocessors(
+    nodes: NodeWithScore[],
+    query: MessageContent,
+  ) {
     let nodesWithScore = nodes;
 
     for (const postprocessor of this.nodePostprocessors) {
@@ -80,12 +84,10 @@ export class RetrieverQueryEngine extends BaseQueryEngine {
     return nodesWithScore;
   }
 
-  private async retrieve(query: string) {
-    const nodes = await this.retriever.retrieve({
-      query,
-      preFilters: this.preFilters,
-    });
+  override async retrieve(query: QueryType) {
+    const nodes = await this.retriever.retrieve(query);
 
-    return await this.applyNodePostprocessors(nodes, query);
+    const messageContent = typeof query === "string" ? query : query.query;
+    return await this.applyNodePostprocessors(nodes, messageContent);
   }
 }

packages/llamaindex/src/index.edge.ts

@@ -1,6 +1,6 @@
 import type { AgentEndEvent, AgentStartEvent } from "./agent/types.js";
-import type { RetrievalEndEvent, RetrievalStartEvent } from "./llm/types.js";
 
+export * from "@llamaindex/core/chat-engine";
 export {
   CallbackManager,
   DEFAULT_BASE_URL,
@@ -35,12 +35,11 @@ export * from "@llamaindex/core/llms";
 export * from "@llamaindex/core/prompts";
 export * from "@llamaindex/core/query-engine";
 export * from "@llamaindex/core/response-synthesizers";
+export * from "@llamaindex/core/retriever";
 export * from "@llamaindex/core/schema";
 
 declare module "@llamaindex/core/global" {
   export interface LlamaIndexEventMaps {
-    "retrieve-start": RetrievalStartEvent;
-    "retrieve-end": RetrievalEndEvent;
     // agent events
     "agent-start": AgentStartEvent;
     "agent-end": AgentEndEvent;
@@ -66,7 +65,6 @@ export * from "./objects/index.js";
 export * from "./OutputParser.js";
 export * from "./postprocessors/index.js";
 export * from "./QuestionGenerator.js";
-export * from "./Retriever.js";
 export * from "./selectors/index.js";
 export * from "./ServiceContext.js";
 export { Settings } from "./Settings.js";

packages/llamaindex/src/indices/BaseIndex.ts

@@ -1,7 +1,7 @@
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import type { BaseRetriever } from "@llamaindex/core/retriever";
 import type { BaseNode, Document } from "@llamaindex/core/schema";
-import type { BaseRetriever } from "../Retriever.js";
 import type { ServiceContext } from "../ServiceContext.js";
 import { nodeParserFromSettingsOrContext } from "../Settings.js";
 import { runTransformations } from "../ingestion/IngestionPipeline.js";

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

@@ -5,7 +5,6 @@ import type {
   NodeWithScore,
 } from "@llamaindex/core/schema";
 import { MetadataMode } from "@llamaindex/core/schema";
-import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { serviceContextFromDefaults } from "../../ServiceContext.js";
 import { RetrieverQueryEngine } from "../../engines/query/index.js";
@@ -29,7 +28,11 @@ import {
   type KeywordExtractPrompt,
   type QueryKeywordExtractPrompt,
 } from "@llamaindex/core/prompts";
-import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
+import type {
+  BaseQueryEngine,
+  QueryBundle,
+} from "@llamaindex/core/query-engine";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import { extractText } from "@llamaindex/core/utils";
 import { llmFromSettingsOrContext } from "../../Settings.js";
 
@@ -48,7 +51,7 @@ export enum KeywordTableRetrieverMode {
 }
 
 // Base Keyword Table Retriever
-abstract class BaseKeywordTableRetriever implements BaseRetriever {
+abstract class BaseKeywordTableRetriever extends BaseRetriever {
   protected index: KeywordTableIndex;
   protected indexStruct: KeywordTable;
   protected docstore: BaseDocumentStore;
@@ -72,6 +75,7 @@ abstract class BaseKeywordTableRetriever implements BaseRetriever {
     maxKeywordsPerQuery: number;
     numChunksPerQuery: number;
   }) {
+    super();
     this.index = index;
     this.indexStruct = index.indexStruct;
     this.docstore = index.docStore;
@@ -87,7 +91,7 @@ abstract class BaseKeywordTableRetriever implements BaseRetriever {
 
   abstract getKeywords(query: string): Promise<string[]>;
 
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
     const keywords = await this.getKeywords(extractText(query));
     const chunkIndicesCount: { [key: string]: number } = {};
     const filteredKeywords = keywords.filter((keyword) =>

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

@@ -2,16 +2,17 @@ import {
   type ChoiceSelectPrompt,
   defaultChoiceSelectPrompt,
 } from "@llamaindex/core/prompts";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import { getResponseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import type {
   BaseNode,
   Document,
   NodeWithScore,
 } from "@llamaindex/core/schema";
-import { extractText, wrapEventCaller } from "@llamaindex/core/utils";
+import { extractText } from "@llamaindex/core/utils";
 import _ from "lodash";
-import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import {
   llmFromSettingsOrContext,
@@ -279,15 +280,15 @@ export type ListRetrieverMode = SummaryRetrieverMode;
 /**
  * Simple retriever for SummaryIndex that returns all nodes
  */
-export class SummaryIndexRetriever implements BaseRetriever {
+export class SummaryIndexRetriever extends BaseRetriever {
   index: SummaryIndex;
 
   constructor(index: SummaryIndex) {
+    super();
     this.index = index;
   }
 
-  @wrapEventCaller
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(queryBundle: QueryBundle): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const nodes = await this.index.docStore.getNodes(nodeIds);
     return nodes.map((node) => ({
@@ -300,7 +301,7 @@ export class SummaryIndexRetriever implements BaseRetriever {
 /**
  * LLM retriever for SummaryIndex which lets you select the most relevant chunks.
  */
-export class SummaryIndexLLMRetriever implements BaseRetriever {
+export class SummaryIndexLLMRetriever extends BaseRetriever {
   index: SummaryIndex;
   choiceSelectPrompt: ChoiceSelectPrompt;
   choiceBatchSize: number;
@@ -317,6 +318,7 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
     parseChoiceSelectAnswerFn?: ChoiceSelectParserFunction,
     serviceContext?: ServiceContext,
   ) {
+    super();
     this.index = index;
     this.choiceSelectPrompt = choiceSelectPrompt || defaultChoiceSelectPrompt;
     this.choiceBatchSize = choiceBatchSize;
@@ -326,7 +328,7 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
     this.serviceContext = serviceContext || index.serviceContext;
   }
 
-  async retrieve({ query }: RetrieveParams): Promise<NodeWithScore[]> {
+  async _retrieve(query: QueryBundle): Promise<NodeWithScore[]> {
     const nodeIds = this.index.indexStruct.nodes;
     const results: NodeWithScore[] = [];
 

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

@@ -2,9 +2,10 @@ import {
   DEFAULT_SIMILARITY_TOP_K,
   type BaseEmbedding,
 } from "@llamaindex/core/embeddings";
-import { Settings } from "@llamaindex/core/global";
 import type { MessageContent } from "@llamaindex/core/llms";
+import type { QueryBundle } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import {
   ImageNode,
   ModalityType,
@@ -14,8 +15,6 @@ import {
   type Document,
   type NodeWithScore,
 } from "@llamaindex/core/schema";
-import { wrapEventCaller } from "@llamaindex/core/utils";
-import type { BaseRetriever, RetrieveParams } from "../../Retriever.js";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { nodeParserFromSettingsOrContext } from "../../Settings.js";
 import { RetrieverQueryEngine } from "../../engines/query/RetrieverQueryEngine.js";
@@ -388,7 +387,7 @@ export type VectorIndexRetrieverOptions = {
   filters?: MetadataFilters;
 };
 
-export class VectorIndexRetriever implements BaseRetriever {
+export class VectorIndexRetriever extends BaseRetriever {
   index: VectorStoreIndex;
   topK: TopKMap;
 
@@ -401,6 +400,7 @@ export class VectorIndexRetriever implements BaseRetriever {
     topK,
     filters,
   }: VectorIndexRetrieverOptions) {
+    super();
     this.index = index;
     this.serviceContext = this.index.serviceContext;
     this.topK = topK ?? {
@@ -417,32 +417,17 @@ export class VectorIndexRetriever implements BaseRetriever {
     this.topK[ModalityType.TEXT] = similarityTopK;
   }
 
-  @wrapEventCaller
-  async retrieve({
-    query,
-    preFilters,
-  }: RetrieveParams): Promise<NodeWithScore[]> {
-    Settings.callbackManager.dispatchEvent("retrieve-start", {
-      query,
-    });
+  async _retrieve(params: QueryBundle): Promise<NodeWithScore[]> {
+    const { query } = params;
     const vectorStores = this.index.vectorStores;
     let nodesWithScores: NodeWithScore[] = [];
 
     for (const type in vectorStores) {
       const vectorStore: VectorStore = vectorStores[type as ModalityType]!;
       nodesWithScores = nodesWithScores.concat(
-        await this.retrieveQuery(
-          query,
-          type as ModalityType,
-          vectorStore,
-          preFilters as MetadataFilters,
-        ),
+        await this.retrieveQuery(query, type as ModalityType, vectorStore),
       );
     }
-    Settings.callbackManager.dispatchEvent("retrieve-end", {
-      query,
-      nodes: nodesWithScores,
-    });
     return nodesWithScores;
   }
 

packages/llamaindex/src/llm/types.ts

@@ -1,10 +0,0 @@
-import type { MessageContent } from "@llamaindex/core/llms";
-import type { NodeWithScore } from "@llamaindex/core/schema";
-
-export type RetrievalStartEvent = {
-  query: MessageContent;
-};
-export type RetrievalEndEvent = {
-  query: MessageContent;
-  nodes: NodeWithScore[];
-};

packages/llamaindex/src/objects/base.ts

@@ -1,8 +1,8 @@
 import type { BaseTool, MessageContent } from "@llamaindex/core/llms";
+import { BaseRetriever } from "@llamaindex/core/retriever";
 import type { BaseNode, Metadata } from "@llamaindex/core/schema";
 import { TextNode } from "@llamaindex/core/schema";
 import { extractText } from "@llamaindex/core/utils";
-import type { BaseRetriever } from "../Retriever.js";
 import type { VectorStoreIndex } from "../indices/vectorStore/index.js";
 
 // Assuming that necessary interfaces and classes (like OT, TextNode, BaseNode, etc.) are defined elsewhere
@@ -49,9 +49,6 @@ export abstract class BaseObjectNodeMapping {
 
 // You will need to implement specific subclasses of BaseObjectNodeMapping as per your project requirements.
 
-// todo: multimodal support
-type QueryType = MessageContent;
-
 export class ObjectRetriever<T = unknown> {
   _retriever: BaseRetriever;
   _objectNodeMapping: BaseObjectNodeMapping;
@@ -70,7 +67,7 @@ export class ObjectRetriever<T = unknown> {
   }
 
   // Translating the retrieve method
-  async retrieve(strOrQueryBundle: QueryType): Promise<T[]> {
+  async retrieve(strOrQueryBundle: MessageContent): Promise<T[]> {
     const nodes = await this.retriever.retrieve({
       query: extractText(strOrQueryBundle),
     });