feat: support uploading pdf, docx, txt

.changeset/late-weeks-sneeze.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Support upload document files: pdf, docx, txt

templates/components/ui/html/chat/hooks/use-config.ts

@@ -3,28 +3,29 @@
 import { useEffect, useMemo, useState } from "react";
 
 export interface ChatConfig {
-  chatAPI?: string;
+  backend?: 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 [config, setConfig] = useState<ChatConfig>();
 
-  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`;
+
   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 {
+    backend: backendOrigin,
+    starterQuestions: config?.starterQuestions,
+  };
 }

templates/types/streaming/express/package.json

@@ -20,7 +20,7 @@
     "dotenv": "^16.3.1",
     "duck-duck-scrape": "^2.2.5",
     "express": "^4.18.2",
-    "llamaindex": "0.4.3",
+    "llamaindex": "0.4.4",
     "pdf2json": "3.0.5",
     "ajv": "^8.12.0",
     "@e2b/code-interpreter": "^0.0.5",

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

@@ -11,8 +11,7 @@ import {
   MessageContent,
   MessageContentDetail,
 } from "llamaindex";
-
-import { CsvFile } from "./stream-helper";
+import { DocumentFile } from "./stream-helper";
 
 export const convertMessageContent = (
   content: string,
@@ -60,17 +59,21 @@ const convertAnnotations = (
         },
       });
     }
-    // convert CSV files to text
-    if (type === "csv" && "csvFiles" in data && Array.isArray(data.csvFiles)) {
-      const rawContents = data.csvFiles.map((csv) => {
-        return "```csv\n" + (csv as CsvFile).content + "\n```";
+    // convert files to text
+    if (
+      type === "document_file" &&
+      "files" in data &&
+      Array.isArray(data.files)
+    ) {
+      const rawContents = data.files.map((file) => {
+        const { filetype, content } = file as DocumentFile;
+        return "```" + `${filetype}\n${content}\n` + "```";
       });
-      const csvContent =
-        "Use data from following CSV raw contents:\n" +
-        rawContents.join("\n\n");
+      const fileContent =
+        `Use data from following raw contents:\n` + rawContents.join("\n\n");
       content.push({
         type: "text",
-        text: csvContent,
+        text: fileContent,
       });
     }
   });

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

@@ -113,9 +113,22 @@ export function createCallbackManager(stream: StreamData) {
   return callbackManager;
 }
 
-export type CsvFile = {
+export type TextEmbedding = {
+  text: string;
+  embedding: number[];
+};
+
+export type ContentFileType = "csv" | "pdf" | "txt" | "docx";
+
+export type DocumentFile = {
+  id: string;
   content: string;
   filename: string;
   filesize: number;
-  id: string;
+  filetype: ContentFileType;
+  embeddings?: TextEmbedding[];
+};
+
+export type DocumentFileData = {
+  files: DocumentFile[];
 };

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

@@ -0,0 +1,36 @@
+import {
+  Document,
+  IngestionPipeline,
+  MetadataMode,
+  Settings,
+  SimpleNodeParser,
+} from "llamaindex";
+import pdf from "pdf-parse";
+
+export async function splitAndEmbed(content: string) {
+  const document = new Document({ text: content });
+  const pipeline = new IngestionPipeline({
+    transformations: [
+      new SimpleNodeParser({
+        chunkSize: Settings.chunkSize,
+        chunkOverlap: Settings.chunkOverlap,
+      }),
+      Settings.embedModel,
+    ],
+  });
+  const nodes = await pipeline.run({ documents: [document] });
+  return nodes.map((node, i) => ({
+    text: node.getContent(MetadataMode.NONE),
+    embedding: node.embedding,
+  }));
+}
+
+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

@@ -11,8 +11,7 @@ import {
   MessageContent,
   MessageContentDetail,
 } from "llamaindex";
-
-import { CsvFile } from "./stream-helper";
+import { DocumentFile } from "./stream-helper";
 
 export const convertMessageContent = (
   content: string,
@@ -60,17 +59,21 @@ const convertAnnotations = (
         },
       });
     }
-    // convert CSV files to text
-    if (type === "csv" && "csvFiles" in data && Array.isArray(data.csvFiles)) {
-      const rawContents = data.csvFiles.map((csv) => {
-        return "```csv\n" + (csv as CsvFile).content + "\n```";
+    // convert files to text
+    if (
+      type === "document_file" &&
+      "files" in data &&
+      Array.isArray(data.files)
+    ) {
+      const rawContents = data.files.map((file) => {
+        const { filetype, content } = file as DocumentFile;
+        return "```" + `${filetype}\n${content}\n` + "```";
       });
-      const csvContent =
-        "Use data from following CSV raw contents:\n" +
-        rawContents.join("\n\n");
+      const fileContent =
+        `Use data from following raw contents:\n` + rawContents.join("\n\n");
       content.push({
         type: "text",
-        text: csvContent,
+        text: fileContent,
       });
     }
   });

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

@@ -113,9 +113,22 @@ export function createCallbackManager(stream: StreamData) {
   return callbackManager;
 }
 
-export type CsvFile = {
+export type TextEmbedding = {
+  text: string;
+  embedding: number[];
+};
+
+export type ContentFileType = "csv" | "pdf" | "txt" | "docx";
+
+export type DocumentFile = {
+  id: string;
   content: string;
   filename: string;
   filesize: number;
-  id: string;
+  filetype: ContentFileType;
+  embeddings?: TextEmbedding[];
+};
+
+export type DocumentFileData = {
+  files: DocumentFile[];
 };

templates/types/streaming/nextjs/app/components/chat-section.tsx

@@ -5,7 +5,7 @@ import { ChatInput, ChatMessages } from "./ui/chat";
 import { useClientConfig } from "./ui/chat/hooks/use-config";
 
 export default function ChatSection() {
-  const { chatAPI } = useClientConfig();
+  const { backend } = useClientConfig();
   const {
     messages,
     input,
@@ -17,7 +17,7 @@ export default function ChatSection() {
     append,
     setInput,
   } = useChat({
-    api: chatAPI,
+    api: `${backend}/api/chat`,
     headers: {
       "Content-Type": "application/json", // using JSON because of vercel/ai 2.2.26
     },