Add streamDownload function and handle model pull in background.ts

src/background.ts

@@ -1,5 +1,80 @@
+import { getOllamaURL, isOllamaRunning } from "~services/ollama"
+
 export {}
 
+const progressHuman = (completed: number, total: number) => {
+  return ((completed / total) * 100).toFixed(0) + "%"
+}
+
+const clearBadge = () => {
+  chrome.action.setBadgeText({ text: "" })
+  chrome.action.setTitle({ title: "" })
+}
+
+const streamDownload = async (url: string, model: string) => {
+  url += "/api/pull"
+  const response = await fetch(url, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json"
+    },
+    body: JSON.stringify({ model, stream: true })
+  })
+
+  const reader = response.body?.getReader()
+
+  const decoder = new TextDecoder()
+
+  let isSuccess = true
+  while (true) {
+    const { done, value } = await reader.read()
+
+    if (done) {
+      break
+    }
+
+    const text = decoder.decode(value)
+    try {
+      const json = JSON.parse(text.trim()) as {
+        status: string
+        total?: number
+        completed?: number
+      }
+      if (json.total && json.completed) {
+        chrome.action.setBadgeText({
+          text: progressHuman(json.completed, json.total)
+        })
+        chrome.action.setBadgeBackgroundColor({ color: "#0000FF" })
+      } else {
+        chrome.action.setBadgeText({ text: "🏋️‍♂️" })
+        chrome.action.setBadgeBackgroundColor({ color: "#FFFFFF" })
+      }
+
+      chrome.action.setTitle({ title: json.status })
+
+      if (json.status === "success") {
+        isSuccess = true
+      }
+    } catch (e) {
+      console.error(e)
+    }
+  }
+
+  if (isSuccess) {
+    chrome.action.setBadgeText({ text: "✅" })
+    chrome.action.setBadgeBackgroundColor({ color: "#00FF00" })
+    chrome.action.setTitle({ title: "Model pulled successfully" })
+  } else {
+    chrome.action.setBadgeText({ text: "❌" })
+    chrome.action.setBadgeBackgroundColor({ color: "#FF0000" })
+    chrome.action.setTitle({ title: "Model pull failed" })
+  }
+
+  setTimeout(() => {
+    clearBadge()
+  }, 5000)
+}
+
 chrome.runtime.onMessage.addListener(async (message) => {
   if (message.type === "sidepanel") {
     chrome.tabs.query({ active: true, currentWindow: true }, async (tabs) => {
@@ -8,6 +83,22 @@ chrome.runtime.onMessage.addListener(async (message) => {
         tabId: tab.id
       })
     })
+  } else if (message.type === "pull_model") {
+    const ollamaURL = await getOllamaURL()
+
+    const isRunning = await isOllamaRunning()
+
+    if (!isRunning) {
+      chrome.action.setBadgeText({ text: "E" })
+      chrome.action.setBadgeBackgroundColor({ color: "#FF0000" })
+      chrome.action.setTitle({ title: "Ollama is not running" })
+      setTimeout(() => {
+        clearBadge()
+      }, 5000)
+    }
+    console.log("Pulling model", message.modelName)
+
+    await streamDownload(ollamaURL, message.modelName)
   }
 })
 

src/contents/ollama-pull.ts

@@ -0,0 +1,54 @@
+import type { PlasmoCSConfig } from "plasmo"
+
+export const config: PlasmoCSConfig = {
+  matches: ["*://ollama.com/library/*"],
+  all_frames: true
+}
+
+const downloadModel = async (modelName: string) => {
+  const ok = confirm(
+    `[Page Assist Extension] Do you want to pull ${modelName} model? This has nothing to do with Ollama.com website. The model will be pulled locally once you confirm.`
+  )
+  if (ok) {
+    alert(
+      `[Page Assist Extension] Pulling ${modelName} model. For more details, check the extension icon.`
+    )
+
+    await chrome.runtime.sendMessage({
+      type: "pull_model",
+      modelName
+    })
+    return true
+  }
+  return false
+}
+
+const downloadSVG = `<svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" stroke-width="1.5" stroke="currentColor" class="h-5 w-5 pageasssist-icon">
+<path stroke-linecap="round" stroke-linejoin="round" d="M3 16.5v2.25A2.25 2.25 0 0 0 5.25 21h13.5A2.25 2.25 0 0 0 21 18.75V16.5M16.5 12 12 16.5m0 0L7.5 12m4.5 4.5V3" />
+</svg>
+`
+const codeDiv = document.querySelectorAll("div.language-none")
+
+for (let i = 0; i < codeDiv.length; i++) {
+  const button = codeDiv[i].querySelector("button")
+  const command = codeDiv[i].querySelector("input")
+  if (button && command) {
+    const newButton = document.createElement("button")
+    newButton.innerHTML = downloadSVG
+    newButton.className = `border-l ${button.className}`
+    newButton.id = `download-${i}-pageassist`
+    const modelName = command?.value
+      .replace("ollama run", "")
+      .replace("ollama pull", "")
+      .trim()
+    newButton.addEventListener("click", () => {
+      downloadModel(modelName)
+    })
+
+    const span = document.createElement("span")
+    span.title = "Download model via Page Assist"
+    span.appendChild(newButton)
+
+    button.parentNode.appendChild(span)
+  }
+}