wip

Author: r2d4

packages/extension/public/icons/icon128.png

@@ -1,8 +1,8 @@
 {
   "manifest_version": 2,
-  "name": "Thiggle Ext",
+  "name": "LlamaTab",
   "version": "1.0",
-  "description": "This is my awesome Chrome extension.",
+  "description": "Run a large language model in your browser.",
   "background": {
     "scripts": [
       "models/sentencepiece.js",
@@ -13,8 +13,13 @@
   },
   "permissions": ["tabs", "storage"],
   "browser_action": {
-    "default_title": "My Extension",
-    "default_popup": "index.html"
+    "default_title": "LlamaTab",
+    "default_popup": "index.html",
+    "default_icon": {
+      "16": "icons/icon16.png",
+      "48": "icons/icon48.png",
+      "128": "icons/icon128.png"
+    }
   },
   "content_security_policy": "script-src 'self' 'unsafe-eval'; object-src 'self'"
 }

packages/extension/public/icons/icon16.png

@@ -4,22 +4,42 @@ import Page from "./Page";
 
 function App() {
   const [api, setApi] = useState(null);
+  const [loadingStatus, setLoadingStatus] = useState({ progress: 0 });
   useEffect(() => {
     // eslint-disable-next-line no-undef
     chrome.runtime.getBackgroundPage((backgroundPage) => {
       setApi(backgroundPage.API);
+      backgroundPage.API.addInitListener(setLoadingStatus);
     });
+    return () => {
+      // eslint-disable-next-line no-undef
+      chrome.runtime.getBackgroundPage((backgroundPage) => {
+        backgroundPage.API.removeInitListener(setLoadingStatus);
+      });
+    };
   }, []);
+
   if (!api) return null;
   return (
     <div className="w-[300px]">
       <ModelProvider
         config={{
           api,
+          persistToLocalStorage: false,
         }}
       >
         <div>
-          <Page api={api} />
+          <Page
+            loadingStatus={loadingStatus}
+            loadedSystemPrompt={api.systemPrompt}
+            loadedPromptList={api.promptTemplates}
+            setPersistedPromptList={(templates) => {
+              api.promptTemplates = templates;
+            }}
+            setPersistedSystemPrompt={(systemPrompt) => {
+              api.systemPrompt = systemPrompt;
+            }}
+          />
         </div>
       </ModelProvider>
     </div>

packages/extension/public/icons/icon48.png

@@ -1,38 +1,62 @@
 "use client";
 import { useLLM } from "@react-llm/headless";
-import { useEffect, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 import AboutPage from "./pages/AboutPage";
 import LoadingPage from "./pages/LoadingPage";
 import MainPage from "./pages/MainPage";
 import NewPromptPage from "./pages/NewPromptPage";
 import OptionsPage from "./pages/OptionsPage";
 
-export default function Page({ api }) {
-  const { init, send, setOnMessage, systemPrompt, setSystemPrompt } = useLLM();
+export default function Page({
+  loadingStatus,
+  loadedSystemPrompt,
+  loadedPromptList,
+  setPersistedPromptList,
+  setPersistedSystemPrompt,
+}) {
+  const {
+    send,
+    setOnMessage,
+    isGenerating,
+    setConversationPrompt,
+    conversation,
+  } = useLLM();
   const [text, setText] = useState("");
-  const [prompt, setPrompt] = useState("");
-  const [promptList, setPromptList] = useState([]);
   const [response, setResponse] = useState("");
   const [page, setPage] = useState("main");
   const [maxTokens, setMaxTokens] = useState(100);
-  const [loadingStatus, setLoadingStatus] = useState({});
+  const [systemPrompt, setSystemPrompt] = useState(loadedSystemPrompt);
+  const [promptList, setPromptList] = useState(loadedPromptList);
+  const [prompt, setPrompt] = useState(loadedPromptList && loadedPromptList[0]);
 
   useEffect(() => {
-    api.addInitListener(setLoadingStatus);
-    setOnMessage((data) => {
-      setResponse(data);
-    });
     return () => {
-      api.removeInitListener(setLoadingStatus);
+      setPersistedPromptList(promptList);
+      setPersistedSystemPrompt(systemPrompt);
     };
-  }, [api, init, setOnMessage]);
+  });
 
-  const handleSubmit = (e) => {
-    e.preventDefault();
-    send(text, maxTokens);
-  };
+  useEffect(() => {
+    setOnMessage(() => (data) => {
+      console.log(data);
+      setResponse(data);
+    });
+  }, [setOnMessage]);
+
+  useEffect(() => {
+    if (conversation?.systemPrompt !== systemPrompt) {
+      setConversationPrompt(systemPrompt);
+    }
+  }, [systemPrompt, conversation]);
 
-  console.log("progress", loadingStatus);
+  const handleSubmit = useCallback(
+    (e) => {
+      e?.preventDefault();
+      const tmpl = prompt ? prompt.replace(/\$TEXT/g, text) : text;
+      send(tmpl, maxTokens);
+    },
+    [send, prompt, text, maxTokens]
+  );
 
   if (loadingStatus.progress < 1) {
     return <LoadingPage progress={loadingStatus.progress} />;
@@ -51,8 +75,6 @@ export default function Page({ api }) {
           setSystemPrompt={setSystemPrompt}
         />
       );
-    case "load":
-      return <LoadingPage />;
     case "newPrompt":
       return (
         <NewPromptPage
@@ -72,6 +94,7 @@ export default function Page({ api }) {
           response={response}
           text={text}
           setText={setText}
+          isGenerating={isGenerating}
         />
       );
   }

packages/extension/public/manifest.json

@@ -15,13 +15,18 @@ const defaultWorkerConfig = {
   maxWindowSize: 2048,
 };
 
+const defaultSystemPrompt =
+  "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions and always follows the users instructions.";
+
 const API = {
   instance: null,
   initialized: false,
   loadingStatus: {
     progress: 0,
   },
   initListeners: [],
+  systemPrompt: defaultSystemPrompt,
+  promptTemplates: ["$TEXT"],
   addInitListener(cb) {
     this.initListeners.push(cb);
     cb(this.loadingStatus);
@@ -45,6 +50,7 @@ const API = {
     });
   },
   generate(request, cb = console.log) {
+    console.log("generate request: ", request);
     this.instance?.generate(request, (resp) => {
       console.log("generate: ", resp);
       cb(resp);

packages/extension/src/App.js

@@ -1,8 +1,20 @@
 import { useLLM } from "@react-llm/headless";
 
 const LoadingPage = ({ progress }) => {
-  const { init } = useLLM();
+  const { init, gpuDevice } = useLLM();
 
+  if (gpuDevice?.unsupportedReason !== null && gpuDevice?.checked) {
+    return (
+      <div className="flex flex-col gap-2 p-4 text-sm">
+        <p>
+          {JSON.stringify(gpuDevice)}
+          Sorry! LLamaTab is not supported on your device. Reason $
+          {gpuDevice.unsupportedReason}. LLamaTab requires a device with WebGPU
+          support.
+        </p>
+      </div>
+    );
+  }
   if (progress === 0)
     return (
       <div className="flex flex-col gap-2 p-4 text-sm">

packages/extension/src/Page.jsx

@@ -1,3 +1,5 @@
+import React, { useEffect } from "react";
+import TextArea from "./Textarea";
 const MainPage = ({
   prompt,
   setPrompt,
@@ -7,7 +9,24 @@ const MainPage = ({
   response,
   text,
   setText,
+  isGenerating,
 }) => {
+  useEffect(() => {
+    const handleKeyPress = (event) => {
+      console.log(event);
+      if (event.metaKey && event.key === "Enter") {
+        event.preventDefault();
+        handleSubmit();
+      }
+    };
+
+    document.addEventListener("keydown", handleKeyPress);
+
+    return () => {
+      document.removeEventListener("keydown", handleKeyPress);
+    };
+  }, [handleSubmit]);
+
   return (
     <form onSubmit={handleSubmit}>
       <div className="flex flex-col m-auto p-3 text-sm gap-2 relative">
@@ -19,60 +38,64 @@ const MainPage = ({
         </div>
         <div>
           <div>Text</div>
-          <textarea
-            className="w-full rounded-md border border-blue-500 p-2 resize-none"
-            value={text}
-            onChange={(e) => setText(e.target.value)}
-          />
+          <TextArea value={text} onChange={(e) => setText(e.target.value)} />
         </div>
         <div>
           <div>Prompt Template</div>
           <select
-            className="w-full rounded-md border border-blue-500 p-1"
-            value={prompt}
+            className="w-full rounded-md border border-gray-300 p-1"
+            defaultValue={prompt}
             onChange={(e) => setPrompt(e.target.value)}
           >
-            {promptList.map((prompt, idx) => (
-              <option key={idx} value={prompt}>
-                {prompt}
+            {promptList.map((item, idx) => (
+              <option key={idx} value={item}>
+                {item}
               </option>
             ))}
           </select>
         </div>
-        <div className="flex w-full items-center">
-          <div className="flex-grow flex gap-2 flex-col">
+        <div className="flex  w-full items-center">
+          <div className="flex-grow flex gap-2 justify-center m-2">
             <div>
               <button
-                className="text-blue-500 "
+                className="text-blue-500"
                 onClick={() => setPage("newPrompt")}
               >
-                Add New Prompt
+                Manage templates
               </button>
             </div>
-            <div className="flex">
+            <div className="">
               <button
                 className="self-end text-blue-500"
                 onClick={() => setPage("options")}
               >
                 Advanced Options
               </button>
             </div>
+            <div className="self-end flex flex-col justify-center">
+              <button
+                type="submit"
+                className="self-end rounded bg-blue-500 px-2 py-2 font-bold text-white hover:bg-blue-700"
+              >
+                Generate
+              </button>
+              <div className="text-xs text-gray-400 m-1">(⌘+Enter)</div>
+            </div>
           </div>
+        </div>
 
-          <div>
-            <button
-              type="submit"
-              className="self-end rounded bg-blue-500 px-4 py-2 font-bold text-white hover:bg-blue-700"
-            >
-              Generate
-            </button>
+        {isGenerating && (
+          <div className="flex flex-col gap-2 items-center py-3">
+            <div className="self-start">
+              <p className="text-bold">Thinking... </p>
+            </div>
           </div>
-        </div>
+        )}
 
         {response && (
           <div>
-            <div>Response</div>
-            <div className="border h-[50px] rounded-md"></div>
+            <div>Assistant</div>
+            <div className="border rounded-md p-3">{response.outputText}</div>
           </div>
         )}
       </div>

packages/extension/src/background/index.js

@@ -1,4 +1,5 @@
 import React, { useState } from "react";
+import TextArea from "./Textarea";
 const NewPromptPage = ({ setPage, promptList, setPromptList }) => {
   const [newPrompt, setNewPrompt] = useState("");
 
@@ -17,7 +18,7 @@ const NewPromptPage = ({ setPage, promptList, setPromptList }) => {
 
   return (
     <div className="flex flex-col gap-2 p-4 text-sm">
-      <div className="self-end mb-4">
+      <div className="self-end">
         <button
           className="text-blue-500 hover:underline"
           onClick={() => setPage("main")}
@@ -27,12 +28,12 @@ const NewPromptPage = ({ setPage, promptList, setPromptList }) => {
       </div>
 
       <div className="mb-4">
-        <input
-          className="w-full rounded-md border border-blue-500 p-2"
-          type="text"
+        <p className="mb-2">
+          A template for the prompt. $TEXT will be replaced by the input text.
+        </p>
+        <TextArea
           value={newPrompt}
           onChange={(e) => setNewPrompt(e.target.value)}
-          placeholder="Add new prompt"
         />
         <div className="flex justify-end">
           <button