fix: streaming for ReAct Agent (#798)

.changeset/yellow-pigs-reply.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: streaming for ReAct Agent

examples/agent/openai-task.ts

@@ -1,63 +1,15 @@
-import { ChatResponseChunk, FunctionTool, OpenAIAgent } from "llamaindex";
+import { ChatResponseChunk, OpenAIAgent } from "llamaindex";
 import { ReadableStream } from "node:stream/web";
-
-const functionTool = FunctionTool.from(
-  () => {
-    console.log("Getting user id...");
-    return crypto.randomUUID();
-  },
-  {
-    name: "get_user_id",
-    description: "Get a random user id",
-  },
-);
-
-const functionTool2 = FunctionTool.from(
-  ({ userId }: { userId: string }) => {
-    console.log("Getting user info...", userId);
-    return `Name: Alex; Address: 1234 Main St, CA; User ID: ${userId}`;
-  },
-  {
-    name: "get_user_info",
-    description: "Get user info",
-    parameters: {
-      type: "object",
-      properties: {
-        userId: {
-          type: "string",
-          description: "The user id",
-        },
-      },
-      required: ["userId"],
-    },
-  },
-);
-
-const functionTool3 = FunctionTool.from(
-  ({ address }: { address: string }) => {
-    console.log("Getting weather...", address);
-    return `${address} is in a sunny location!`;
-  },
-  {
-    name: "get_weather",
-    description: "Get the current weather for a location",
-    parameters: {
-      type: "object",
-      properties: {
-        address: {
-          type: "string",
-          description: "The address",
-        },
-      },
-      required: ["address"],
-    },
-  },
-);
+import {
+  getCurrentIDTool,
+  getUserInfoTool,
+  getWeatherTool,
+} from "./utils/tools";
 
 async function main() {
   // Create an OpenAIAgent with the function tools
   const agent = new OpenAIAgent({
-    tools: [functionTool, functionTool2, functionTool3],
+    tools: [getCurrentIDTool, getUserInfoTool, getWeatherTool],
   });
 
   const task = await agent.createTask(

examples/agent/react-task.ts

@@ -0,0 +1,40 @@
+import { ChatResponseChunk, ReActAgent } from "llamaindex";
+import { ReadableStream } from "node:stream/web";
+import {
+  getCurrentIDTool,
+  getUserInfoTool,
+  getWeatherTool,
+} from "./utils/tools";
+
+async function main() {
+  // Create an OpenAIAgent with the function tools
+  const agent = new ReActAgent({
+    tools: [getCurrentIDTool, getUserInfoTool, getWeatherTool],
+  });
+
+  const task = await agent.createTask(
+    "What is my current address weather based on my profile?",
+    true,
+  );
+
+  for await (const stepOutput of task) {
+    const stream = stepOutput.output as ReadableStream<ChatResponseChunk>;
+    if (stepOutput.isLast) {
+      for await (const chunk of stream) {
+        process.stdout.write(chunk.delta);
+      }
+      process.stdout.write("\n");
+    } else {
+      // handing function call
+      console.log("handling function call...");
+      for await (const chunk of stream) {
+        console.log("debug:", JSON.stringify(chunk.raw));
+      }
+    }
+    console.log("---");
+  }
+}
+
+void main().then(() => {
+  console.log("Done");
+});

examples/agent/utils/tools.ts

@@ -0,0 +1,54 @@
+import { FunctionTool } from "llamaindex";
+
+export const getCurrentIDTool = FunctionTool.from(
+  () => {
+    console.log("Getting user id...");
+    return crypto.randomUUID();
+  },
+  {
+    name: "get_user_id",
+    description: "Get a random user id",
+  },
+);
+
+export const getUserInfoTool = FunctionTool.from(
+  ({ userId }: { userId: string }) => {
+    console.log("Getting user info...", userId);
+    return `Name: Alex; Address: 1234 Main St, CA; User ID: ${userId}`;
+  },
+  {
+    name: "get_user_info",
+    description: "Get user info",
+    parameters: {
+      type: "object",
+      properties: {
+        userId: {
+          type: "string",
+          description: "The user id",
+        },
+      },
+      required: ["userId"],
+    },
+  },
+);
+
+export const getWeatherTool = FunctionTool.from(
+  ({ address }: { address: string }) => {
+    console.log("Getting weather...", address);
+    return `${address} is in a sunny location!`;
+  },
+  {
+    name: "get_weather",
+    description: "Get the current weather for a location",
+    parameters: {
+      type: "object",
+      properties: {
+        address: {
+          type: "string",
+          description: "The address",
+        },
+      },
+      required: ["address"],
+    },
+  },
+);

packages/core/e2e/node/react.e2e.ts

@@ -27,3 +27,23 @@ await test("react agent", async (t) => {
     ok(extractText(response.message.content).includes("72"));
   });
 });
+
+await test("react agent stream", async (t) => {
+  await mockLLMEvent(t, "react-agent-stream");
+  await t.test("get weather", async () => {
+    const agent = new ReActAgent({
+      tools: [getWeatherTool],
+    });
+
+    const stream = await agent.chat({
+      stream: true,
+      message: "What is the weather like in San Francisco?",
+    });
+
+    let content = "";
+    for await (const { response } of stream) {
+      content += response.delta;
+    }
+    ok(content.includes("72"));
+  });
+});