add timeout for streaming

Author: thucpn

helpers/env-variables.ts

@@ -276,6 +276,12 @@ const getEngineEnvs = (): EnvVar[] => {
         "The number of similar embeddings to return when retrieving documents.",
       value: "3",
     },
+    {
+      name: "STREAM_TIMEOUT",
+      description:
+        "The time in milliseconds to wait for the stream to return a response.",
+      value: "60000",
+    },
   ];
 };
 

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

@@ -7,9 +7,12 @@ import {
   LlamaIndexStream,
   convertMessageContent,
 } from "./llamaindex-stream";
-import { createCallbackManager } from "./stream-helper";
+import { createCallbackManager, createStreamTimeout } from "./stream-helper";
 
 export const chat = async (req: Request, res: Response) => {
+  // Init Vercel AI StreamData and timeout
+  const vercelStreamData = new StreamData();
+  const streamTimeout = createStreamTimeout(vercelStreamData);
   try {
     const {
       messages,
@@ -28,9 +31,6 @@ export const chat = async (req: Request, res: Response) => {
     // Convert message content from Vercel/AI format to LlamaIndex/OpenAI format
     const userMessageContent = convertMessageContent(userMessage.content, data);
 
-    // Init Vercel AI StreamData
-    const vercelStreamData = new StreamData();
-
     // Setup callbacks
     const callbackManager = createCallbackManager(vercelStreamData);
 
@@ -57,5 +57,7 @@ export const chat = async (req: Request, res: Response) => {
     return res.status(500).json({
       detail: (error as Error).message,
     });
+  } finally {
+    clearTimeout(streamTimeout);
   }
 };

templates/types/streaming/express/src/controllers/stream-helper.ts

@@ -75,6 +75,15 @@ export function appendToolData(
   });
 }
 
+export function createStreamTimeout(stream: StreamData) {
+  const timeout = Number(process.env.STREAM_TIMEOUT ?? 1000 * 60 * 5); // default to 5 minutes
+  const t = setTimeout(() => {
+    appendEventData(stream, `Stream timed out after ${timeout / 1000} seconds`);
+    stream.close();
+  }, timeout);
+  return t;
+}
+
 export function createCallbackManager(stream: StreamData) {
   const callbackManager = new CallbackManager();
 

templates/types/streaming/nextjs/app/api/chat/route.ts

@@ -9,7 +9,7 @@ import {
   LlamaIndexStream,
   convertMessageContent,
 } from "./llamaindex-stream";
-import { createCallbackManager } from "./stream-helper";
+import { createCallbackManager, createStreamTimeout } from "./stream-helper";
 
 initObservability();
 initSettings();
@@ -18,6 +18,10 @@ export const runtime = "nodejs";
 export const dynamic = "force-dynamic";
 
 export async function POST(request: NextRequest) {
+  // Init Vercel AI StreamData and timeout
+  const vercelStreamData = new StreamData();
+  const streamTimeout = createStreamTimeout(vercelStreamData);
+
   try {
     const body = await request.json();
     const {
@@ -40,9 +44,6 @@ export async function POST(request: NextRequest) {
     // Convert message content from Vercel/AI format to LlamaIndex/OpenAI format
     const userMessageContent = convertMessageContent(userMessage.content, data);
 
-    // Init Vercel AI StreamData
-    const vercelStreamData = new StreamData();
-
     // Setup callbacks
     const callbackManager = createCallbackManager(vercelStreamData);
 
@@ -75,5 +76,7 @@ export async function POST(request: NextRequest) {
         status: 500,
       },
     );
+  } finally {
+    clearTimeout(streamTimeout);
   }
 }

templates/types/streaming/nextjs/app/api/chat/stream-helper.ts

@@ -75,6 +75,15 @@ export function appendToolData(
   });
 }
 
+export function createStreamTimeout(stream: StreamData) {
+  const timeout = Number(process.env.STREAM_TIMEOUT ?? 1000 * 60 * 5); // default to 5 minutes
+  const t = setTimeout(() => {
+    appendEventData(stream, `Stream timed out after ${timeout / 1000} seconds`);
+    stream.close();
+  }, timeout);
+  return t;
+}
+
 export function createCallbackManager(stream: StreamData) {
   const callbackManager = new CallbackManager();