feat: support uploading pdf, docx, txt

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

@@ -16,7 +16,7 @@ import {
 } from "llamaindex";
 
 import { AgentStreamChatResponse } from "llamaindex/agent/base";
-import { CsvFile, appendSourceData } from "./stream-helper";
+import { CsvFile, PdfFile, appendSourceData } from "./stream-helper";
 
 type LlamaIndexResponse =
   | AgentStreamChatResponse<ToolCallLLMMessageOptions>
@@ -81,6 +81,18 @@ const convertAnnotations = (
         text: csvContent,
       });
     }
+    // convert PDF files to text
+    if (type === "pdf" && "pdfFiles" in data && Array.isArray(data.pdfFiles)) {
+      const pdfContent = data.pdfFiles
+        .map((pdf) => {
+          return "```pdf\n" + (pdf as PdfFile).content + "\n```";
+        })
+        .join("\n\n");
+      content.push({
+        type: "text",
+        text: pdfContent,
+      });
+    }
   });
 
   return content;

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

@@ -118,3 +118,16 @@ export type CsvFile = {
   filesize: number;
   id: string;
 };
+
+export type TextEmbedding = {
+  text: string;
+  embedding: number[];
+};
+
+export type PdfFile = {
+  id: string;
+  content: string;
+  filename: string;
+  filesize: number;
+  embeddings: TextEmbedding[];
+};

templates/types/streaming/nextjs/app/api/chat/embed/embeddings.ts

@@ -0,0 +1,28 @@
+import { Document, MetadataMode, Settings, SimpleNodeParser } from "llamaindex";
+import pdf from "pdf-parse";
+
+export async function splitAndEmbed(document: string) {
+  const nodeParser = new SimpleNodeParser({
+    chunkSize: Settings.chunkSize,
+    chunkOverlap: Settings.chunkOverlap,
+  });
+  const nodes = nodeParser.getNodesFromDocuments([
+    new Document({ text: document }),
+  ]);
+  const texts = nodes.map((node) => node.getContent(MetadataMode.EMBED));
+  const embeddings = await Settings.embedModel.getTextEmbeddingsBatch(texts);
+  return nodes.map((node, i) => ({
+    text: node.getContent(MetadataMode.NONE),
+    embedding: embeddings[i],
+  }));
+}
+
+export async function getPdfDetail(rawPdf: string) {
+  const pdfBuffer = Buffer.from(rawPdf.split(",")[1], "base64");
+  const content = (await pdf(pdfBuffer)).text;
+  const embeddings = await splitAndEmbed(content);
+  return {
+    content,
+    embeddings,
+  };
+}

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

@@ -0,0 +1,28 @@
+import { NextRequest, NextResponse } from "next/server";
+import { initSettings } from "../engine/settings";
+import { getPdfDetail } from "./embeddings";
+
+initSettings();
+
+export async function POST(request: NextRequest) {
+  try {
+    const { pdf }: { pdf: string } = await request.json();
+    if (!pdf) {
+      return NextResponse.json(
+        { error: "pdf is required in the request body" },
+        { status: 400 },
+      );
+    }
+    const pdfDetail = await getPdfDetail(pdf);
+    return NextResponse.json(pdfDetail);
+  } catch (error) {
+    console.error("[Embed API]", error);
+    return NextResponse.json(
+      { error: (error as Error).message },
+      { status: 500 },
+    );
+  }
+}
+
+export const runtime = "nodejs";
+export const dynamic = "force-dynamic";

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

@@ -16,7 +16,7 @@ import {
 } from "llamaindex";
 
 import { AgentStreamChatResponse } from "llamaindex/agent/base";
-import { CsvFile, appendSourceData } from "./stream-helper";
+import { CsvFile, PdfFile, appendSourceData } from "./stream-helper";
 
 type LlamaIndexResponse =
   | AgentStreamChatResponse<ToolCallLLMMessageOptions>
@@ -81,6 +81,18 @@ const convertAnnotations = (
         text: csvContent,
       });
     }
+    // convert PDF files to text
+    if (type === "pdf" && "pdfFiles" in data && Array.isArray(data.pdfFiles)) {
+      const pdfContent = data.pdfFiles
+        .map((pdf) => {
+          return "```pdf\n" + (pdf as PdfFile).content + "\n```";
+        })
+        .join("\n\n");
+      content.push({
+        type: "text",
+        text: pdfContent,
+      });
+    }
   });
 
   return content;

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

@@ -118,3 +118,16 @@ export type CsvFile = {
   filesize: number;
   id: string;
 };
+
+export type TextEmbedding = {
+  text: string;
+  embedding: number[];
+};
+
+export type PdfFile = {
+  id: string;
+  content: string;
+  filename: string;
+  filesize: number;
+  embeddings: TextEmbedding[];
+};

templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx

@@ -7,8 +7,10 @@ import FileUploader from "../file-uploader";
 import { Input } from "../input";
 import UploadCsvPreview from "../upload-csv-preview";
 import UploadImagePreview from "../upload-image-preview";
+import UploadPdfPreview from "../upload-pdf-preview";
 import { ChatHandler } from "./chat.interface";
 import { useCsv } from "./hooks/use-csv";
+import { usePdf } from "./hooks/use-pdf ";
 
 export default function ChatInput(
   props: Pick<
@@ -26,9 +28,10 @@ export default function ChatInput(
 ) {
   const [imageUrl, setImageUrl] = useState<string | null>(null);
   const { files: csvFiles, upload, remove, reset } = useCsv();
+  const { pdf, setPdf, uploadAndEmbed } = usePdf();
 
   const getAnnotations = () => {
-    if (!imageUrl && csvFiles.length === 0) return undefined;
+    if (!imageUrl && csvFiles.length === 0 && !pdf) return undefined;
     const annotations: MessageAnnotation[] = [];
     if (imageUrl) {
       annotations.push({
@@ -49,6 +52,22 @@ export default function ChatInput(
         },
       });
     }
+    if (pdf) {
+      annotations.push({
+        type: MessageAnnotationType.PDF,
+        data: {
+          pdfFiles: [
+            {
+              id: pdf.id,
+              content: pdf.content,
+              filename: pdf.filename,
+              filesize: pdf.filesize,
+              embeddings: pdf.embeddings,
+            },
+          ],
+        },
+      });
+    }
     return annotations as JSONValue[];
   };
 
@@ -74,6 +93,7 @@ export default function ChatInput(
       handleSubmitWithAnnotations(e, annotations);
       imageUrl && setImageUrl(null);
       csvFiles.length && reset();
+      pdf && setPdf(null);
       return;
     }
     props.handleSubmit(e);
@@ -84,7 +104,7 @@ export default function ChatInput(
   const readContent = async (file: File): Promise<string> => {
     const content = await new Promise<string>((resolve, reject) => {
       const reader = new FileReader();
-      if (file.type.startsWith("image/")) {
+      if (file.type.startsWith("image/") || file.type === "application/pdf") {
         reader.readAsDataURL(file);
       } else {
         reader.readAsText(file);
@@ -113,6 +133,16 @@ export default function ChatInput(
     }
   };
 
+  const handleUploadPdfFile = async (file: File) => {
+    const base64 = await readContent(file);
+    await uploadAndEmbed({
+      id: uuidv4(),
+      filename: file.name,
+      filesize: file.size,
+      pdfBase64: base64,
+    });
+  };
+
   const handleUploadFile = async (file: File) => {
     try {
       if (file.type.startsWith("image/")) {
@@ -125,6 +155,13 @@ export default function ChatInput(
         }
         return await handleUploadCsvFile(file);
       }
+      if (file.type === "application/pdf") {
+        if (pdf) {
+          alert("You can only upload one pdf file at a time.");
+          return;
+        }
+        return await handleUploadPdfFile(file);
+      }
       props.onFileUpload?.(file);
     } catch (error: any) {
       props.onFileError?.(error.message);
@@ -152,6 +189,7 @@ export default function ChatInput(
           })}
         </div>
       )}
+      {pdf && <UploadPdfPreview pdf={pdf} onRemove={() => setPdf(null)} />}
       <div className="flex w-full items-start justify-between gap-4 ">
         <Input
           autoFocus

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/index.tsx

@@ -3,13 +3,15 @@ import { Check, Copy } from "lucide-react";
 import { Message } from "ai";
 import { Fragment } from "react";
 import { Button } from "../../button";
+import UploadPdfPreview from "../../upload-pdf-preview";
 import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
 import {
   CsvData,
   EventData,
   ImageData,
   MessageAnnotation,
   MessageAnnotationType,
+  PDFData,
   SourceData,
   ToolData,
   getAnnotationData,
@@ -45,6 +47,10 @@ function ChatMessageContent({
     annotations,
     MessageAnnotationType.CSV,
   );
+  const pdfData = getAnnotationData<PDFData>(
+    annotations,
+    MessageAnnotationType.PDF,
+  );
   const eventData = getAnnotationData<EventData>(
     annotations,
     MessageAnnotationType.EVENTS,
@@ -74,6 +80,12 @@ function ChatMessageContent({
       order: 2,
       component: csvData[0] ? <CsvContent data={csvData[0]} /> : null,
     },
+    {
+      order: 3,
+      component: pdfData[0] ? (
+        <UploadPdfPreview pdf={pdfData[0].pdfFiles[0]} />
+      ) : null,
+    },
     {
       order: -1,
       component: toolData[0] ? <ChatTools data={toolData[0]} /> : null,

templates/types/streaming/nextjs/app/components/ui/chat/hooks/use-config.ts

@@ -4,27 +4,33 @@ import { useEffect, useMemo, useState } from "react";
 
 export interface ChatConfig {
   chatAPI?: string;
+  embedAPI?: string;
   starterQuestions?: string[];
 }
 
-export function useClientConfig() {
-  const API_ROUTE = "/api/chat/config";
+export function useClientConfig(): ChatConfig {
   const chatAPI = process.env.NEXT_PUBLIC_CHAT_API;
   const [config, setConfig] = useState<ChatConfig>({
     chatAPI,
   });
 
-  const configAPI = useMemo(() => {
-    const backendOrigin = chatAPI ? new URL(chatAPI).origin : "";
-    return `${backendOrigin}${API_ROUTE}`;
+  const backendOrigin = useMemo(() => {
+    return chatAPI ? new URL(chatAPI).origin : "";
   }, [chatAPI]);
 
+  const configAPI = `${backendOrigin}/api/chat/config`;
+  const embedAPI = `${backendOrigin}/api/chat/embed`;
+
   useEffect(() => {
     fetch(configAPI)
       .then((response) => response.json())
       .then((data) => setConfig({ ...data, chatAPI }))
       .catch((error) => console.error("Error fetching config", error));
   }, [chatAPI, configAPI]);
 
-  return config;
+  return {
+    chatAPI,
+    embedAPI,
+    starterQuestions: config.starterQuestions,
+  };
 }

templates/types/streaming/nextjs/app/components/ui/chat/hooks/use-pdf .ts

@@ -0,0 +1,42 @@
+"use client";
+
+import { useState } from "react";
+import { PdfFile } from "..";
+import { useClientConfig } from "./use-config";
+
+export function usePdf() {
+  const { embedAPI } = useClientConfig();
+  const [pdf, setPdf] = useState<PdfFile | null>(null);
+
+  const getPdfDetail = async (
+    pdfBase64: string,
+  ): Promise<Pick<PdfFile, "content" | "embeddings">> => {
+    if (!embedAPI) throw new Error("Embed API is not defined");
+    const response = await fetch(embedAPI, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        pdf: pdfBase64,
+      }),
+    });
+    if (!response.ok) throw new Error("Failed to get pdf detail");
+    const data = await response.json();
+    return data;
+  };
+
+  const uploadAndEmbed = async (pdf: {
+    id: string;
+    filename: string;
+    filesize: number;
+    pdfBase64: string;
+  }) => {
+    const { pdfBase64, ...rest } = pdf;
+    const pdfDetail = await getPdfDetail(pdfBase64);
+    setPdf({ ...pdfDetail, ...rest });
+    return pdfDetail;
+  };
+
+  return { pdf, setPdf, uploadAndEmbed };
+}