feat: use query bundle

packages/core/e2e/fixtures/embeddings/OpenAIEmbedding.ts

@@ -0,0 +1,33 @@
+import { BaseNode, SimilarityType, type BaseEmbedding } from "llamaindex";
+
+export class OpenAIEmbedding implements BaseEmbedding {
+ embedBatchSize = 512;
+
+ async getQueryEmbedding(query: string) {
+ return [0];
+ }
+
+ async getTextEmbedding(text: string) {
+ return [0];
+ }
+
+ async getTextEmbeddings(texts: string[]) {
+ return [[0]];
+ }
+
+ async getTextEmbeddingsBatch(texts: string[]) {
+ return [[0]];
+ }
+
+ similarity(
+ embedding1: number[],
+ embedding2: number[],
+ mode?: SimilarityType,
+ ) {
+ return 1;
+ }
+
+ async transform(nodes: BaseNode[], _options?: any): Promise<BaseNode[]> {
+ return nodes;
+ }
+}

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

@@ -1,4 +1,3 @@
-import { faker } from "@faker-js/faker";
 import type {
  ChatResponse,
  ChatResponseChunk,
@@ -9,6 +8,9 @@ import type {
  LLMCompletionParamsNonStreaming,
  LLMCompletionParamsStreaming,
 } from "llamaindex/llm/types";
+import { extractText } from "llamaindex/llm/utils";
+import { strictEqual } from "node:assert";
+import { llmCompleteMockStorage } from "../../node/utils.js";
 
 export function getOpenAISession() {
  return {};
@@ -40,32 +42,62 @@ export class OpenAI implements LLM {
  | LLMChatParamsStreaming<Record<string, unknown>>
  | LLMChatParamsNonStreaming<Record<string, unknown>>,
  ): unknown {
- if (params.stream) {
- return {
- [Symbol.asyncIterator]: async function* () {
- yield {
- delta: faker.word.words(),
- } satisfies ChatResponseChunk;
- },
- };
+ if (llmCompleteMockStorage.llmEventStart.length > 0) {
+ const chatMessage =
+ llmCompleteMockStorage.llmEventStart.shift()!["messages"];
+ strictEqual(chatMessage.length, params.messages.length);
+ for (let i = 0; i < chatMessage.length; i++) {
+ strictEqual(chatMessage[i].role, params.messages[i].role);
+ strictEqual(chatMessage[i].content, params.messages[i].content);
+ }
+
+ if (llmCompleteMockStorage.llmEventEnd.length > 0) {
+ const response =
+ llmCompleteMockStorage.llmEventEnd.shift()!["response"];
+ if (params.stream) {
+ const content = response.message.content as string;
+ // maybe this is not the correct way to split the content, but it's good enough for now
+ const tokens = content.split("");
+ return {
+ [Symbol.asyncIterator]: async function* () {
+ const delta = tokens.shift();
+ if (delta) {
+ yield {
+ delta,
+ } as ChatResponseChunk;
+ }
+ },
+ };
+ } else {
+ return response;
+ }
+ }
  }
- return {
- get raw(): never {
- throw new Error("not implemented");
- },
- message: {
- content: faker.lorem.paragraph(),
- role: "assistant",
- },
- } satisfies ChatResponse;
+ throw new Error("Method not implemented.");
  }
  complete(
  params: LLMCompletionParamsStreaming,
  ): Promise<AsyncIterable<CompletionResponse>>;
  complete(
  params: LLMCompletionParamsNonStreaming,
  ): Promise<CompletionResponse>;
- async complete(params: unknown): Promise<unknown> {
+ async complete(
+ params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+ ): Promise<AsyncIterable<CompletionResponse> | CompletionResponse> {
+ if (llmCompleteMockStorage.llmEventStart.length > 0) {
+ const chatMessage =
+ llmCompleteMockStorage.llmEventStart.shift()!["messages"];
+ strictEqual(chatMessage.length, 1);
+ strictEqual(chatMessage[0].role, "user");
+ strictEqual(chatMessage[0].content, params.prompt);
+ }
+ if (llmCompleteMockStorage.llmEventEnd.length > 0) {
+ const response = llmCompleteMockStorage.llmEventEnd.shift()!["response"];
+ return {
+ raw: response,
+ text: extractText(response.message.content),
+ } satisfies CompletionResponse;
+ }
  throw new Error("Method not implemented.");
  }
 }

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

@@ -1,84 +1,30 @@
 /* eslint-disable @typescript-eslint/no-floating-promises */
 import { consola } from "consola";
 import {
+ Document,
  OpenAI,
  OpenAIAgent,
+ QueryEngineTool,
  Settings,
+ SubQuestionQueryEngine,
+ VectorStoreIndex,
  type LLM,
- type LLMEndEvent,
- type LLMStartEvent,
 } from "llamaindex";
 import { ok } from "node:assert";
-import type { WriteStream } from "node:fs";
-import { createWriteStream } from "node:fs";
-import { mkdir } from "node:fs/promises";
-import { join } from "node:path";
-import { after, before, beforeEach, describe, test } from "node:test";
-import { inspect } from "node:util";
+import { before, test } from "node:test";
+import { mockLLMEvent } from "./utils.js";
 
 let llm: LLM;
-let fsStream: WriteStream;
 before(async () => {
- const logUrl = new URL(
- join(
- "..",
- "logs",
- `basic.e2e.${new Date().toISOString().replace(/:/g, "-").replace(/\./g, "-")}.log`,
- ),
- import.meta.url,
- );
- await mkdir(new URL(".", logUrl), { recursive: true });
- fsStream = createWriteStream(logUrl, {
- encoding: "utf-8",
- });
-});
-
-after(() => {
- fsStream.end();
-});
-
-beforeEach((s) => {
- fsStream.write("start: " + s.name + "\n");
-});
-
-const llmEventStartHandler = (event: LLMStartEvent) => {
- const { payload } = event.detail;
- fsStream.write(
- "llmEventStart: " +
- inspect(payload, {
- depth: Infinity,
- }) +
- "\n",
- );
-};
-
-const llmEventEndHandler = (event: LLMEndEvent) => {
- const { payload } = event.detail;
- fsStream.write(
- "llmEventEnd: " +
- inspect(payload, {
- depth: Infinity,
- }) +
- "\n",
- );
-};
-
-before(() => {
  Settings.llm = new OpenAI({
  model: "gpt-3.5-turbo",
  });
  llm = Settings.llm;
- Settings.callbackManager.on("llm-start", llmEventStartHandler);
- Settings.callbackManager.on("llm-end", llmEventEndHandler);
-});
-
-after(() => {
- Settings.callbackManager.off("llm-start", llmEventStartHandler);
- Settings.callbackManager.off("llm-end", llmEventEndHandler);
 });
 
-describe("llm", () => {
- test("llm.chat", async () => {
+test("llm", async (t) => {
+ await mockLLMEvent(t, "llm");
+ await t.test("llm.chat", async () => {
  const response = await llm.chat({
  messages: [
  {
@@ -91,7 +37,7 @@ describe("llm", () => {
  ok(typeof response.message.content === "string");
  });
 
- test("stream llm.chat", async () => {
+ await t.test("stream llm.chat", async () => {
  const iter = await llm.chat({
  stream: true,
  messages: [
@@ -108,8 +54,9 @@ describe("llm", () => {
  });
 });
 
-describe("agent", () => {
- test("agent.chat", async () => {
+test("agent", async (t) => {
+ await mockLLMEvent(t, "agent");
+ await t.test("chat", async () => {
  const agent = new OpenAIAgent({
  tools: [
  {
@@ -137,3 +84,33 @@ describe("agent", () => {
  ok(typeof result.response === "string");
  });
 });
+
+test("queryEngine", async (t) => {
+ await mockLLMEvent(t, "queryEngine_subquestion");
+ await t.test("subquestion", async () => {
+ const document = new Document({
+ text: "Bill Gates stole from Apple.\n Steve Jobs stole from Xerox.",
+ });
+ const index = await VectorStoreIndex.fromDocuments([document]);
+
+ const queryEngineTools = [
+ new QueryEngineTool({
+ queryEngine: index.asQueryEngine(),
+ metadata: {
+ name: "bill_gates_idea",
+ description: "Get what Bill Gates idea from.",
+ },
+ }),
+ ];
+
+ const queryEngine = SubQuestionQueryEngine.fromDefaults({
+ queryEngineTools,
+ });
+
+ const { response } = await queryEngine.query({
+ query: "What did Bill Gates steal from?",
+ });
+
+ ok(response.includes("Apple"));
+ });
+});