feat: update route handler of NextJs streaming template to response stream

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

@@ -12,6 +12,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";
@@ -210,7 +211,7 @@ export class OpenAIAgentWorker implements AgentWorker {
     task: Task,
     mode: ChatResponseMode,
     llmChatKwargs: any,
-  ): Promise<AgentChatResponse> {
+  ): Promise<AgentChatResponse | AsyncIterable<AgentChatResponse>> {
     if (mode === ChatResponseMode.WAIT) {
       const chatResponse = (await this.llm.chat({
         stream: false,
@@ -219,7 +220,25 @@ export class OpenAIAgentWorker implements AgentWorker {
 
       return this._processMessage(task, chatResponse) as AgentChatResponse;
     } else {
-      throw new Error("Not implemented");
+      const stream = await this.llm.chat({
+        stream: true,
+        ...llmChatKwargs,
+      });
+      return streamConverter(
+        streamReducer({
+          stream,
+          initialValue: "",
+          reducer: (accumulator, part) => (accumulator += part.delta),
+          finished: (accumulator) => {
+            task.extraState.newMemory.put({
+              content: accumulator,
+              role: "assistant",
+            });
+          },
+        }),
+        (r: ChatResponseChunk) =>
+          new AgentChatResponse(r.delta, task.extraState.sources),
+      );
     }
   }
 

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

@@ -244,10 +244,18 @@ export class AgentRunner extends BaseAgentRunner {
       );
     }
 
+    // if not an AgentChatResponse or an AsyncIterable<AgentChatResponse> throw an error
     if (!(stepOutput.output instanceof AgentChatResponse)) {
-      throw new Error(
-        `When \`isLast\` is True, cur_step_output.output must be AGENT_CHAT_RESPONSE_TYPE: ${stepOutput.output}`,
-      );
+      const { value, done } = await (
+        stepOutput.output as AsyncIterable<AgentChatResponse>
+      )
+        [Symbol.asyncIterator]()
+        .next();
+      if (!(value 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 +270,16 @@ export class AgentRunner extends BaseAgentRunner {
   protected async _chat({
     message,
     toolChoice,
+    mode,
   }: ChatEngineAgentParams & { mode: ChatResponseMode }) {
     const task = this.createTask(message as string);
 
     let resultOutput;
 
     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 +303,10 @@ export class AgentRunner extends BaseAgentRunner {
     message,
     chatHistory,
     toolChoice,
-  }: ChatEngineAgentParams): Promise<AgentChatResponse> {
+    mode = ChatResponseMode.WAIT,
+  }: ChatEngineAgentParams & {
+    mode?: ChatResponseMode;
+  }): Promise<AgentChatResponse> {
     if (!toolChoice) {
       toolChoice = this.defaultToolChoice;
     }
@@ -308,7 +315,7 @@ export class AgentRunner extends BaseAgentRunner {
       message,
       chatHistory,
       toolChoice,
-      mode: ChatResponseMode.WAIT,
+      mode,
     });
 
     return chatResponse;

packages/create-llama/templates/components/vectordbs/typescript/none/index.ts

@@ -1,10 +1,12 @@
 import {
-  ContextChatEngine,
   LLM,
-  serviceContextFromDefaults,
+  OpenAI,
+  OpenAIAgent,
+  QueryEngineTool,
   SimpleDocumentStore,
-  storageContextFromDefaults,
   VectorStoreIndex,
+  serviceContextFromDefaults,
+  storageContextFromDefaults,
 } from "llamaindex";
 import { CHUNK_OVERLAP, CHUNK_SIZE, STORAGE_CACHE_DIR } from "./constants.mjs";
 
@@ -32,13 +34,22 @@ async function getDataSource(llm: LLM) {
   });
 }
 
-export async function createChatEngine(llm: LLM) {
+export async function createChatEngine(llm: OpenAI) {
   const index = await getDataSource(llm);
-  const retriever = index.asRetriever();
-  retriever.similarityTopK = 3;
-
-  return new ContextChatEngine({
-    chatModel: llm,
-    retriever,
+  const queryEngine = index.asQueryEngine();
+  const queryEngineTool = new QueryEngineTool({
+    queryEngine: queryEngine,
+    metadata: {
+      name: "data_query_engine",
+      description: `A query engine for documents in storage folder: ${STORAGE_CACHE_DIR}`,
+    },
   });
+
+  const agent = new OpenAIAgent({
+    tools: [queryEngineTool],
+    verbose: true,
+    llm,
+  });
+
+  return agent;
 }

packages/create-llama/templates/types/streaming/nextjs/app/api/chat/route.ts

@@ -1,5 +1,10 @@
 import { StreamingTextResponse } from "ai";
-import { ChatMessage, MessageContent, OpenAI } from "llamaindex";
+import {
+  ChatMessage,
+  ChatResponseMode,
+  MessageContent,
+  OpenAI,
+} from "llamaindex";
 import { NextRequest, NextResponse } from "next/server";
 import { createChatEngine } from "./engine";
 import { LlamaIndexStream } from "./llamaindex-stream";
@@ -54,11 +59,11 @@ export async function POST(request: NextRequest) {
       data?.imageUrl,
     );
 
-    // Calling LlamaIndex's ChatEngine to get a streamed response
+    // Calling chatEngine to get a streamed response
     const response = await chatEngine.chat({
       message: userMessageContent,
       chatHistory: messages,
-      stream: true,
+      mode: ChatResponseMode.STREAM,
     });
 
     // Transform LlamaIndex stream to Vercel/AI format