Add support for stuff QA using context in active tab

examples/chrome-extension/src/background.ts

@@ -1,7 +1,7 @@
 import {ChatRestModule} from "@mlc-ai/web-llm";
 
 // TODO: Surface this as an option to the user 
-const useWebGPU = true;
+const useWebGPU = false;
 
 var cm: ChatRestModule;
 if (!useWebGPU) {
@@ -14,10 +14,19 @@ const generateProgressCallback = (_step: number, message: string) => {
     chrome.runtime.sendMessage({ answer: message });
 };
 
-// listen for a request message from the content script
+var context = "";
 chrome.runtime.onMessage.addListener(async function (request) {
     // check if the request contains a message that the user sent a new message
     if (request.input) {
-        const response = await cm.generate(request.input, generateProgressCallback);
+        var inp = request.input;
+        if (context.length > 0) {
+            inp = "Use only the following context when answering the question at the end. Don't use any other knowledge.\n"+ context + "\n\nQuestion: " + request.input + "\n\nHelpful Answer: ";
+        }
+        console.log("Input:", inp);
+        const response = await cm.generate(inp, generateProgressCallback);
+    }
+    if (request.context) {
+        context = request.context;
+        console.log("Got context:", context);
     }
 });

examples/chrome-extension/src/content.js

@@ -0,0 +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});
+    });
+  });

examples/chrome-extension/src/example.html

@@ -0,0 +1,11 @@
+In the year 2154, humanity had colonized several planets in the distant reaches of the galaxy. The planet of Xylophia-IV was one of the most remote and inhospitable, with temperatures often dropping to -200 degrees Celsius. Despite these harsh conditions, a team of scientists had established a research station on the planet to study the unique geological formations and exotic flora and fauna.
+
+One day, while conducting a routine survey of the planet's surface, the team discovered an strange object buried deep in the ice. As they examined it closer, they realized it was a small, metallic capsule with a glowing blue symbol etched onto its surface.
+
+The team's leader, a brilliant scientist named Dr. Maria Rodriguez, was immediately intrigued by the capsule's mysterious origins. She ordered her team to bring it back to the research station for further analysis.
+
+After weeks of studying the capsule, the team finally cracked the code to the symbol etched onto its surface. It was a message from an alien race, warning Earth of an impending attack from an unknown threat.
+
+The team was shocked and dismayed by the news, but they knew they had to act quickly to warn the rest of humanity. They transmitted the message to the nearest space station, which relayed it to Earth's government.
+
+As the threat of attack loomed near, the team remained on high alert, ready to face whatever dangers lay ahead. They had uncovered a secrets of the universe, and now they were determined to protect their planet and its inhabitants at all costs.

examples/chrome-extension/src/manifest.json

@@ -9,13 +9,24 @@
     "64": "icons/icon-64.png",
     "128": "icons/icon-128.png"
   },
-  "content_security_policy": "style-src-elem 'self' https://cdnjs.cloudflare.com; font-src 'self' https://cdnjs.cloudflare.com; script-src 'self' 'unsafe-eval' 'wasm-unsafe-eval'; default-src 'self' data:; connect-src 'self' data: https://huggingface.co https://cdn-lfs.huggingface.co https://raw.githubusercontent.com",
+  "content_security_policy": "style-src-elem 'self' https://cdnjs.cloudflare.com; font-src 'self' https://cdnjs.cloudflare.com; script-src 'self' 'unsafe-eval' 'wasm-unsafe-eval'; default-src 'self' data:; connect-src 'self' data: http://localhost:8000 https://huggingface.co https://cdn-lfs.huggingface.co https://raw.githubusercontent.com",
   "browser_action": {
     "default_popup": "popup.html"
   },
+  "content_scripts": [
+    {
+      "matches": ["<all_urls>"],
+      "js": ["content.js"]
+    }
+  ],
+  "background": {
+    "scripts": ["background.ts"],
+    "persistent": false
+  },
   "permissions": [
       "storage",
       "tabs",
-      "webNavigation"
+      "webNavigation",
+      "activeTab"
   ]
 }

examples/chrome-extension/src/popup.ts

@@ -8,7 +8,7 @@ import './popup.css';
 import {ChatModule, AppConfig, InitProgressReport} from "@mlc-ai/web-llm";
 
 // TODO: Surface this as an option to the user
-const useWebGPU = true;
+const useWebGPU = false;
 
 var cm: ChatModule;
 const generateProgressCallback = (_step: number, message: string) => {
@@ -110,4 +110,17 @@ function updateAnswer(answer: string) {
     document.getElementById("timestamp")!.innerText = time;
     // Hide loading indicator
     document.getElementById("loading-indicator")!.style.display = "none";
-}
+}
+
+
+// Grab the page contents when the popup is opened
+window.onload = function() {
+    chrome.tabs.query({currentWindow: true,active: true}, function(tabs){
+        var port = chrome.tabs.connect(tabs[0].id,{name: "channelName"});
+        port.postMessage({});
+        port.onMessage.addListener(function(msg) {
+            console.log("Page contents:", msg.contents);
+            chrome.runtime.sendMessage({ context: msg.contents });
+        });
+    });
+}