fix: streaming issues with LLMAgent (#1692)

.changeset/afraid-boxes-approve.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+fix: streaming issues with LLMAgent

examples/chat-engine/streaming.ts

@@ -0,0 +1,25 @@
+import { LLMAgent } from "llamaindex";
+import { stdin as input, stdout as output } from "node:process";
+import readline from "node:readline/promises";
+
+const agent = new LLMAgent({ tools: [] });
+
+(async () => {
+  const rl = readline.createInterface({ input, output });
+
+  while (true) {
+    const query = await rl.question("User: ");
+
+    const startTime = Date.now();
+    const stream = await agent.chat({ message: query, stream: true });
+    const timeToGetFirstChunk = Date.now() - startTime;
+    process.stdout.write(
+      `Time to get first chunk from LLMAgent: ${timeToGetFirstChunk}ms\n`,
+    );
+    process.stdout.write("Assistant with LLMAgent: ");
+    for await (const chunk of stream) {
+      process.stdout.write(chunk.response);
+    }
+    process.stdout.write("\n");
+  }
+})();

packages/core/src/agent/base.ts

@@ -193,7 +193,7 @@ export abstract class AgentWorker<
     return new ReadableStream<
       TaskStepOutput<AI, Store, AdditionalMessageOptions, AdditionalChatOptions>
     >({
-      start: async (controller) => {
+      pull: async (controller) => {
         for await (const stepOutput of taskOutputStream) {
           this.#taskSet.add(stepOutput.taskStep);
           if (stepOutput.isLast) {
@@ -209,23 +209,31 @@ export abstract class AgentWorker<
             }
             const { output, taskStep } = stepOutput;
             if (output instanceof ReadableStream) {
-              const [pipStream, finalStream] = output.tee();
-              stepOutput.output = finalStream;
-              const reader = pipStream.getReader();
-              const { value } = await reader.read();
-              reader.releaseLock();
-              let content: string = value!.delta;
-              for await (const chunk of pipStream) {
-                content += chunk.delta;
-              }
-              taskStep.context.store.messages = [
-                ...taskStep.context.store.messages,
-                {
-                  role: "assistant",
-                  content,
-                  options: value!.options,
-                },
-              ];
+              let content = "";
+              let options: AdditionalMessageOptions | undefined = undefined;
+              const transformedStream = output.pipeThrough(
+                new TransformStream({
+                  transform(chunk, controller) {
+                    content += chunk.delta;
+                    if (!options && chunk.options) {
+                      options = chunk.options;
+                    }
+                    controller.enqueue(chunk); // Pass the chunk through unchanged
+                  },
+                  // When stream finishes, store the accumulated message in context
+                  flush() {
+                    taskStep.context.store.messages = [
+                      ...taskStep.context.store.messages,
+                      {
+                        role: "assistant",
+                        content,
+                        options,
+                      },
+                    ];
+                  },
+                }),
+              );
+              stepOutput.output = transformedStream;
             }
             controller.enqueue(stepOutput);
             controller.close();

packages/core/tests/agents.test.ts

@@ -1,4 +1,5 @@
-import { validateAgentParams } from "@llamaindex/core/agent";
+import { LLMAgent, validateAgentParams } from "@llamaindex/core/agent";
+import { MockLLM } from "@llamaindex/core/utils";
 import { expect, test } from "vitest";
 import { ZodError } from "zod";
 
@@ -33,3 +34,62 @@ test("validate agent params", () => {
     ],
   });
 });
+
+test("LLMAgent streaming: first chunk should be available immediately", async () => {
+  const responseMessage =
+    "This is a very long response message that should take a while to stream";
+  const timeBetweenToken = 20; // delay time between tokens
+
+  const agent = new LLMAgent({
+    tools: [],
+    llm: new MockLLM({ responseMessage, timeBetweenToken }),
+  });
+
+  const startTime = Date.now();
+  const stream = await agent.chat({ message: "Hello", stream: true });
+
+  let fullResponse = "";
+  let timeToGetFirstChunk: number | undefined;
+
+  for await (const chunk of stream) {
+    expect(chunk).toHaveProperty("delta");
+    fullResponse += chunk.delta;
+    if (timeToGetFirstChunk === undefined) {
+      timeToGetFirstChunk = Date.now() - startTime;
+    }
+  }
+
+  expect(fullResponse).toBe(responseMessage);
+
+  // the first chunk should be available immediately and no need the whole response to be sent
+  expect(timeToGetFirstChunk).toBeLessThan(500);
+});
+
+test("LLMAgent create task: first task should be executed immediately", async () => {
+  const responseMessage =
+    "This is a very long response message that should take a while to stream";
+  const timeBetweenToken = 20; // delay time between tokens
+
+  const agent = new LLMAgent({
+    tools: [],
+    llm: new MockLLM({ responseMessage, timeBetweenToken }),
+  });
+
+  const startTime = Date.now();
+  const task = agent.createTask("Write a long paragraph", true, false, []);
+
+  let timeToGetFirstChunk: number | undefined;
+  let output: ReadableStream | undefined;
+  for await (const stepOutput of task) {
+    if (timeToGetFirstChunk === undefined) {
+      timeToGetFirstChunk = Date.now() - startTime;
+    }
+    if (stepOutput.output instanceof ReadableStream) {
+      output = stepOutput.output;
+    }
+  }
+
+  expect(timeToGetFirstChunk).toBeLessThan(500);
+  expect(output).toBeDefined();
+  expect(output).toBeInstanceOf(ReadableStream);
+});