feat: OpenAI Agent Stream (#597)

.changeset/khaki-rivers-unite.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+Add streaming to agents

examples/agent/stream_openai_agent.ts

@@ -0,0 +1,77 @@
+import { FunctionTool, OpenAIAgent } from "llamaindex";
+
+// Define a function to sum two numbers
+function sumNumbers({ a, b }: { a: number; b: number }): number {
+  return a + b;
+}
+
+// Define a function to divide two numbers
+function divideNumbers({ a, b }: { a: number; b: number }): number {
+  return a / b;
+}
+
+// Define the parameters of the sum function as a JSON schema
+const sumJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The first number",
+    },
+    b: {
+      type: "number",
+      description: "The second number",
+    },
+  },
+  required: ["a", "b"],
+};
+
+const divideJSON = {
+  type: "object",
+  properties: {
+    a: {
+      type: "number",
+      description: "The dividend",
+    },
+    b: {
+      type: "number",
+      description: "The divisor",
+    },
+  },
+  required: ["a", "b"],
+};
+
+async function main() {
+  // Create a function tool from the sum function
+  const functionTool = new FunctionTool(sumNumbers, {
+    name: "sumNumbers",
+    description: "Use this function to sum two numbers",
+    parameters: sumJSON,
+  });
+
+  // Create a function tool from the divide function
+  const functionTool2 = new FunctionTool(divideNumbers, {
+    name: "divideNumbers",
+    description: "Use this function to divide two numbers",
+    parameters: divideJSON,
+  });
+
+  // Create an OpenAIAgent with the function tools
+  const agent = new OpenAIAgent({
+    tools: [functionTool, functionTool2],
+    verbose: false,
+  });
+
+  const stream = await agent.chat({
+    message: "Divide 16 by 2 then add 20",
+    stream: true,
+  });
+
+  for await (const chunk of stream.response) {
+    process.stdout.write(chunk.response);
+  }
+}
+
+main().then(() => {
+  console.log("\nDone");
+});

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

@@ -1,10 +1,12 @@
 // Assuming that the necessary interfaces and classes (like BaseTool, OpenAI, ChatMessage, CallbackManager, etc.) are defined elsewhere
 
 import { randomUUID } from "@llamaindex/env";
+import { Response } from "../../Response.js";
 import type { CallbackManager } from "../../callbacks/CallbackManager.js";
 import {
   AgentChatResponse,
   ChatResponseMode,
+  StreamingAgentChatResponse,
 } from "../../engines/chat/types.js";
 import type {
   ChatMessage,
@@ -12,6 +14,7 @@ import type {
   ChatResponseChunk,
 } from "../../llm/index.js";
 import { OpenAI } from "../../llm/index.js";
+import { streamConverter, streamReducer } from "../../llm/utils.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
 import type { ObjectRetriever } from "../../objects/base.js";
 import type { ToolOutput } from "../../tools/types.js";
@@ -192,13 +195,40 @@ export class OpenAIAgentWorker implements AgentWorker {
   private _processMessage(
     task: Task,
     chatResponse: ChatResponse,
-  ): AgentChatResponse | AsyncIterable<ChatResponseChunk> {
+  ): AgentChatResponse {
     const aiMessage = chatResponse.message;
     task.extraState.newMemory.put(aiMessage);
 
     return new AgentChatResponse(aiMessage.content, task.extraState.sources);
   }
 
+  private async _getStreamAiResponse(
+    task: Task,
+    llmChatKwargs: any,
+  ): Promise<StreamingAgentChatResponse> {
+    const stream = await this.llm.chat({
+      stream: true,
+      ...llmChatKwargs,
+    });
+
+    const iterator = streamConverter(
+      streamReducer({
+        stream,
+        initialValue: "",
+        reducer: (accumulator, part) => (accumulator += part.delta),
+        finished: (accumulator) => {
+          task.extraState.newMemory.put({
+            content: accumulator,
+            role: "assistant",
+          });
+        },
+      }),
+      (r: ChatResponseChunk) => new Response(r.delta),
+    );
+
+    return new StreamingAgentChatResponse(iterator, task.extraState.sources);
+  }
+
   /**
    * Get agent response.
    * @param task: task
@@ -210,7 +240,7 @@ export class OpenAIAgentWorker implements AgentWorker {
     task: Task,
     mode: ChatResponseMode,
     llmChatKwargs: any,
-  ): Promise<AgentChatResponse> {
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse> {
     if (mode === ChatResponseMode.WAIT) {
       const chatResponse = (await this.llm.chat({
         stream: false,
@@ -218,9 +248,11 @@ export class OpenAIAgentWorker implements AgentWorker {
       })) as unknown as ChatResponse;
 
       return this._processMessage(task, chatResponse) as AgentChatResponse;
-    } else {
-      throw new Error("Not implemented");
+    } else if (mode === ChatResponseMode.STREAM) {
+      return this._getStreamAiResponse(task, llmChatKwargs);
     }
+
+    throw new Error("Invalid mode");
   }
 
   /**

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

@@ -4,6 +4,7 @@ import type { ChatEngineAgentParams } from "../../engines/chat/index.js";
 import {
   AgentChatResponse,
   ChatResponseMode,
+  StreamingAgentChatResponse,
 } from "../../engines/chat/index.js";
 import type { ChatMessage, LLM } from "../../llm/index.js";
 import { ChatMemoryBuffer } from "../../memory/ChatMemoryBuffer.js";
@@ -231,23 +232,26 @@ export class AgentRunner extends BaseAgentRunner {
     taskId: string,
     stepOutput: TaskStepOutput,
     kwargs?: any,
-  ): Promise<AgentChatResponse> {
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse> {
     if (!stepOutput) {
       stepOutput =
         this.getCompletedSteps(taskId)[
           this.getCompletedSteps(taskId).length - 1
         ];
     }
+
     if (!stepOutput.isLast) {
       throw new Error(
         "finalizeResponse can only be called on the last step output",
       );
     }
 
-    if (!(stepOutput.output instanceof AgentChatResponse)) {
-      throw new Error(
-        `When \`isLast\` is True, cur_step_output.output must be AGENT_CHAT_RESPONSE_TYPE: ${stepOutput.output}`,
-      );
+    if (!(stepOutput.output instanceof StreamingAgentChatResponse)) {
+      if (!(stepOutput.output instanceof AgentChatResponse)) {
+        throw new Error(
+          `When \`isLast\` is True, cur_step_output.output must be AGENT_CHAT_RESPONSE_TYPE: ${stepOutput.output}`,
+        );
+      }
     }
 
     this.agentWorker.finalizeTask(this.getTask(taskId), kwargs);
@@ -262,20 +266,32 @@ export class AgentRunner extends BaseAgentRunner {
   protected async _chat({
     message,
     toolChoice,
-  }: ChatEngineAgentParams & { mode: ChatResponseMode }) {
+    stream,
+  }: ChatEngineAgentParams): Promise<AgentChatResponse>;
+  protected async _chat({
+    message,
+    toolChoice,
+    stream,
+  }: ChatEngineAgentParams & {
+    stream: true;
+  }): Promise<StreamingAgentChatResponse>;
+  protected async _chat({
+    message,
+    toolChoice,
+    stream,
+  }: ChatEngineAgentParams): Promise<
+    AgentChatResponse | StreamingAgentChatResponse
+  > {
     const task = this.createTask(message as string);
 
     let resultOutput;
 
+    const mode = stream ? ChatResponseMode.STREAM : ChatResponseMode.WAIT;
+
     while (true) {
-      const curStepOutput = await this._runStep(
-        task.taskId,
-        undefined,
-        ChatResponseMode.WAIT,
-        {
-          toolChoice,
-        },
-      );
+      const curStepOutput = await this._runStep(task.taskId, undefined, mode, {
+        toolChoice,
+      });
 
       if (curStepOutput.isLast) {
         resultOutput = curStepOutput;
@@ -299,7 +315,26 @@ export class AgentRunner extends BaseAgentRunner {
     message,
     chatHistory,
     toolChoice,
-  }: ChatEngineAgentParams): Promise<AgentChatResponse> {
+    stream,
+  }: ChatEngineAgentParams & {
+    stream?: false;
+  }): Promise<AgentChatResponse>;
+  public async chat({
+    message,
+    chatHistory,
+    toolChoice,
+    stream,
+  }: ChatEngineAgentParams & {
+    stream: true;
+  }): Promise<StreamingAgentChatResponse>;
+  public async chat({
+    message,
+    chatHistory,
+    toolChoice,
+    stream,
+  }: ChatEngineAgentParams): Promise<
+    AgentChatResponse | StreamingAgentChatResponse
+  > {
     if (!toolChoice) {
       toolChoice = this.defaultToolChoice;
     }
@@ -308,7 +343,7 @@ export class AgentRunner extends BaseAgentRunner {
       message,
       chatHistory,
       toolChoice,
-      mode: ChatResponseMode.WAIT,
+      stream,
     });
 
     return chatResponse;

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

@@ -1,4 +1,7 @@
-import type { AgentChatResponse } from "../../engines/chat/index.js";
+import type {
+  AgentChatResponse,
+  StreamingAgentChatResponse,
+} from "../../engines/chat/index.js";
 import type { Task, TaskStep, TaskStepOutput } from "../types.js";
 import { BaseAgent } from "../types.js";
 
@@ -57,7 +60,7 @@ export abstract class BaseAgentRunner extends BaseAgent {
     taskId: string,
     stepOutput: TaskStepOutput,
     kwargs?: any,
-  ): Promise<AgentChatResponse>;
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse>;
 
   abstract undoStep(taskId: string): void;
 }

packages/core/src/agent/types.ts

@@ -1,6 +1,7 @@
 import type {
   AgentChatResponse,
   ChatEngineAgentParams,
+  StreamingAgentChatResponse,
 } from "../engines/chat/index.js";
 import type { QueryEngineParamsNonStreaming } from "../types.js";
 
@@ -12,11 +13,15 @@ export interface AgentWorker {
 }
 
 interface BaseChatEngine {
-  chat(params: ChatEngineAgentParams): Promise<AgentChatResponse>;
+  chat(
+    params: ChatEngineAgentParams,
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse>;
 }
 
 interface BaseQueryEngine {
-  query(params: QueryEngineParamsNonStreaming): Promise<AgentChatResponse>;
+  query(
+    params: QueryEngineParamsNonStreaming,
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse>;
 }
 
 /**
@@ -31,7 +36,10 @@ export abstract class BaseAgent implements BaseChatEngine, BaseQueryEngine {
     return [];
   }
 
-  abstract chat(params: ChatEngineAgentParams): Promise<AgentChatResponse>;
+  abstract chat(
+    params: ChatEngineAgentParams,
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse>;
+
   abstract reset(): void;
 
   /**
@@ -41,7 +49,7 @@ export abstract class BaseAgent implements BaseChatEngine, BaseQueryEngine {
    */
   async query(
     params: QueryEngineParamsNonStreaming,
-  ): Promise<AgentChatResponse> {
+  ): Promise<AgentChatResponse | StreamingAgentChatResponse> {
     // Handle non-streaming query
     const agentResponse = await this.chat({
       message: params.query,

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

@@ -27,6 +27,7 @@ export interface ChatEngineParamsNonStreaming extends ChatEngineParamsBase {
 
 export interface ChatEngineAgentParams extends ChatEngineParamsBase {
   toolChoice?: string | Record<string, any>;
+  stream?: boolean;
 }
 
 /**
@@ -86,3 +87,20 @@ export class AgentChatResponse {
     return this.response ?? "";
   }
 }
+
+export class StreamingAgentChatResponse {
+  response: AsyncIterable<Response>;
+
+  sources: ToolOutput[];
+  sourceNodes?: BaseNode[];
+
+  constructor(
+    response: AsyncIterable<Response>,
+    sources?: ToolOutput[],
+    sourceNodes?: BaseNode[],
+  ) {
+    this.response = response;
+    this.sources = sources ?? [];
+    this.sourceNodes = sourceNodes ?? [];
+  }
+}