Database and search submission

app/api/embedding/get/route.ts

@@ -0,0 +1,48 @@
+// app/api/embedding/get/route.ts
+import EmbeddingStorage from "../../../utils/embeddingStorage";
+import { storageMethod } from "../../../utils/env";
+import { StorageMethod } from "../../../utils/types";
+
+async function handleRequest(req: Request): Promise<Response> {
+  let db: EmbeddingStorage | undefined;
+  try {
+    const body = await req.json();
+    let highlights;
+
+    if (storageMethod === StorageMethod.sqlite) {
+      db = new EmbeddingStorage();
+      highlights = await db.searchEmbedding(body.pdfId, body.query);
+    } else {
+    console.log("Databse not initialized");
+  }
+
+    return new Response(JSON.stringify(highlights), {
+      status: 200,
+      headers: { "Content-Type": "application/json" },
+    });
+  } catch (error) {
+    console.error("Error in handleRequest:", error);
+    return new Response(
+      JSON.stringify({
+        error: "Internal Server Error",
+        details: error.message,
+      }),
+      {
+        status: 500,
+        headers: { "Content-Type": "application/json" },
+      }
+    );
+  } finally {
+    if (db) {
+      try {
+        await db.close();
+      } catch (closeError) {
+        console.error("Error closing database:", closeError);
+      }
+    }
+  }
+}
+
+export async function POST(req: Request): Promise<Response> {
+  return handleRequest(req);
+}

app/api/embedding/upload/route.ts

@@ -0,0 +1,57 @@
+// app/api/embedding/upload/route.ts
+import EmbeddingStorage from "../../../utils/embeddingStorage";
+import { storageMethod } from "../../../utils/env";
+import { StorageMethod, StoredEmbedding } from "../../../utils/types";
+
+async function handleRequest(
+  req: Request,
+  action: (body: any, db?: EmbeddingStorage) => Promise<void>
+): Promise<Response> {
+  let db: EmbeddingStorage | undefined;
+  try {
+    const body = await req.json();
+    if (storageMethod === StorageMethod.sqlite) {
+      db = new EmbeddingStorage();
+    }
+    await action(body, db);
+    return new Response(null, { status: 200 });
+  } catch (error) {
+    console.error(error);
+    return new Response(null, { status: 500 });
+  } finally {
+    if (db) {
+      await db.close();
+    }
+  }
+}
+
+async function saveEmbedding(body: any, db?: EmbeddingStorage): Promise<void> {
+  if (db) {
+    if (Array.isArray(body.embeddings)) {
+      await db.saveBulkEmbedding(body.embeddings);
+    } else {
+      await db.saveEmbedding(body.embeddings);
+    }
+  } else {
+    console.log("Databse not initialized");
+  }
+}
+
+async function removeEmbedding(
+  body: any,
+  db?: EmbeddingStorage
+): Promise<void> {
+  if (db) {
+    await db.deleteEmbedding(body.pdfId, body.id);
+  } else {
+    console.log("Databse not initialized");
+  }
+}
+
+export async function POST(req: Request): Promise<Response> {
+  return handleRequest(req, saveEmbedding);
+}
+
+export async function DELETE(req: Request): Promise<Response> {
+  return handleRequest(req, removeEmbedding);
+}

app/api/ocr/route.ts

@@ -0,0 +1,38 @@
+// app/api/ocr/route.ts
+
+import { googleApiKey } from "../../utils/env";
+
+// API for server to call the Google Vision API
+async function handleRequest(req: Request): Promise<Response> {
+  try {
+    const body = await req.json();
+
+    const googleResponse = await fetch("https://vision.googleapis.com/v1/images:annotate?key=" + googleApiKey, {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ requests: body.requests })
+    });
+    const googleResponseValue = await googleResponse.json();
+
+    return new Response(JSON.stringify(googleResponseValue), {
+      status: 200,
+      headers: { "Content-Type": "application/json" },
+    });
+  } catch (error) {
+    console.error("Error in handleRequest:", error);
+    return new Response(
+      JSON.stringify({
+        error: "Internal Server Error",
+        details: error.message,
+      }),
+      {
+        status: 500,
+        headers: { "Content-Type": "application/json" },
+      }
+    );
+  }
+}
+
+export async function POST(req: Request): Promise<Response> {
+  return handleRequest(req);
+}

app/components/App.tsx

@@ -6,10 +6,10 @@ import KeywordSearch from "./KeywordSearch";
 import PdfViewer from "./PdfViewer";
 import { Header } from "./Header";
 import Spinner from "./Spinner";
-import { convertPdfToImages, searchPdf } from "../utils/pdfUtils";
+import { convertPdfToImages, searchPdf, createSearchablePDF } from "../utils/pdfUtils";
 import type { IHighlight } from "react-pdf-highlighter";
 import HighlightUploader from "./HighlightUploader";
-import { StoredHighlight, StorageMethod } from "../utils/types";
+import { StoredHighlight, StorageMethod, StoredEmbedding } from "../utils/types";
 import {
   IHighlightToStoredHighlight,
   StoredHighlightToIHighlight,
@@ -49,17 +49,36 @@ export default function App() {
     //   perform OCR,
     //   convert output back to PDF
     //   update file url with new PDF url
-    const i = await convertPdfToImages(file);
-    const worker = await createWorker("eng");
-    const res = await worker.recognize(
-      i[0],
-      { pdfTitle: "ocr-out" },
-      { pdf: true }
-    );
-    const pdf = res.data.pdf;
-    if (pdf) {
+    const images = await convertPdfToImages(file);
+
+    const base64Pages = images.map((item): string => {return item.split(',')[1]});
+
+    const requests = base64Pages.map(base64 => ({
+      image: { content: base64 },
+      features: [{ type: 'DOCUMENT_TEXT_DETECTION' }]
+    }));
+
+    const fullAnnotations = [];
+
+    // Maximum number of pages per API call is 16
+    const maxPages = 16;
+    for (let i = 0; i < requests.length; i += maxPages) {
+      const googleOcrRes = await fetch("/api/ocr", {
+          method: "POST",
+          headers: { "Content-Type": "application/json" },
+          body: JSON.stringify({ requests: requests.slice(i, Math.min(i + maxPages, requests.length)) }),
+      });
+
+      const body = await googleOcrRes.json();
+
+      for (let j = 0; j < body.responses.length; ++j) {
+        fullAnnotations.push(body.responses[j].fullTextAnnotation);
+      }
+    }
+
+    if (fullAnnotations.length > 0) {
       // Update file url if OCR success
-      const blob = new Blob([new Uint8Array(pdf)], { type: "application/pdf" });
+      const blob = await createSearchablePDF(base64Pages, fullAnnotations);
       const fileOcrUrl = URL.createObjectURL(blob);
       setPdfOcrUrl(fileOcrUrl);
 

app/globals.css

@@ -8,7 +8,7 @@
   --background-end-rgb: 255, 255, 255;
 }
 
-@media (prefers-color-scheme: dark) {
+@media (prefers-color-scheme: light) {
   :root {
     --foreground-rgb: 255, 255, 255;
     --background-start-rgb: 0, 0, 0;

app/utils/embedding.ts

@@ -0,0 +1,25 @@
+// app/utils/embedding.ts
+
+import { openAiKey } from "./env";
+
+export async function embedText(text: string): Promise<number[]> {
+  const res = await fetch("https://api.openai.com/v1/embeddings", {
+    method: "POST",
+    headers: {
+      Authorization: "Bearer ${process.env.}" + openAiKey,
+      "Content-Type": "application/json"
+    },
+    body: JSON.stringify({
+      model: "text-embedding-ada-002",
+      input: text
+    }),
+  });
+
+  const data = await res.json();
+
+  if (!data) {
+    return [];
+  }
+
+  return data.data[0].embedding;
+}

app/utils/embeddingStorage.ts

@@ -0,0 +1,152 @@
+// app/utils/embeddingStorage.ts
+
+import sqlite3 from "sqlite3"
+import path from "path";
+import { StoredEmbedding } from "./types";
+import { embedText } from "./embedding";
+
+function euclideanDistance(a: number[], b: number[]) {
+  return Math.sqrt(a.reduce((sum, ai, i) => sum + Math.pow(ai - b[i], 2), 0));
+}
+
+class EmbeddingDatabase {
+  private db: sqlite3.Database;
+  private tableName: string = "pages";
+
+  private migrationPromise: Promise<void>;
+
+  constructor() {
+    this.db = new sqlite3.Database(
+      path.join(process.cwd(), "embeddings.db"),
+      sqlite3.OPEN_READWRITE | sqlite3.OPEN_CREATE,
+      (error) => {
+        if (error) {
+          console.error("Error opening database:", error.message);
+        } else {
+          console.log("Connected to embeddings db!");
+        }
+      }
+    );
+    this.migrationPromise = this.migrate();
+  }
+
+  private migrate(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      const sql = `
+        CREATE TABLE ${this.tableName} (
+          id UUID PRIMARY KEY,
+          pdfId INT,
+          pageNum INT,
+          text TEXT,
+          embedding VECTOR(1536)
+        );
+      `;
+      this.db.run(sql, (err) => {
+        if (err) {
+          console.error("Error creating table:", err.message);
+          reject(err);
+        } else {
+          console.log("Highlights table created or already exists");
+          resolve();
+        }
+      });
+    });
+  }
+
+  private async ensureMigrated(): Promise<void> {
+    await this.migrationPromise;
+  }
+
+  async saveEmbedding(embedding: StoredEmbedding): Promise<void> {
+    await this.ensureMigrated();
+    const sql = `INSERT OR REPLACE INTO ${this.tableName} (id, pdfId, pageNum, text, embedding) VALUES (?, ?, ?, ?, ?)`;
+    return new Promise((resolve, reject) => {
+      this.db.run(sql, Object.values(embedding), (error) => {
+        if (error) reject(error);
+        else resolve();
+      });
+    });
+  }
+
+  async saveBulkEmbedding(embeddings: StoredEmbedding[]): Promise<void> {
+    await this.ensureMigrated();
+    const sql = `INSERT OR REPLACE INTO ${this.tableName} (id, pdfId, pageNum, text, embedding) VALUES (?, ?, ?, ?, ?)`;
+    return new Promise((resolve, reject) => {
+      this.db.serialize(() => {
+        this.db.run("BEGIN TRANSACTION");
+        const stmt = this.db.prepare(sql);
+        embeddings.forEach((embedding) => {
+          stmt.run(Object.values(embedding));
+        });
+        stmt.finalize((error) => {
+          if (error) {
+            this.db.run("ROLLBACK");
+            reject(error);
+          } else {
+            this.db.run("COMMIT", (commitError) => {
+              if (commitError) reject(commitError);
+              else resolve();
+            });
+          }
+        });
+      });
+    });
+  }
+
+  async getEmbeddingForPDF(pdfId: string): Promise<StoredEmbedding[]> {
+    await this.ensureMigrated();
+    const sql = `SELECT * FROM ${this.tableName} WHERE pdfId = ?`;
+    return new Promise((resolve, reject) => {
+      this.db.all(sql, [pdfId], (error, rows) => {
+        if (error) reject(error);
+        else resolve(rows as StoredEmbedding[]);
+      });
+    });
+  }
+
+  async searchEmbedding(query: string) {
+    const embedding = await embedText(query);
+
+    await this.ensureMigrated();
+    const sql = `
+      SELECT *
+        FROM pages
+    `;
+
+    const res: StoredEmbedding[] = await new Promise((resolve, reject) => {
+      this.db.all(sql, [], (error, rows) => {
+        if (error) reject(error);
+        else resolve(rows as StoredEmbedding[]);
+      });
+    });
+
+    return res.map(r => ({
+      ...r,
+      distance: euclideanDistance(r.embedding, embedding)
+    }))
+    .sort((a, b) => a.distance - b.distance);
+  }
+
+  async deleteEmbedding(pdfId: string, id: string): Promise<void> {
+    await this.ensureMigrated();
+    const sql = `DELETE FROM ${this.tableName} WHERE pdfId = ? AND id = ?`;
+    return new Promise((resolve, reject) => {
+      this.db.run(sql, [pdfId, id], (error) => {
+        if (error) reject(error);
+        else resolve();
+      });
+    });
+  }
+
+  async close(): Promise<void> {
+    await this.ensureMigrated();
+    return new Promise((resolve, reject) => {
+      this.db.close((error) => {
+        if (error) reject(error);
+        else resolve();
+      });
+    });
+  }
+}
+
+export default EmbeddingDatabase;