feat: implement csv upload

.changeset/bright-turkeys-melt.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Add CSV upload

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

@@ -1,32 +1,23 @@
 import { Message, StreamData, streamToResponse } from "ai";
 import { Request, Response } from "express";
-import { ChatMessage, MessageContent, Settings } from "llamaindex";
+import { ChatMessage, Settings } from "llamaindex";
 import { createChatEngine } from "./engine/chat";
-import { LlamaIndexStream } from "./llamaindex-stream";
-import { createCallbackManager } from "./stream-helper";
-
-const convertMessageContent = (
-  textMessage: string,
-  imageUrl: string | undefined,
-): MessageContent => {
-  if (!imageUrl) return textMessage;
-  return [
-    {
-      type: "text",
-      text: textMessage,
-    },
-    {
-      type: "image_url",
-      image_url: {
-        url: imageUrl,
-      },
-    },
-  ];
-};
+import {
+  DataParserOptions,
+  LlamaIndexStream,
+  convertMessageContent,
+} from "./llamaindex-stream";
+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, data }: { messages: Message[]; data: any } = req.body;
+    const {
+      messages,
+      data,
+    }: { messages: Message[]; data: DataParserOptions | undefined } = req.body;
     const userMessage = messages.pop();
     if (!messages || !userMessage || userMessage.role !== "user") {
       return res.status(400).json({
@@ -38,13 +29,7 @@ export const chat = async (req: Request, res: Response) => {
     const chatEngine = await createChatEngine();
 
     // Convert message content from Vercel/AI format to LlamaIndex/OpenAI format
-    const userMessageContent = convertMessageContent(
-      userMessage.content,
-      data?.imageUrl,
-    );
-
-    // Init Vercel AI StreamData
-    const vercelStreamData = new StreamData();
+    const userMessageContent = convertMessageContent(userMessage.content, data);
 
     // Setup callbacks
     const callbackManager = createCallbackManager(vercelStreamData);
@@ -61,7 +46,8 @@ export const chat = async (req: Request, res: Response) => {
     // Return a stream, which can be consumed by the Vercel/AI client
     const stream = LlamaIndexStream(response, vercelStreamData, {
       parserOptions: {
-        image_url: data?.imageUrl,
+        imageUrl: data?.imageUrl,
+        uploadedCsv: data?.uploadedCsv,
       },
     });
 
@@ -71,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/llamaindex-stream.ts

@@ -6,35 +6,74 @@ import {
   type AIStreamCallbacksAndOptions,
 } from "ai";
 import {
+  MessageContent,
   Metadata,
   NodeWithScore,
   Response,
   ToolCallLLMMessageOptions,
 } from "llamaindex";
 
 import { AgentStreamChatResponse } from "llamaindex/agent/base";
-import { appendImageData, appendSourceData } from "./stream-helper";
+import {
+  UploadedCsv,
+  appendCsvData,
+  appendImageData,
+  appendSourceData,
+} from "./stream-helper";
 
 type LlamaIndexResponse =
   | AgentStreamChatResponse<ToolCallLLMMessageOptions>
   | Response;
 
-type ParserOptions = {
-  image_url?: string;
+export type DataParserOptions = {
+  imageUrl?: string;
+  uploadedCsv?: UploadedCsv;
+};
+
+export const convertMessageContent = (
+  textMessage: string,
+  additionalData?: DataParserOptions,
+): MessageContent => {
+  if (additionalData?.imageUrl) {
+    return [
+      {
+        type: "text",
+        text: textMessage,
+      },
+      {
+        type: "image_url",
+        image_url: {
+          url: additionalData?.imageUrl,
+        },
+      },
+    ];
+  }
+
+  if (additionalData?.uploadedCsv) {
+    const csvContent =
+      "Use the following CSV data:\n" +
+      "```csv\n" +
+      additionalData.uploadedCsv.content +
+      "\n```";
+    return `${csvContent}\n\n${textMessage}`;
+  }
+
+  return textMessage;
 };
 
 function createParser(
   res: AsyncIterable<LlamaIndexResponse>,
   data: StreamData,
-  opts?: ParserOptions,
+  opts?: DataParserOptions,
 ) {
   const it = res[Symbol.asyncIterator]();
   const trimStartOfStream = trimStartOfStreamHelper();
 
   let sourceNodes: NodeWithScore<Metadata>[] | undefined;
   return new ReadableStream<string>({
     start() {
-      appendImageData(data, opts?.image_url);
+      appendImageData(data, opts?.imageUrl);
+      appendCsvData(data, opts?.uploadedCsv);
     },
     async pull(controller): Promise<void> {
       const { value, done } = await it.next();
@@ -72,7 +111,7 @@ export function LlamaIndexStream(
   data: StreamData,
   opts?: {
     callbacks?: AIStreamCallbacksAndOptions;
-    parserOptions?: ParserOptions;
+    parserOptions?: DataParserOptions;
   },
 ): ReadableStream<Uint8Array> {
   return createParser(response, data, opts?.parserOptions)

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

@@ -17,6 +17,20 @@ export function appendImageData(data: StreamData, imageUrl?: string) {
   });
 }
 
+export type UploadedCsv = {
+  content: string;
+  filename: string;
+  filesize: number;
+};
+
+export function appendCsvData(data: StreamData, uploadedCsv?: UploadedCsv) {
+  if (!uploadedCsv) return;
+  data.appendMessageAnnotation({
+    type: "csv",
+    data: uploadedCsv,
+  });
+}
+
 export function appendSourceData(
   data: StreamData,
   sourceNodes?: NodeWithScore<Metadata>[],
@@ -65,6 +79,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/llamaindex-stream.ts

@@ -6,35 +6,74 @@ import {
   type AIStreamCallbacksAndOptions,
 } from "ai";
 import {
+  MessageContent,
   Metadata,
   NodeWithScore,
   Response,
   ToolCallLLMMessageOptions,
 } from "llamaindex";
 
 import { AgentStreamChatResponse } from "llamaindex/agent/base";
-import { appendImageData, appendSourceData } from "./stream-helper";
+import {
+  UploadedCsv,
+  appendCsvData,
+  appendImageData,
+  appendSourceData,
+} from "./stream-helper";
 
 type LlamaIndexResponse =
   | AgentStreamChatResponse<ToolCallLLMMessageOptions>
   | Response;
 
-type ParserOptions = {
-  image_url?: string;
+export type DataParserOptions = {
+  imageUrl?: string;
+  uploadedCsv?: UploadedCsv;
+};
+
+export const convertMessageContent = (
+  textMessage: string,
+  additionalData?: DataParserOptions,
+): MessageContent => {
+  if (additionalData?.imageUrl) {
+    return [
+      {
+        type: "text",
+        text: textMessage,
+      },
+      {
+        type: "image_url",
+        image_url: {
+          url: additionalData?.imageUrl,
+        },
+      },
+    ];
+  }
+
+  if (additionalData?.uploadedCsv) {
+    const csvContent =
+      "Use the following CSV data:\n" +
+      "```csv\n" +
+      additionalData.uploadedCsv.content +
+      "\n```";
+    return `${csvContent}\n\n${textMessage}`;
+  }
+
+  return textMessage;
 };
 
 function createParser(
   res: AsyncIterable<LlamaIndexResponse>,
   data: StreamData,
-  opts?: ParserOptions,
+  opts?: DataParserOptions,
 ) {
   const it = res[Symbol.asyncIterator]();
   const trimStartOfStream = trimStartOfStreamHelper();
 
   let sourceNodes: NodeWithScore<Metadata>[] | undefined;
   return new ReadableStream<string>({
     start() {
-      appendImageData(data, opts?.image_url);
+      appendImageData(data, opts?.imageUrl);
+      appendCsvData(data, opts?.uploadedCsv);
     },
     async pull(controller): Promise<void> {
       const { value, done } = await it.next();
@@ -72,7 +111,7 @@ export function LlamaIndexStream(
   data: StreamData,
   opts?: {
     callbacks?: AIStreamCallbacksAndOptions;
-    parserOptions?: ParserOptions;
+    parserOptions?: DataParserOptions;
   },
 ): ReadableStream<Uint8Array> {
   return createParser(response, data, opts?.parserOptions)