chore: use `unpdf` (#849)

.changeset/dry-cameras-travel.md

+---
+"llamaindex": patch
+"@llamaindex/core-test": patch
+---
+
+chore: use `unpdf`

packages/core/package.json

@@ -52,18 +52,23 @@
     "openai": "^4.48.1",
     "papaparse": "^5.4.1",
     "pathe": "^1.1.2",
-    "pdf2json": "3.1.3",
     "pg": "^8.12.0",
     "pgvector": "^0.1.8",
     "portkey-ai": "^0.1.16",
     "rake-modified": "^1.0.8",
     "string-strip-html": "^13.4.8",
+    "unpdf": "^0.10.1",
     "wikipedia": "^2.1.2",
     "wink-nlp": "^2.3.0"
   },
   "peerDependencies": {
     "@notionhq/client": "^2.2.15"
   },
+  "peerDependenciesMeta": {
+    "@notionhq/client": {
+      "optional": true
+    }
+  },
   "devDependencies": {
     "@notionhq/client": "^2.2.15",
     "@swc/cli": "^0.3.12",

packages/core/src/readers/PDFReader.ts

+import { fs } from "@llamaindex/env";
 import { Document } from "../Node.js";
 import { FileReader } from "./type.js";
 
@@ -5,29 +6,30 @@ import { FileReader } from "./type.js";
  * Read the text of a PDF
  */
 export class PDFReader extends FileReader {
-  async loadDataAsContent(fileContent: Buffer): Promise<Document[]> {
-    const pages = await readPDF(fileContent);
-    return pages.map((text, page) => {
+  async loadData(file: string): Promise<Document[]> {
+    const content = await fs.readFile(file);
+    return this.loadDataAsContent(new Uint8Array(content.buffer));
+  }
+
+  async loadDataAsContent(content: Uint8Array): Promise<Document[]> {
+    const { totalPages, text } = await readPDF(content);
+    return text.map((text, page) => {
       const metadata = {
         page_number: page + 1,
+        total_pages: totalPages,
       };
       return new Document({ text, metadata });
     });
   }
 }
 
-async function readPDF(data: Buffer): Promise<string[]> {
-  const parser = await import("pdf2json").then(
-    ({ default: Pdfparser }) => new Pdfparser(null, true),
-  );
-  const text = await new Promise<string>((resolve, reject) => {
-    parser.on("pdfParser_dataError", (error) => {
-      reject(error);
-    });
-    parser.on("pdfParser_dataReady", () => {
-      resolve((parser as any).getRawTextContent() as string);
-    });
-    parser.parseBuffer(data);
-  });
-  return text.split(/----------------Page \(\d+\) Break----------------/g);
+async function readPDF(data: Uint8Array): Promise<{
+  totalPages: number;
+  text: string[];
+}> {
+  const { extractText } = await import("unpdf");
+  return (await extractText(data)) as {
+    totalPages: number;
+    text: string[];
+  };
 }

packages/core/tests/readers/.snap/basic.pdf.snap

+Sample PDF
+
+This is a simple PDF file. Fun fun fun.
+
+Lorem ipsum dolor sit amet, consectetuer adipiscing elit. Phasellus facilisis odio sed mi.
+
+Curabitur suscipit. Nullam vel nisi. Etiam semper ipsum ut lectus. Proin aliquam, erat eget
+
+pharetra commodo, eros mi condimentum quam, sed commodo justo quam ut velit.
+
+Integer
+ 
+a
+ 
+erat.
+ 
+Cras
+ 
+laoreet
+ 
+ligula
+ 
+cursus
+ 
+enim.
+ 
+Aenean
+ 
+scelerisque
+ 
+velit
+ 
+et
+ 
+tellus.
+
+Vestibulum dictum aliquet sem. Nulla facilisi. Vestibulum accumsan ante vitae elit. Nulla
+
+erat dolor, blandit in, rutrum quis, semper pulvinar, enim. Nullam varius congue risus.
+
+Vivamus sollicitudin, metus ut interdum eleifend, nisi tellus pellentesque elit, tristique
+
+accumsan eros quam et risus. Suspendisse libero odio, mattis sit amet, aliquet eget,
+
+hendrerit vel, nulla. Sed vitae augue. Aliquam erat volutpat. Aliquam feugiat vulputate nisl.
+
+Suspendisse quis nulla pretium ante pretium mollis. Proin velit ligula, sagittis at, egestas a,
+
+pulvinar quis, nisl.
+
+Pellentesque sit amet lectus. Praesent pulvinar, nunc quis iaculis sagittis, justo quam
+
+lobortis tortor, sed vestibulum dui metus venenatis est. Nunc cursus ligula. Nulla facilisi.
+
+Phasellus ullamcorper consectetuer ante. Duis tincidunt, urna id condimentum luctus, nibh
+
+ante vulputate sapien, id sagittis massa orci ut enim. Pellentesque vestibulum convallis
+
+sem. Nulla consequat quam ut nisl. Nullam est. Curabitur tincidunt dapibus lorem. Proin
+
+velit turpis, scelerisque sit amet, iaculis nec, rhoncus ac, ipsum. Phasellus lorem arcu,
+
+feugiat eu, gravida eu, consequat molestie, ipsum. Nullam vel est ut ipsum volutpat
+
+feugiat. Aenean pellentesque.
+
+In mauris. Pellentesque dui nisi, iaculis eu, rhoncus in, venenatis ac, ante. Ut odio justo,
+
+scelerisque vel, facilisis non, commodo a, pede. Cras nec massa sit amet tortor volutpat
+
+varius. Donec lacinia, neque a luctus aliquet, pede massa imperdiet ante, at varius lorem
+
+pede sed sapien. Fusce erat nibh, aliquet in, eleifend eget, commodo eget, erat. Fusce
+
+consectetuer. Cras risus tortor, porttitor nec, tristique sed, convallis semper, eros. Fusce
+
+vulputate ipsum a mauris. Phasellus mollis. Curabitur sed urna. Aliquam nec sapien non
+
+nibh pulvinar convallis. Vivamus facilisis augue quis quam. Proin cursus aliquet metus.
+
+Suspendisse lacinia. Nulla at tellus ac turpis eleifend scelerisque. Maecenas a pede vitae
+
+enim commodo interdum. Donec odio. Sed sollicitudin dui vitae justo.
+
+Morbi elit nunc, facilisis a, mollis a, molestie at, lectus. Suspendisse eget mauris eu tellus
+
+molestie cursus. Duis ut magna at justo dignissim condimentum. Cum sociis natoque
+
+penatibus et magnis dis parturient montes, nascetur ridiculus mus. Vivamus varius. Ut sit
+
+amet diam suscipit mauris ornare aliquam. Sed varius. Duis arcu. Etiam tristique massa
+
+eget dui. Phasellus congue. Aenean est erat, tincidunt eget, venenatis quis, commodo at,
+
+quam.
\ No newline at end of file

packages/core/tests/readers/pdf-reader.test.ts

+import { PDFReader } from "llamaindex";
+import { describe, expect, test } from "vitest";
+
+describe("pdf reader", () => {
+  const reader = new PDFReader();
+  test("basic.pdf", async () => {
+    const documents = await reader.loadData("../../../examples/data/basic.pdf");
+    expect(documents.length).toBe(1);
+    expect(documents[0].metadata).toEqual({
+      page_number: 1,
+      total_pages: 1,
+    });
+    await expect(documents[0].text).toMatchFileSnapshot(
+      "./.snap/basic.pdf.snap",
+    );
+  });
+  test("brk-2022.pdf", async () => {
+    const documents = await reader.loadData(
+      "../../../examples/data/brk-2022.pdf",
+    );
+    expect(documents.length).toBe(140);
+  });
+  test("manga.pdf", async () => {
+    const documents = await reader.loadData("../../../examples/data/manga.pdf");
+    expect(documents.length).toBe(4);
+  });
+});