diff --git a/src/background.ts b/src/background.ts index add4d6b..43b8dde 100644 --- a/src/background.ts +++ b/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) } }) diff --git a/src/contents/ollama-pull.ts b/src/contents/ollama-pull.ts new file mode 100644 index 0000000..e8c2336 --- /dev/null +++ b/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 = ` + + +` +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) + } +}