[Examples] Modify chrome extension (#557)

- updated manifest.json to include content.js (enabling reading of page contents as context) - extended extension functionality to include all available MLC models - aesthetics changes to the UI, - included loading text and loading bar that appears as models are loaded in - included a Now Chatting With text that displays the model that we are currently chatting with
mlc-ai · Aug 31, 2024 · 1fe4fc6 · 1fe4fc6
1 parent 0cfb945
commit 1fe4fc6
Show file tree

Hide file tree

Showing 7 changed files with 161 additions and 25 deletions.
diff --git a/examples/chrome-extension/README.md b/examples/chrome-extension/README.md
@@ -9,4 +9,4 @@ npm install
 npm run build
 ```
 
-This will create a new directory at `chrome-extension/dist/`. To load the extension into Chrome, go to Extensions > Manage Extensions and select Load Unpacked. Add the `chrome-extension/dist/` directory. You can now pin the extension to your toolbar and use it to chat with your favorite model!
+This will create a new directory at `chrome-extension/dist/`. To load the extension into Chrome, go to Extensions > Manage Extensions and select Load Unpacked. Add the `chrome-extension/dist/` directory. You can now pin the extension to your toolbar and use the drop-down menu to chat with your favorite model!
diff --git a/examples/chrome-extension/package.json b/examples/chrome-extension/package.json
@@ -1,6 +1,6 @@
 {
   "name": "chrome-extension",
-  "version": "1.0.0",
+  "version": "1.0.1",
   "description": "",
   "private": true,
   "scripts": {

diff --git a/examples/chrome-extension/src/content.js b/examples/chrome-extension/src/content.js
@@ -1,6 +1,6 @@
 // Only the content script is able to access the DOM
 chrome.runtime.onConnect.addListener(function (port) {
   port.onMessage.addListener(function (msg) {
-    port.postMessage({ contents: document.body.innerHTML });
+    port.postMessage({ contents: document.body.innerText });
   });
 });
diff --git a/examples/chrome-extension/src/manifest.json b/examples/chrome-extension/src/manifest.json
@@ -1,7 +1,7 @@
 {
   "manifest_version": 3,
   "name": "MLCBot",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "Chat with your browser",
   "icons": {
     "16": "icons/icon-16.png",
@@ -16,5 +16,12 @@
     "default_title": "MLCBot",
     "default_popup": "popup.html"
   },
-  "permissions": ["storage", "tabs", "webNavigation"]
+  "content_scripts": [
+    {
+      "matches": ["<all_urls>"],
+      "js": ["content.js"]
+    }
+  ],
+  "permissions": ["storage", "tabs", "webNavigation", "activeTab", "scripting"],
+  "host_permissions": ["http://*/", "https://*/"]
 }
diff --git a/examples/chrome-extension/src/popup.css b/examples/chrome-extension/src/popup.css
@@ -21,7 +21,7 @@ body {
   margin: 0;
   padding: 0.5rem;
   background-color: #778da9;
-  width: 320px;
+  width: 335px;
   font-size: small;
 }
 
@@ -32,7 +32,7 @@ p {
 /* LOADING BAR */
 #loadingContainer {
   margin-bottom: 15px;
-  width: 300px;
+  width: 315px;
   height: 8px;
 }
 
@@ -55,7 +55,7 @@ p {
   margin-right: 0.5rem;
 }
 
-/* SUBMIT BUTTON */
+/* BUTTON */
 .btn {
   background-color: #1b263b;
   color: white;

diff --git a/examples/chrome-extension/src/popup.html b/examples/chrome-extension/src/popup.html
@@ -10,8 +10,12 @@
     />
   </head>
   <body>
-    <div id="loadingContainer"></div>
-
+    <select id="model-selection"></select>
+    <div id="loadingBox">
+      <p id="init-label">Initializing model...</p>
+      <div id="loadingContainer"></div>
+    </div>
+    <p id="model-name"></p>
     <div class="input-container form-group">
       <input
         type="search"

diff --git a/examples/chrome-extension/src/popup.ts b/examples/chrome-extension/src/popup.ts
@@ -11,22 +11,41 @@ import {
   InitProgressReport,
   CreateMLCEngine,
   ChatCompletionMessageParam,
+  prebuiltAppConfig,
 } from "@mlc-ai/web-llm";
 import { ProgressBar, Line } from "progressbar.js";
 
+// modified setLabel to not throw error
+function setLabel(id: string, text: string) {
+  const label = document.getElementById(id);
+  if (label != null) {
+    label.innerText = text;
+  }
+}
+
+function getElementAndCheck(id: string): HTMLElement {
+  const element = document.getElementById(id);
+  if (element == null) {
+    throw Error("Cannot find element " + id);
+  }
+  return element;
+}
+
 const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
 
-const queryInput = document.getElementById("query-input")!;
-const submitButton = document.getElementById("submit-button")!;
+const queryInput = getElementAndCheck("query-input")!;
+const submitButton = getElementAndCheck("submit-button")!;
+const modelName = getElementAndCheck("model-name");
 
 let context = "";
-let isLoadingParams = false;
+let modelDisplayName = "";
 
+// throws runtime.lastError if you refresh extension AND try to access a webpage that is already open
 fetchPageContents();
 
 (<HTMLButtonElement>submitButton).disabled = true;
 
-const progressBar: ProgressBar = new Line("#loadingContainer", {
+let progressBar: ProgressBar = new Line("#loadingContainer", {
   strokeWidth: 4,
   easing: "easeInOut",
   duration: 1400,
@@ -36,8 +55,10 @@ const progressBar: ProgressBar = new Line("#loadingContainer", {
   svgStyle: { width: "100%", height: "100%" },
 });
 
-const initProgressCallback = (report: InitProgressReport) => {
-  console.log(report.text, report.progress);
+let isLoadingParams = true;
+
+let initProgressCallback = (report: InitProgressReport) => {
+  setLabel("init-label", report.text);
   progressBar.animate(report.progress, {
     duration: 50,
   });
@@ -46,29 +67,67 @@ const initProgressCallback = (report: InitProgressReport) => {
   }
 };
 
-// const selectedModel = "TinyLlama-1.1B-Chat-v0.4-q4f16_1-MLC-1k";
-const selectedModel = "Mistral-7B-Instruct-v0.2-q4f16_1-MLC";
+// initially selected model
+let selectedModel = "Qwen2-0.5B-Instruct-q4f16_1-MLC";
+
+// populate model-selection
+const modelSelector = getElementAndCheck(
+  "model-selection",
+) as HTMLSelectElement;
+for (let i = 0; i < prebuiltAppConfig.model_list.length; ++i) {
+  const model = prebuiltAppConfig.model_list[i];
+  const opt = document.createElement("option");
+  opt.value = model.model_id;
+  opt.innerHTML = model.model_id;
+  opt.selected = false;
+
+  // set initial selection as the initially selected model
+  if (model.model_id == selectedModel) {
+    opt.selected = true;
+  }
+
+  modelSelector.appendChild(opt);
+}
+
+modelName.innerText = "Loading initial model...";
 const engine: MLCEngineInterface = await CreateMLCEngine(selectedModel, {
   initProgressCallback: initProgressCallback,
 });
-const chatHistory: ChatCompletionMessageParam[] = [];
+modelName.innerText = "Now chatting with " + modelDisplayName;
 
-isLoadingParams = true;
+let chatHistory: ChatCompletionMessageParam[] = [];
 
 function enableInputs() {
   if (isLoadingParams) {
     sleep(500);
-    (<HTMLButtonElement>submitButton).disabled = false;
-    const loadingBarContainer = document.getElementById("loadingContainer")!;
-    loadingBarContainer.remove();
-    queryInput.focus();
     isLoadingParams = false;
   }
+
+  // remove loading bar and loading bar descriptors, if exists
+  const initLabel = document.getElementById("init-label");
+  initLabel?.remove();
+  const loadingBarContainer = document.getElementById("loadingContainer")!;
+  loadingBarContainer?.remove();
+  queryInput.focus();
+
+  const modelNameArray = selectedModel.split("-");
+  modelDisplayName = modelNameArray[0];
+  let j = 1;
+  while (j < modelNameArray.length && modelNameArray[j][0] != "q") {
+    modelDisplayName = modelDisplayName + "-" + modelNameArray[j];
+    j++;
+  }
 }
 
+let requestInProgress = false;
+
 // Disable submit button if input field is empty
 queryInput.addEventListener("keyup", () => {
-  if ((<HTMLInputElement>queryInput).value === "") {
+  if (
+    (<HTMLInputElement>queryInput).value === "" ||
+    requestInProgress ||
+    isLoadingParams
+  ) {
     (<HTMLButtonElement>submitButton).disabled = true;
   } else {
     (<HTMLButtonElement>submitButton).disabled = false;
@@ -85,6 +144,9 @@ queryInput.addEventListener("keyup", (event) => {
 
 // Listen for clicks on submit button
 async function handleClick() {
+  requestInProgress = true;
+  (<HTMLButtonElement>submitButton).disabled = true;
+
   // Get the message from the input field
   const message = (<HTMLInputElement>queryInput).value;
   console.log("message", message);
@@ -123,9 +185,72 @@ async function handleClick() {
   const response = await engine.getMessage();
   chatHistory.push({ role: "assistant", content: await engine.getMessage() });
   console.log("response", response);
+
+  requestInProgress = false;
+  (<HTMLButtonElement>submitButton).disabled = false;
 }
 submitButton.addEventListener("click", handleClick);
 
+// listen for changes in modelSelector
+async function handleSelectChange() {
+  if (isLoadingParams) {
+    return;
+  }
+
+  modelName.innerText = "";
+
+  const initLabel = document.createElement("p");
+  initLabel.id = "init-label";
+  initLabel.innerText = "Initializing model...";
+  const loadingContainer = document.createElement("div");
+  loadingContainer.id = "loadingContainer";
+
+  const loadingBox = getElementAndCheck("loadingBox");
+  loadingBox.appendChild(initLabel);
+  loadingBox.appendChild(loadingContainer);
+
+  isLoadingParams = true;
+  (<HTMLButtonElement>submitButton).disabled = true;
+
+  if (requestInProgress) {
+    engine.interruptGenerate();
+  }
+  engine.resetChat();
+  chatHistory = [];
+  await engine.unload();
+
+  selectedModel = modelSelector.value;
+
+  progressBar = new Line("#loadingContainer", {
+    strokeWidth: 4,
+    easing: "easeInOut",
+    duration: 1400,
+    color: "#ffd166",
+    trailColor: "#eee",
+    trailWidth: 1,
+    svgStyle: { width: "100%", height: "100%" },
+  });
+
+  initProgressCallback = (report: InitProgressReport) => {
+    setLabel("init-label", report.text);
+    progressBar.animate(report.progress, {
+      duration: 50,
+    });
+    if (report.progress == 1.0) {
+      enableInputs();
+    }
+  };
+
+  engine.setInitProgressCallback(initProgressCallback);
+
+  requestInProgress = true;
+  modelName.innerText = "Reloading with new model...";
+  await engine.reload(selectedModel);
+  requestInProgress = false;
+  modelName.innerText = "Now chatting with " + modelDisplayName;
+}
+modelSelector.addEventListener("change", handleSelectChange);
+
 // Listen for messages from the background script
 chrome.runtime.onMessage.addListener(({ answer, error }) => {
   if (answer) {