feat: use LlamaIndexAdapter

Author: marcusschiesser

templates/components/llamaindex/typescript/streaming/stream.ts

@@ -1,4 +1,4 @@
-import { JSONValue, Message, StreamData, streamToResponse } from "ai";
+import { JSONValue, LlamaIndexAdapter, Message, StreamData } from "ai";
 import { Request, Response } from "express";
 import { ChatMessage, Settings } from "llamaindex";
 import { createChatEngine } from "./engine/chat";
@@ -10,7 +10,7 @@ import {
   createCallbackManager,
   createStreamTimeout,
 } from "./llamaindex/streaming/events";
-import { LlamaIndexStream } from "./llamaindex/streaming/stream";
+import { generateNextQuestions } from "./llamaindex/streaming/suggestion";
 
 export const chat = async (req: Request, res: Response) => {
   // Init Vercel AI StreamData and timeout
@@ -56,24 +56,37 @@ export const chat = async (req: Request, res: Response) => {
 
     // Setup callbacks
     const callbackManager = createCallbackManager(vercelStreamData);
+    const chatHistory: ChatMessage[] = messages as ChatMessage[];
 
     // Calling LlamaIndex's ChatEngine to get a streamed response
     const response = await Settings.withCallbackManager(callbackManager, () => {
       return chatEngine.chat({
         message: userMessageContent,
-        chatHistory: messages as ChatMessage[],
+        chatHistory,
         stream: true,
       });
     });
 
-    // Return a stream, which can be consumed by the Vercel/AI client
-    const stream = LlamaIndexStream(
-      response,
-      vercelStreamData,
-      messages as ChatMessage[],
-    );
+    const onFinal = (content: string) => {
+      chatHistory.push({ role: "assistant", content: content });
+      generateNextQuestions(chatHistory)
+        .then((questions: string[]) => {
+          if (questions.length > 0) {
+            vercelStreamData.appendMessageAnnotation({
+              type: "suggested_questions",
+              data: questions,
+            });
+          }
+        })
+        .finally(() => {
+          vercelStreamData.close();
+        });
+    };
 
-    return streamToResponse(stream, res, {}, vercelStreamData);
+    return LlamaIndexAdapter.toDataStreamResponse(response, {
+      data: vercelStreamData,
+      callbacks: { onFinal },
+    });
   } catch (error) {
     console.error("[LlamaIndex]", error);
     return res.status(500).json({

templates/types/streaming/express/src/controllers/chat.controller.ts

@@ -1,5 +1,5 @@
 import { initObservability } from "@/app/observability";
-import { JSONValue, Message, StreamData, StreamingTextResponse } from "ai";
+import { JSONValue, LlamaIndexAdapter, Message, StreamData } from "ai";
 import { ChatMessage, Settings } from "llamaindex";
 import { NextRequest, NextResponse } from "next/server";
 import { createChatEngine } from "./engine/chat";
@@ -12,7 +12,7 @@ import {
   createCallbackManager,
   createStreamTimeout,
 } from "./llamaindex/streaming/events";
-import { LlamaIndexStream } from "./llamaindex/streaming/stream";
+import { generateNextQuestions } from "./llamaindex/streaming/suggestion";
 
 initObservability();
 initSettings();
@@ -69,25 +69,37 @@ export async function POST(request: NextRequest) {
 
     // Setup callbacks
     const callbackManager = createCallbackManager(vercelStreamData);
+    const chatHistory: ChatMessage[] = messages as ChatMessage[];
 
     // Calling LlamaIndex's ChatEngine to get a streamed response
     const response = await Settings.withCallbackManager(callbackManager, () => {
       return chatEngine.chat({
         message: userMessageContent,
-        chatHistory: messages as ChatMessage[],
+        chatHistory,
         stream: true,
       });
     });
 
-    // Transform LlamaIndex stream to Vercel/AI format
-    const stream = LlamaIndexStream(
-      response,
-      vercelStreamData,
-      messages as ChatMessage[],
-    );
+    const onFinal = (content: string) => {
+      chatHistory.push({ role: "assistant", content: content });
+      generateNextQuestions(chatHistory)
+        .then((questions: string[]) => {
+          if (questions.length > 0) {
+            vercelStreamData.appendMessageAnnotation({
+              type: "suggested_questions",
+              data: questions,
+            });
+          }
+        })
+        .finally(() => {
+          vercelStreamData.close();
+        });
+    };
 
-    // Return a StreamingTextResponse, which can be consumed by the Vercel/AI client
-    return new StreamingTextResponse(stream, {}, vercelStreamData);
+    return LlamaIndexAdapter.toDataStreamResponse(response, {
+      data: vercelStreamData,
+      callbacks: { onFinal },
+    });
   } catch (error) {
     console.error("[LlamaIndex]", error);
     return NextResponse.json(