Add use llama parse for file loader in TS template

Author: leehuwuj

templates/components/llamaindex/typescript/documents/documents.ts

@@ -9,11 +9,9 @@ import {
   storageContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
-import { DocxReader } from "llamaindex/readers/DocxReader";
-import { PDFReader } from "llamaindex/readers/PDFReader";
-import { TextFileReader } from "llamaindex/readers/TextFileReader";
 import crypto from "node:crypto";
 import { getDataSource } from "../../engine";
+import { getExtractors } from "../../engine/loader";
 
 const MIME_TYPE_TO_EXT: Record<string, string> = {
   "application/pdf": "pdf",
@@ -58,23 +56,14 @@ async function runPipeline(
 }
 
 async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
-  console.log(`Processing uploaded document of type: ${mimeType}`);
-  switch (mimeType) {
-    case "application/pdf": {
-      const pdfReader = new PDFReader();
-      return await pdfReader.loadDataAsContent(new Uint8Array(fileBuffer));
-    }
-    case "text/plain": {
-      const textReader = new TextFileReader();
-      return await textReader.loadDataAsContent(fileBuffer);
-    }
-    case "application/vnd.openxmlformats-officedocument.wordprocessingml.document": {
-      const docxReader = new DocxReader();
-      return await docxReader.loadDataAsContent(fileBuffer);
-    }
-    default:
-      throw new Error(`Unsupported document type: ${mimeType}`);
+  const extractors = getExtractors();
+  const reader = extractors[MIME_TYPE_TO_EXT[mimeType]];
+
+  if (!reader) {
+    throw new Error(`Unsupported document type: ${mimeType}`);
   }
+  console.log(`Processing uploaded document of type: ${mimeType}`);
+  return await reader.loadDataAsContent(fileBuffer);
 }
 
 async function saveDocument(fileBuffer: Buffer, mimeType: string) {

templates/components/loaders/typescript/file/loader.ts

@@ -1,7 +1,14 @@
-import { SimpleDirectoryReader } from "llamaindex";
+import {
+  FILE_EXT_TO_READER,
+  SimpleDirectoryReader,
+} from "llamaindex/readers/SimpleDirectoryReader";
 
 export const DATA_DIR = "./data";
 
+export function getExtractors() {
+  return FILE_EXT_TO_READER;
+}
+
 export async function getDocuments() {
   return await new SimpleDirectoryReader().loadData({
     directoryPath: DATA_DIR,

templates/components/loaders/typescript/llama_parse/loader.ts

@@ -1,19 +1,26 @@
+import { LlamaParseReader } from "llamaindex/readers/LlamaParseReader";
 import {
   FILE_EXT_TO_READER,
-  LlamaParseReader,
   SimpleDirectoryReader,
-} from "llamaindex";
+} from "llamaindex/readers/SimpleDirectoryReader";
 
 export const DATA_DIR = "./data";
 
+export function getExtractors() {
+  const llamaParseParser = new LlamaParseReader({ resultType: "markdown" });
+  const extractors = FILE_EXT_TO_READER;
+  // Change all the supported extractors to LlamaParse
+  for (const key in extractors) {
+    extractors[key] = llamaParseParser;
+  }
+  return extractors;
+}
+
 export async function getDocuments() {
   const reader = new SimpleDirectoryReader();
-  // Load PDFs using LlamaParseReader
+  const extractors = getExtractors();
   return await reader.loadData({
     directoryPath: DATA_DIR,
-    fileExtToReader: {
-      ...FILE_EXT_TO_READER,
-      pdf: new LlamaParseReader({ resultType: "markdown" }),
-    },
+    fileExtToReader: extractors,
   });
 }

templates/types/streaming/nextjs/package.json

@@ -24,7 +24,7 @@
     "duck-duck-scrape": "^2.2.5",
     "formdata-node": "^6.0.3",
     "got": "^14.4.1",
-    "llamaindex": "0.4.6",
+    "llamaindex": "0.5.3",
     "lucide-react": "^0.294.0",
     "next": "^14.2.4",
     "react": "^18.2.0",