Skip to content

Commit 1fe4fc6

Browse files
AMKCodeAMKCode
authored
[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
1 parent 0cfb945 commit 1fe4fc6

File tree

7 files changed

+161
-25
lines changed

7 files changed

+161
-25
lines changed

‎examples/chrome-extension/README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,4 @@ npm install
99
npm run build
1010
```
1111

12-
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!
12+
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!

‎examples/chrome-extension/package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "chrome-extension",
3-
"version": "1.0.0",
3+
"version": "1.0.1",
44
"description": "",
55
"private": true,
66
"scripts": {

‎examples/chrome-extension/src/content.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// Only the content script is able to access the DOM
22
chrome.runtime.onConnect.addListener(function (port) {
33
port.onMessage.addListener(function (msg) {
4-
port.postMessage({ contents: document.body.innerHTML });
4+
port.postMessage({ contents: document.body.innerText });
55
});
66
});

‎examples/chrome-extension/src/manifest.json

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"manifest_version": 3,
33
"name": "MLCBot",
4-
"version": "0.1.0",
4+
"version": "0.1.1",
55
"description": "Chat with your browser",
66
"icons": {
77
"16": "icons/icon-16.png",
@@ -16,5 +16,12 @@
1616
"default_title": "MLCBot",
1717
"default_popup": "popup.html"
1818
},
19-
"permissions": ["storage", "tabs", "webNavigation"]
19+
"content_scripts": [
20+
{
21+
"matches": ["<all_urls>"],
22+
"js": ["content.js"]
23+
}
24+
],
25+
"permissions": ["storage", "tabs", "webNavigation", "activeTab", "scripting"],
26+
"host_permissions": ["http://*/", "https://*/"]
2027
}

‎examples/chrome-extension/src/popup.css

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ body {
2121
margin: 0;
2222
padding: 0.5rem;
2323
background-color: #778da9;
24-
width: 320px;
24+
width: 335px;
2525
font-size: small;
2626
}
2727

@@ -32,7 +32,7 @@ p {
3232
/* LOADING BAR */
3333
#loadingContainer {
3434
margin-bottom: 15px;
35-
width: 300px;
35+
width: 315px;
3636
height: 8px;
3737
}
3838

@@ -55,7 +55,7 @@ p {
5555
margin-right: 0.5rem;
5656
}
5757

58-
/* SUBMIT BUTTON */
58+
/* BUTTON */
5959
.btn {
6060
background-color: #1b263b;
6161
color: white;

‎examples/chrome-extension/src/popup.html

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,12 @@
1010
/>
1111
</head>
1212
<body>
13-
<div id="loadingContainer"></div>
14-
13+
<select id="model-selection"></select>
14+
<div id="loadingBox">
15+
<p id="init-label">Initializing model...</p>
16+
<div id="loadingContainer"></div>
17+
</div>
18+
<p id="model-name"></p>
1519
<div class="input-container form-group">
1620
<input
1721
type="search"

‎examples/chrome-extension/src/popup.ts

Lines changed: 140 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,22 +11,41 @@ import {
1111
InitProgressReport,
1212
CreateMLCEngine,
1313
ChatCompletionMessageParam,
14+
prebuiltAppConfig,
1415
} from "@mlc-ai/web-llm";
1516
import { ProgressBar, Line } from "progressbar.js";
1617

18+
// modified setLabel to not throw error
19+
function setLabel(id: string, text: string) {
20+
const label = document.getElementById(id);
21+
if (label != null) {
22+
label.innerText = text;
23+
}
24+
}
25+
26+
function getElementAndCheck(id: string): HTMLElement {
27+
const element = document.getElementById(id);
28+
if (element == null) {
29+
throw Error("Cannot find element " + id);
30+
}
31+
return element;
32+
}
33+
1734
const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
1835

19-
const queryInput = document.getElementById("query-input")!;
20-
const submitButton = document.getElementById("submit-button")!;
36+
const queryInput = getElementAndCheck("query-input")!;
37+
const submitButton = getElementAndCheck("submit-button")!;
38+
const modelName = getElementAndCheck("model-name");
2139

2240
let context = "";
23-
let isLoadingParams = false;
41+
let modelDisplayName = "";
2442

43+
// throws runtime.lastError if you refresh extension AND try to access a webpage that is already open
2544
fetchPageContents();
2645

2746
(<HTMLButtonElement>submitButton).disabled = true;
2847

29-
const progressBar: ProgressBar = new Line("#loadingContainer", {
48+
let progressBar: ProgressBar = new Line("#loadingContainer", {
3049
strokeWidth: 4,
3150
easing: "easeInOut",
3251
duration: 1400,
@@ -36,8 +55,10 @@ const progressBar: ProgressBar = new Line("#loadingContainer", {
3655
svgStyle: { width: "100%", height: "100%" },
3756
});
3857

39-
const initProgressCallback = (report: InitProgressReport) => {
40-
console.log(report.text, report.progress);
58+
let isLoadingParams = true;
59+
60+
let initProgressCallback = (report: InitProgressReport) => {
61+
setLabel("init-label", report.text);
4162
progressBar.animate(report.progress, {
4263
duration: 50,
4364
});
@@ -46,29 +67,67 @@ const initProgressCallback = (report: InitProgressReport) => {
4667
}
4768
};
4869

49-
// const selectedModel = "TinyLlama-1.1B-Chat-v0.4-q4f16_1-MLC-1k";
50-
const selectedModel = "Mistral-7B-Instruct-v0.2-q4f16_1-MLC";
70+
// initially selected model
71+
let selectedModel = "Qwen2-0.5B-Instruct-q4f16_1-MLC";
72+
73+
// populate model-selection
74+
const modelSelector = getElementAndCheck(
75+
"model-selection",
76+
) as HTMLSelectElement;
77+
for (let i = 0; i < prebuiltAppConfig.model_list.length; ++i) {
78+
const model = prebuiltAppConfig.model_list[i];
79+
const opt = document.createElement("option");
80+
opt.value = model.model_id;
81+
opt.innerHTML = model.model_id;
82+
opt.selected = false;
83+
84+
// set initial selection as the initially selected model
85+
if (model.model_id == selectedModel) {
86+
opt.selected = true;
87+
}
88+
89+
modelSelector.appendChild(opt);
90+
}
91+
92+
modelName.innerText = "Loading initial model...";
5193
const engine: MLCEngineInterface = await CreateMLCEngine(selectedModel, {
5294
initProgressCallback: initProgressCallback,
5395
});
54-
const chatHistory: ChatCompletionMessageParam[] = [];
96+
modelName.innerText = "Now chatting with " + modelDisplayName;
5597

56-
isLoadingParams = true;
98+
let chatHistory: ChatCompletionMessageParam[] = [];
5799

58100
function enableInputs() {
59101
if (isLoadingParams) {
60102
sleep(500);
61-
(<HTMLButtonElement>submitButton).disabled = false;
62-
const loadingBarContainer = document.getElementById("loadingContainer")!;
63-
loadingBarContainer.remove();
64-
queryInput.focus();
65103
isLoadingParams = false;
66104
}
105+
106+
// remove loading bar and loading bar descriptors, if exists
107+
const initLabel = document.getElementById("init-label");
108+
initLabel?.remove();
109+
const loadingBarContainer = document.getElementById("loadingContainer")!;
110+
loadingBarContainer?.remove();
111+
queryInput.focus();
112+
113+
const modelNameArray = selectedModel.split("-");
114+
modelDisplayName = modelNameArray[0];
115+
let j = 1;
116+
while (j < modelNameArray.length && modelNameArray[j][0] != "q") {
117+
modelDisplayName = modelDisplayName + "-" + modelNameArray[j];
118+
j++;
119+
}
67120
}
68121

122+
let requestInProgress = false;
123+
69124
// Disable submit button if input field is empty
70125
queryInput.addEventListener("keyup", () => {
71-
if ((<HTMLInputElement>queryInput).value === "") {
126+
if (
127+
(<HTMLInputElement>queryInput).value === "" ||
128+
requestInProgress ||
129+
isLoadingParams
130+
) {
72131
(<HTMLButtonElement>submitButton).disabled = true;
73132
} else {
74133
(<HTMLButtonElement>submitButton).disabled = false;
@@ -85,6 +144,9 @@ queryInput.addEventListener("keyup", (event) => {
85144

86145
// Listen for clicks on submit button
87146
async function handleClick() {
147+
requestInProgress = true;
148+
(<HTMLButtonElement>submitButton).disabled = true;
149+
88150
// Get the message from the input field
89151
const message = (<HTMLInputElement>queryInput).value;
90152
console.log("message", message);
@@ -123,9 +185,72 @@ async function handleClick() {
123185
const response = await engine.getMessage();
124186
chatHistory.push({ role: "assistant", content: await engine.getMessage() });
125187
console.log("response", response);
188+
189+
requestInProgress = false;
190+
(<HTMLButtonElement>submitButton).disabled = false;
126191
}
127192
submitButton.addEventListener("click", handleClick);
128193

194+
// listen for changes in modelSelector
195+
async function handleSelectChange() {
196+
if (isLoadingParams) {
197+
return;
198+
}
199+
200+
modelName.innerText = "";
201+
202+
const initLabel = document.createElement("p");
203+
initLabel.id = "init-label";
204+
initLabel.innerText = "Initializing model...";
205+
const loadingContainer = document.createElement("div");
206+
loadingContainer.id = "loadingContainer";
207+
208+
const loadingBox = getElementAndCheck("loadingBox");
209+
loadingBox.appendChild(initLabel);
210+
loadingBox.appendChild(loadingContainer);
211+
212+
isLoadingParams = true;
213+
(<HTMLButtonElement>submitButton).disabled = true;
214+
215+
if (requestInProgress) {
216+
engine.interruptGenerate();
217+
}
218+
engine.resetChat();
219+
chatHistory = [];
220+
await engine.unload();
221+
222+
selectedModel = modelSelector.value;
223+
224+
progressBar = new Line("#loadingContainer", {
225+
strokeWidth: 4,
226+
easing: "easeInOut",
227+
duration: 1400,
228+
color: "#ffd166",
229+
trailColor: "#eee",
230+
trailWidth: 1,
231+
svgStyle: { width: "100%", height: "100%" },
232+
});
233+
234+
initProgressCallback = (report: InitProgressReport) => {
235+
setLabel("init-label", report.text);
236+
progressBar.animate(report.progress, {
237+
duration: 50,
238+
});
239+
if (report.progress == 1.0) {
240+
enableInputs();
241+
}
242+
};
243+
244+
engine.setInitProgressCallback(initProgressCallback);
245+
246+
requestInProgress = true;
247+
modelName.innerText = "Reloading with new model...";
248+
await engine.reload(selectedModel);
249+
requestInProgress = false;
250+
modelName.innerText = "Now chatting with " + modelDisplayName;
251+
}
252+
modelSelector.addEventListener("change", handleSelectChange);
253+
129254
// Listen for messages from the background script
130255
chrome.runtime.onMessage.addListener(({ answer, error }) => {
131256
if (answer) {

0 commit comments

Comments
 (0)