refactor: import readable stream

Author: thucpn

templates/components/multiagent/typescript/express/chat.controller.ts

@@ -5,10 +5,7 @@ import {
   retrieveMessageContent,
 } from "./llamaindex/streaming/annotations";
 import { createWorkflow } from "./workflow/factory";
-import {
-  createStreamFromWorkflowContext,
-  streamToAsyncIterable,
-} from "./workflow/stream";
+import { createStreamFromWorkflowContext } from "./workflow/stream";
 
 export const chat = async (req: Request, res: Response) => {
   try {
@@ -30,13 +27,10 @@ export const chat = async (req: Request, res: Response) => {
 
     const { stream, dataStream } =
       await createStreamFromWorkflowContext(context);
-    const streamIterable = streamToAsyncIterable(stream);
-    const streamResponse = LlamaIndexAdapter.toDataStreamResponse(
-      streamIterable,
-      {
-        data: dataStream,
-      },
-    );
+
+    const streamResponse = LlamaIndexAdapter.toDataStreamResponse(stream, {
+      data: dataStream,
+    });
     if (streamResponse.body) {
       const reader = streamResponse.body.getReader();
       while (true) {

templates/components/multiagent/typescript/nextjs/route.ts

@@ -8,10 +8,7 @@ import {
   retrieveMessageContent,
 } from "./llamaindex/streaming/annotations";
 import { createWorkflow } from "./workflow/factory";
-import {
-  createStreamFromWorkflowContext,
-  streamToAsyncIterable,
-} from "./workflow/stream";
+import { createStreamFromWorkflowContext } from "./workflow/stream";
 
 initObservability();
 initSettings();
@@ -44,8 +41,7 @@ export async function POST(request: NextRequest) {
     });
     const { stream, dataStream } =
       await createStreamFromWorkflowContext(context);
-    const streamIterable = streamToAsyncIterable(stream);
-    return LlamaIndexAdapter.toDataStreamResponse(streamIterable, {
+    return LlamaIndexAdapter.toDataStreamResponse(stream, {
       data: dataStream,
     });
   } catch (error) {

templates/components/multiagent/typescript/workflow/stream.ts

@@ -3,30 +3,26 @@ import {
   WorkflowContext,
   WorkflowEvent,
 } from "@llamaindex/workflow";
-import {
-  StreamData,
-  createStreamDataTransformer,
-  parseDataStreamPart,
-} from "ai";
+import { StreamData } from "ai";
 import { ChatResponseChunk, EngineResponse } from "llamaindex";
+import { ReadableStream } from "stream/web";
 import { AgentRunEvent } from "./type";
 
 export async function createStreamFromWorkflowContext<Input, Output, Context>(
   context: WorkflowContext<Input, Output, Context>,
-): Promise<{ stream: ReadableStream<string>; dataStream: StreamData }> {
+): Promise<{ stream: ReadableStream<EngineResponse>; dataStream: StreamData }> {
   const dataStream = new StreamData();
-  const encoder = new TextEncoder();
   let generator: AsyncGenerator<ChatResponseChunk> | undefined;
 
   const closeStreams = (controller: ReadableStreamDefaultController) => {
     controller.close();
     dataStream.close();
   };
 
-  const mainStream = new ReadableStream({
+  const stream = new ReadableStream<EngineResponse>({
     async start(controller) {
       // Kickstart the stream by sending an empty string
-      controller.enqueue(encoder.encode(""));
+      controller.enqueue({ delta: "" } as EngineResponse);
     },
     async pull(controller) {
       while (!generator) {
@@ -45,18 +41,14 @@ export async function createStreamFromWorkflowContext<Input, Output, Context>(
         closeStreams(controller);
         return;
       }
-      if (chunk.delta) {
-        controller.enqueue(encoder.encode(chunk.delta));
+      const delta = chunk.delta ?? "";
+      if (delta) {
+        controller.enqueue({ delta } as EngineResponse);
       }
     },
   });
 
-  return {
-    stream: mainStream
-      .pipeThrough(createStreamDataTransformer())
-      .pipeThrough(new TextDecoderStream()),
-    dataStream,
-  };
+  return { stream, dataStream };
 }
 
 function handleEvent(
@@ -75,25 +67,3 @@ function handleEvent(
     });
   }
 }
-
-export function streamToAsyncIterable(stream: ReadableStream<string>) {
-  const streamIterable: AsyncIterable<EngineResponse> = {
-    [Symbol.asyncIterator]() {
-      const reader = stream.getReader();
-      return {
-        async next() {
-          const { done, value } = await reader.read();
-          if (done) {
-            return { done: true, value: undefined };
-          }
-          const delta = parseDataStreamPart(value)?.value.toString() || "";
-          return {
-            done: false,
-            value: { delta } as unknown as EngineResponse,
-          };
-        },
-      };
-    },
-  };
-  return streamIterable;
-}