Merge pull request #152 from andrewnguonly/1.0.13

PR for app version `1.0.13`

README.md

@@ -219,22 +219,32 @@ Note: Content that is highlighted will not be cached in the vector store cache.
 - `cmd + k`: Clear all messages.
 - `cmd + ;`: Open/close Chat History panel.
 - `ctrl + c`: Cancel request (LLM request/streaming or embeddings generation)
+- `ctrl + x`: Remove file attachment.
 
 ## Multimodal
 
 Lumos supports multimodal models! Images that are present on the current page will be downloaded and bound to the model for prompting. See documentation and examples [here](./docs/multimodal.md).
 
 ## File Attachments
 
-File attachments can be uploaded to Lumos. The contents of a file will be parsed and processed through Lumos's RAG workflow (similar to processing page content). By default, the text content of a file will be parsed if the extension type is not listed below (e.g. `.py`).
+File attachments can be uploaded to Lumos. The contents of a file will be parsed and processed through Lumos's RAG workflow (similar to processing page content). By default, the text content of a file will be parsed if the extension type is not explicitly listed below.
 
 Supported extension types:
 - `.csv`
 - `.json`
-- `.txt`
+- `.pdf`
+- any plain text file format (`.txt`, `.md`, `.py`, etc)
 
 Note: If an attachment is present, page content will not be parsed. Remove the file attachment to resume parsing page content.
 
+### Image Files
+
+Image files will be processed through Lumos's [multimodal workflow](./docs/multimodal.md) (requires multimodal model).
+
+Supported image types:
+- `.jpeg`, `.jpg`
+- `.png`
+
 ## Tools (Experimental)
 
 Lumos invokes [Tools](https://js.langchain.com/docs/modules/agents/tools/) automatically based on the provided prompt. See documentation and examples [here](./docs/tools.md).

manifest.json

@@ -1,5 +1,5 @@
 {
-    "version": "1.0.12",
+    "version": "1.0.13",
     "manifest_version": 3,
     "name": "Lumos",
     "description": "An LLM co-pilot for browsing the web, powered by local LLMs. Your prompts never leave the browser.",

package.json

@@ -1,6 +1,6 @@
 {
   "name": "lumos",
-  "version": "1.0.12",
+  "version": "1.0.13",
   "private": true,
   "dependencies": {
     "@chatscope/chat-ui-kit-react": "^1.10.1",
@@ -24,6 +24,7 @@
     "fuse.js": "^7.0.0",
     "langchain": "^0.1.21",
     "markdown-to-jsx": "^7.4.1",
+    "pdf-parse": "^1.1.1",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-syntax-highlighter": "^15.5.0",

src/components/ChatBar.tsx

@@ -31,6 +31,7 @@ import { v4 as uuidv4 } from "uuid";
 
 import { getContentConfig } from "../contentConfig";
 import { useThemeContext } from "../contexts/ThemeContext";
+import { getBase64Str, getExtension } from "../document_loaders/util";
 import {
   CHAT_CONTAINER_HEIGHT_MAX,
   CHAT_CONTAINER_HEIGHT_MIN,
@@ -130,10 +131,22 @@ const ChatBar: React.FC = () => {
 
     if (fileUploaded) {
       const reader = new FileReader();
-      reader.onloadend = () => {
+      reader.onloadend = async () => {
+        let base64 = reader.result as string;
+
+        // Override base64 encoded string with content from document loaded
+        // files. Loading files on the client side (as opposed to the background
+        // script) is a workaround for using document loaders that require
+        // DOM/browser APIs.
+        const extensions = [".pdf"];
+        const extension = getExtension(fileUploaded.name, true);
+        if (extensions.includes(extension)) {
+          base64 = await getBase64Str(fileUploaded);
+        }
+
         const attachment: Attachment = {
           name: fileUploaded.name,
-          base64: reader.result as string,
+          base64: base64,
           lastModified: fileUploaded.lastModified,
         };
 
@@ -272,7 +285,7 @@ const ChatBar: React.FC = () => {
         if (activeTabUrl.protocol === "chrome:" || attachment) {
           // skip script injection for chrome:// urls or if an attachment is present
           const result = new Array(1);
-          result[0] = { result: [prompt, false, []] };
+          result[0] = { result: ["", false, []] };
           return result;
         } else {
           return chrome.scripting.executeScript({
@@ -387,6 +400,10 @@ const ChatBar: React.FC = () => {
           // cancel request
           cancelRequest();
           break;
+        case "x":
+          // remove attachment
+          handleAttachmentDelete();
+          break;
       }
     }
   };
@@ -633,7 +650,10 @@ const ChatBar: React.FC = () => {
         )}
         <div style={{ flex: 1 }}></div>
         {attachment && (
-          <Tooltip placement="top" title={`Unattach ${attachment.name}`}>
+          <Tooltip
+            placement="top"
+            title={`Unattach ${attachment.name} (ctrl + x)`}
+          >
             <IconButton
               disabled={submitDisabled}
               onClick={handleAttachmentDelete}

src/document_loaders/csv.ts

@@ -2,10 +2,16 @@ import { dsvFormat } from "d3-dsv";
 import { CSVLoader } from "langchain/document_loaders/fs/csv";
 
 export class CSVPackedLoader extends CSVLoader {
+  /**
+   * This function is copied from the CSVLoader class with a few
+   * modifications so that it's able to run in a Chrome extension
+   * context.
+   */
   public async parse(raw: string): Promise<string[]> {
     const { column, separator = "," } = this.options;
 
     const psv = dsvFormat(separator);
+    // cannot use psv.parse(), unsafe-eval is not allowed
     let parsed = psv.parseRows(raw.trim());
 
     if (column !== undefined) {

src/document_loaders/dynamic_file.ts

@@ -2,6 +2,8 @@ import { Document } from "@langchain/core/documents";
 import { BaseDocumentLoader } from "langchain/document_loaders/base";
 import { TextLoader } from "langchain/document_loaders/fs/text";
 
+import { getExtension } from "./util";
+
 export interface LoadersMapping {
   [extension: string]: (file: File) => BaseDocumentLoader;
 }
@@ -27,8 +29,7 @@ export class DynamicFileLoader extends BaseDocumentLoader {
 
   public async load(): Promise<Document[]> {
     const documents: Document[] = [];
-    const extension =
-      "." + this.file.name.split(".").pop()?.toLowerCase() || "";
+    const extension = getExtension(this.file.name, true);
     let loader;
 
     if (extension !== "" && extension in this.loaders) {

src/document_loaders/util.ts

@@ -0,0 +1,56 @@
+import { Document } from "@langchain/core/documents";
+import { JSONLoader } from "langchain/document_loaders/fs/json";
+import { WebPDFLoader } from "langchain/document_loaders/web/pdf";
+
+import { Attachment } from "../components/ChatBar";
+import { CSVPackedLoader } from "../document_loaders/csv";
+import { DynamicFileLoader } from "../document_loaders/dynamic_file";
+
+export const getExtension = (fileName: string, withDot = false): string => {
+  const extension = fileName.split(".").pop()?.toLowerCase() || "";
+  return withDot ? `.${extension}` : extension;
+};
+
+export const getBase64Str = async (file: File): Promise<string> => {
+  const loader = new DynamicFileLoader(file, {
+    // add more loaders that require DOM/browser APIs here
+    ".pdf": (file) => new WebPDFLoader(file, { splitPages: false }),
+  });
+  const docs = await loader.load();
+
+  // construct new base64 encoded string
+  let pageContent = "";
+  for (const doc of docs) {
+    pageContent += doc.pageContent + "\n\n";
+  }
+  const utf8Bytes = new TextEncoder().encode(pageContent);
+  let binary = "";
+  utf8Bytes.forEach((byte) => {
+    binary += String.fromCharCode(byte);
+  });
+
+  return `data:${file.type};base64,${btoa(binary)}`;
+};
+
+export const getDocuments = async (
+  attachment: Attachment,
+): Promise<Document[]> => {
+  const base64 = attachment.base64;
+  const byteString = atob(base64.split(",")[1]);
+  const mimeString = base64.split(",")[0].split(":")[1].split(";")[0];
+  const ab = new ArrayBuffer(byteString.length);
+  const ia = new Uint8Array(ab);
+  for (let i = 0; i < byteString.length; i++) {
+    ia[i] = byteString.charCodeAt(i);
+  }
+  const blob = new Blob([ab], { type: mimeString });
+  const file = new File([blob], attachment.name, { type: mimeString });
+
+  const loader = new DynamicFileLoader(file, {
+    // add more loaders that don't require DOM/browser APIs here
+    ".csv": (file) => new CSVPackedLoader(file),
+    ".json": (file) => new JSONLoader(file),
+  });
+
+  return await loader.load();
+};

src/scripts/background.ts

@@ -16,14 +16,11 @@ import {
 import { Runnable, RunnableSequence } from "@langchain/core/runnables";
 import { ConsoleCallbackHandler } from "@langchain/core/tracers/console";
 import { IterableReadableStream } from "@langchain/core/utils/stream";
-import { JSONLoader } from "langchain/document_loaders/fs/json";
-import { TextLoader } from "langchain/document_loaders/fs/text";
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
 import { formatDocumentsAsString } from "langchain/util/document";
 
 import { Attachment, LumosMessage } from "../components/ChatBar";
-import { CSVPackedLoader } from "../document_loaders/csv";
-import { DynamicFileLoader } from "../document_loaders/dynamic_file";
+import { getDocuments, getExtension } from "../document_loaders/util";
 import {
   DEFAULT_KEEP_ALIVE,
   getLumosOptions,
@@ -113,34 +110,28 @@ const createDocuments = async (
   chunkSize: number,
   chunkOverlap: number,
 ): Promise<Document[]> => {
+  const documents: Document[] = [];
+
   if (attachments.length > 0) {
-    // Convert base64 to Blob
-    const attachment = attachments[0];
-    const base64 = attachment.base64;
-    const byteString = atob(base64.split(",")[1]);
-    const mimeString = base64.split(",")[0].split(":")[1].split(";")[0];
-    const ab = new ArrayBuffer(byteString.length);
-    const ia = new Uint8Array(ab);
-    for (let i = 0; i < byteString.length; i++) {
-      ia[i] = byteString.charCodeAt(i);
+    for (const attachment of attachments) {
+      const extension = getExtension(attachment.name);
+      if (!SUPPORTED_IMG_FORMATS.includes(extension)) {
+        // only add non-image attachments
+        documents.push(...(await getDocuments(attachment)));
+      }
     }
-    const blob = new Blob([ab], { type: mimeString });
-    const file = new File([blob], attachment.name, { type: mimeString });
+  }
 
-    const loader = new DynamicFileLoader(file, {
-      ".csv": (file) => new CSVPackedLoader(file),
-      ".json": (file) => new JSONLoader(file),
-      ".txt": (file) => new TextLoader(file),
-    });
-    return await loader.load();
-  } else {
+  if (context !== "") {
     // split page content into overlapping documents
     const splitter = new RecursiveCharacterTextSplitter({
       chunkSize: chunkSize,
       chunkOverlap: chunkOverlap,
     });
-    return await splitter.createDocuments([context]);
+    documents.push(...(await splitter.createDocuments([context])));
   }
+
+  return documents;
 };
 
 const downloadImages = async (imageURLs: string[]): Promise<string[]> => {
@@ -167,16 +158,14 @@ const downloadImages = async (imageURLs: string[]): Promise<string[]> => {
 
     if (response.ok) {
       const blob = await response.blob();
-      let base64String: string = await new Promise((resolve) => {
+      const base64String: string = await new Promise((resolve) => {
         const reader = new FileReader();
         reader.readAsDataURL(blob);
         reader.onloadend = () => {
           resolve(reader.result as string);
         };
       });
 
-      // remove leading data url prefix `data:*/*;base64,`
-      base64String = base64String.split(",")[1];
       base64EncodedImages.push(base64String);
     } else {
       console.log(`Failed to download image: ${url}`);
@@ -235,7 +224,7 @@ const getMessages = async (
       base64EncodedImages.forEach((image) => {
         content.push({
           type: "image_url",
-          image_url: `data:image/*;base64,${image}`,
+          image_url: image,
         });
       });
 
@@ -370,7 +359,19 @@ chrome.runtime.onMessage.addListener(async (request) => {
         CLS_IMG_TRIGGER,
       ))
     ) {
-      base64EncodedImages = await downloadImages(request.imageURLs);
+      // first, try to get images from attachments
+      if (attachments.length > 0) {
+        for (const attachment of attachments) {
+          const extension = getExtension(attachment.name);
+          if (SUPPORTED_IMG_FORMATS.includes(extension)) {
+            base64EncodedImages.push(attachment.base64);
+          }
+        }
+      }
+      // then, try to download images from URLs
+      if (base64EncodedImages.length === 0) {
+        base64EncodedImages = await downloadImages(request.imageURLs);
+      }
     } else if (
       options.toolConfig["Calculator"].enabled &&
       (await classifyPrompt(
@@ -476,7 +477,7 @@ chrome.runtime.onMessage.addListener(async (request) => {
   }
 
   // process parsed context
-  if (request.context) {
+  if (request.context || request.attachments) {
     context = request.context;
     attachments = request.attachments;
     console.log(`Received context: ${context}`);