From a872f9eec29d323d642c9daf32a82cfb63a5c805 Mon Sep 17 00:00:00 2001 From: Salman Paracha Date: Wed, 2 Jul 2025 14:52:32 -0700 Subject: [PATCH] working version of the extension --- .../PreferenceBasedModelSelector.js | 15 +- .../src/scripts/background.js | 36 --- .../src/scripts/content.js | 253 ++++++++++++++---- .../src/scripts/pageFetchOverride.js | 82 ++---- 4 files changed, 228 insertions(+), 158 deletions(-) delete mode 100644 demos/use_cases/chatgpt-preference-model-selector/src/scripts/background.js diff --git a/demos/use_cases/chatgpt-preference-model-selector/src/components/PreferenceBasedModelSelector.js b/demos/use_cases/chatgpt-preference-model-selector/src/components/PreferenceBasedModelSelector.js index bd8f5684..e13bd867 100644 --- a/demos/use_cases/chatgpt-preference-model-selector/src/components/PreferenceBasedModelSelector.js +++ b/demos/use_cases/chatgpt-preference-model-selector/src/components/PreferenceBasedModelSelector.js @@ -129,17 +129,24 @@ export default function PreferenceBasedModelSelector() { // Save settings: generate name slug and store tuples const handleSave = () => { - const tuples = preferences.map((p) => { - const slug = p.usage.split(/\s+/).slice(0, 3).join('-').toLowerCase() || 'route'; - return { name: slug, usage: p.usage, model: p.model }; - }); + // Only keep valid preferences with non-empty usage + const tuples = preferences + .filter(p => p.usage?.trim()) + .map((p) => { + const slug = p.usage.split(/\s+/).slice(0, 3).join('-').toLowerCase() || 'route'; + return { name: slug, usage: p.usage, model: p.model }; + }); + chrome.storage.sync.set({ routingEnabled, preferences: tuples, defaultModel }, () => { console.log('[PBMS] Saved tuples:', tuples); }); + // Notify content script chrome.tabs.query({ active: true, currentWindow: true }, (tabs) => { chrome.tabs.sendMessage(tabs[0].id, { action: 'applyModelSelection', model: defaultModel }); }); + + // Close the modal window.parent.postMessage({ action: 'CLOSE_PBMS_MODAL' }, '*'); }; diff --git a/demos/use_cases/chatgpt-preference-model-selector/src/scripts/background.js b/demos/use_cases/chatgpt-preference-model-selector/src/scripts/background.js deleted file mode 100644 index 2efa549a..00000000 --- a/demos/use_cases/chatgpt-preference-model-selector/src/scripts/background.js +++ /dev/null @@ -1,36 +0,0 @@ -// background.js -chrome.runtime.onMessage.addListener((msg, sender) => { - if (msg.type === 'ARCHGW_STREAM') { - const { url, init, port } = msg; - - fetch(url, { - method: init.method || 'POST', - headers: init.headers, - body: init.body, - credentials: init.credentials ?? 'same-origin', - // ...copy other init options as needed - }).then(res => { - const reader = res.body.getReader(); - function read() { - reader.read().then(({ done, value }) => { - if (done) { - port.postMessage({ done: true }); - port.close(); - } else { - // Send each chunk (Uint8Array) to the content script - port.postMessage({ chunk: value.buffer }, [value.buffer]); - read(); - } - }); - } - read(); - }).catch(err => { - // In case of error, signal done - port.postMessage({ done: true }); - port.close(); - }); - - // Indicate we’re handling this asynchronously - return true; - } -}); diff --git a/demos/use_cases/chatgpt-preference-model-selector/src/scripts/content.js b/demos/use_cases/chatgpt-preference-model-selector/src/scripts/content.js index 9a34d691..5b202cc1 100644 --- a/demos/use_cases/chatgpt-preference-model-selector/src/scripts/content.js +++ b/demos/use_cases/chatgpt-preference-model-selector/src/scripts/content.js @@ -1,17 +1,107 @@ -// content.js - (() => { const TAG = '[ModelSelector]'; - /**─────────────────────── 0️⃣ Broadcast initial settings ───────────────────────**/ - chrome.storage.sync.get(['preferences','defaultModel'], settings => { - window.postMessage({ type: 'PBMS_SETTINGS', settings }, '*'); - }); - chrome.storage.onChanged.addListener(() => { - chrome.storage.sync.get(['preferences','defaultModel'], settings => { - window.postMessage({ type: 'PBMS_SETTINGS', settings }, '*'); + /**─────────────────────── 🔧 Utility: Scrape Current Messages ───────────────────────**/ + function getMessagesFromDom() { + const bubbles = [...document.querySelectorAll('[data-message-author-role]')]; + + return bubbles + .map(b => { + const role = b.getAttribute('data-message-author-role'); + const content = + role === 'assistant' + ? (b.querySelector('.markdown')?.innerText ?? b.innerText ?? '').trim() + : (b.innerText ?? '').trim(); + return content ? { role, content } : null; + }) + .filter(Boolean); + } + + /**─────────────────────── 🔧 Utility: Prepare Request to Proxy ───────────────────────**/ + function prepareProxyRequest(messages, routes, maxTokenLength = 2048) { + const SYSTEM_PROMPT_TEMPLATE = ` + You are a helpful assistant designed to find the best suited route. + You are provided with route description within XML tags: + + {routes} + + + + {conversation} + + + Your task is to decide which route is best suit with user intent on the conversation in XML tags. Follow the instruction: + 1. If the latest intent from user is irrelevant or user intent is full filled, response with other route {"route": "other"}. + 2. You must analyze the route descriptions and find the best match route for user latest intent. + 3. You only response the name of the route that best matches the user's request, use the exact name in the . + + Based on your analysis, provide your response in the following JSON formats if you decide to match any route: + {"route": "route_name"} + `; + const TOKEN_DIVISOR = 4; + + const filteredMessages = messages.filter( + m => m.role !== 'system' && m.role !== 'tool' && m.content?.trim() + ); + + let tokenCount = SYSTEM_PROMPT_TEMPLATE.length / TOKEN_DIVISOR; + const selected = []; + + for (let i = filteredMessages.length - 1; i >= 0; i--) { + const msg = filteredMessages[i]; + tokenCount += msg.content.length / TOKEN_DIVISOR; + + if (tokenCount > maxTokenLength) { + if (msg.role === 'user') selected.push(msg); + break; + } + + selected.push(msg); + } + + if (selected.length === 0 && filteredMessages.length > 0) { + selected.push(filteredMessages[filteredMessages.length - 1]); + } + + const selectedOrdered = selected.reverse(); + + const systemPrompt = SYSTEM_PROMPT_TEMPLATE + .replace('{routes}', JSON.stringify(routes, null, 2)) + .replace('{conversation}', JSON.stringify(selectedOrdered, null, 2)); + + return systemPrompt; + } + + /**─────────────────────── 🔧 Get routes from storage ───────────────────────**/ + function getRoutesFromStorage() { + return new Promise(resolve => { + chrome.storage.sync.get(['preferences'], ({ preferences }) => { + if (!preferences || !Array.isArray(preferences)) { + console.warn('[ModelSelector] No preferences found in storage'); + return resolve([]); + } + + const routes = preferences.map(p => ({ + name: p.name, + description: p.usage + })); + + resolve(routes); + }); }); - }); + } + + + /**─────────────────────── 🔧 Get model ID by route name ───────────────────────**/ + function getModelIdForRoute(routeName) { + return new Promise(resolve => { + chrome.storage.sync.get(['preferences'], ({ preferences }) => { + const match = (preferences || []).find(p => p.name === routeName); + if (match) resolve(match.model); + else resolve(null); + }); + }); + } /**─────────────────────── 1️⃣ Inject page-context fetch override ───────────────────────**/ (function injectPageFetchOverride() { @@ -25,89 +115,136 @@ (document.head || document.documentElement).appendChild(s); })(); - /**─────────────────────── 2️⃣ Handle proxied fetch from the page ───────────────────────**/ + /**─────────────────────── 2️⃣ Intercept fetch and reroute via Ollama ───────────────────────**/ window.addEventListener('message', ev => { - console.log('[ModelSelector] page→content message', ev.data, ev.ports); - if (ev.source !== window || ev.data?.type !== 'ARCHGW_FETCH') return; - const { url, init, originalRequestUrl } = ev.data; + const { url, init } = ev.data; const port = ev.ports[0]; (async () => { try { - console.log('[ModelSelector] Fetching model recommendation from local proxy...'); - //const res = await fetch(url, init); - //const json = await res.json(); + console.log(`${TAG} Intercepted fetch from page:`, url); - //console.log('[ModelSelector] Proxy responded with:', json); - - const targetModel = 'o4-mini-high'; - if (!targetModel) { - console.warn('[ModelSelector] No model returned from proxy, using default fetch'); - port.postMessage({ done: true }); - return; - } - - // ✅ Extract the original fetch request body from init.body let originalBody = {}; try { originalBody = JSON.parse(init.body); } catch { - console.warn('[ModelSelector] Could not parse original fetch body'); + console.warn(`${TAG} Could not parse original fetch body`); } - // ✅ Patch the model in the request - originalBody.model = targetModel; + const scrapedMessages = getMessagesFromDom(); + const routes = await getRoutesFromStorage(); + const prompt = prepareProxyRequest(scrapedMessages, routes); - console.log(`[ModelSelector] Updating model in request → ${targetModel}`); + // 🔁 Call Ollama router + let selectedRoute = null; + try { + const res = await fetch('http://localhost:11434/api/generate', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model: 'hf.co/katanemo/Arch-Router-1.5B.gguf:Q4_K_M', + prompt: prompt, + temperature: 0.1, + stream: false + }) + }); - // ✅ Resume the request to the real backend - const upstreamRes = await fetch('/backend-api/conversation', { + if (res.ok) { + const data = await res.json(); + console.log(`${TAG} Ollama router response:`, data.response); + try { + let parsed = data.response; + if (typeof data.response === 'string') { + try { + parsed = JSON.parse(data.response); + } catch (jsonErr) { + // Try to recover from single quotes + const safe = data.response.replace(/'/g, '"'); + parsed = JSON.parse(safe); + } + } + selectedRoute = parsed.route || null; + if (!selectedRoute) console.warn(`${TAG} Route missing in parsed response`); + } catch (e) { + console.warn(`${TAG} Failed to parse or extract route from response`, e); + } + } + else { + console.warn(`${TAG} Ollama router failed:`, res.status); + } + } catch (err) { + console.error(`${TAG} Ollama request error`, err); + } + + let targetModel = null; + if (selectedRoute) { + targetModel = await getModelIdForRoute(selectedRoute); + console.log(`${TAG} Resolved model for route "${selectedRoute}" →`, targetModel); + } + + // 🧠 Replace model if we found one + const modifiedBody = { ...originalBody }; + if (targetModel) { + modifiedBody.model = targetModel; + console.log(`${TAG} Overriding request with model: ${targetModel}`); + } else { + console.warn(`${TAG} No route/model override applied`); + } + + const upstreamRes = await fetch(url, { method: init.method, headers: init.headers, credentials: init.credentials, - body: JSON.stringify(originalBody) + body: JSON.stringify(modifiedBody) }); - // ✅ Stream the upstream response back to the page const reader = upstreamRes.body.getReader(); - while (true) { const { done, value } = await reader.read(); if (done) { port.postMessage({ done: true }); break; - } else { - port.postMessage({ chunk: value.buffer }, [value.buffer]); } + port.postMessage({ chunk: value.buffer }, [value.buffer]); } } catch (err) { - console.error('[ModelSelector] proxy fetch error', err); + console.error(`${TAG} Proxy fetch error`, err); port.postMessage({ done: true }); } })(); }); - /**─────────────────────── 3️⃣ DOM patch for model selector label ───────────────────────**/ let desiredModel = null; function patchDom() { - if (desiredModel == null) return; + if (!desiredModel) return; + const btn = document.querySelector('[data-testid="model-switcher-dropdown-button"]'); if (!btn) return; - const span = btn.querySelector('span.text-token-text-tertiary') || btn.querySelector('span'); + + const span = btn.querySelector('div > span'); const wantLabel = `Model selector, current model is ${desiredModel}`; + if (span && span.textContent !== desiredModel) span.textContent = desiredModel; - if (btn.getAttribute('aria-label') !== wantLabel) btn.setAttribute('aria-label', wantLabel); + if (btn.getAttribute('aria-label') !== wantLabel) { + btn.setAttribute('aria-label', wantLabel); + } } + const observer = new MutationObserver(patchDom); observer.observe(document.body || document.documentElement, { subtree: true, childList: true, characterData: true, attributes: true }); + chrome.storage.sync.get(['defaultModel'], ({ defaultModel }) => { - if (defaultModel) { desiredModel = defaultModel; patchDom(); } + if (defaultModel) { + desiredModel = defaultModel; + patchDom(); + } }); + chrome.runtime.onMessage.addListener(msg => { if (msg.action === 'applyModelSelection' && msg.model) { desiredModel = msg.model; @@ -121,31 +258,35 @@ const overlay = document.createElement('div'); overlay.id = 'pbms-overlay'; Object.assign(overlay.style, { - position:'fixed', top:0, left:0, - width:'100vw', height:'100vh', - background:'rgba(0,0,0,0.4)', - display:'flex', alignItems:'center', justifyContent:'center', - zIndex:2147483647 + position: 'fixed', top: 0, left: 0, + width: '100vw', height: '100vh', + background: 'rgba(0,0,0,0.4)', + display: 'flex', alignItems: 'center', justifyContent: 'center', + zIndex: 2147483647 }); const iframe = document.createElement('iframe'); iframe.src = chrome.runtime.getURL('index.html'); - Object.assign(iframe.style,{ - width:'500px', height:'600px', - border:0, borderRadius:'8px', - boxShadow:'0 4px 16px rgba(0,0,0,0.2)', - background:'white', zIndex:2147483648 + Object.assign(iframe.style, { + width: '500px', height: '600px', + border: 0, borderRadius: '8px', + boxShadow: '0 4px 16px rgba(0,0,0,0.2)', + background: 'white', zIndex: 2147483648 }); - overlay.addEventListener('click', e => e.target===overlay && overlay.remove()); + overlay.addEventListener('click', e => e.target === overlay && overlay.remove()); overlay.appendChild(iframe); document.body.appendChild(overlay); } + function interceptDropdown(ev) { if (!ev.target.closest('button[aria-haspopup="menu"]')) return; - ev.preventDefault(); ev.stopPropagation(); + ev.preventDefault(); + ev.stopPropagation(); showModal(); } + document.addEventListener('pointerdown', interceptDropdown, true); document.addEventListener('mousedown', interceptDropdown, true); + window.addEventListener('message', ev => { if (ev.data?.action === 'CLOSE_PBMS_MODAL') { document.getElementById('pbms-overlay')?.remove(); diff --git a/demos/use_cases/chatgpt-preference-model-selector/src/scripts/pageFetchOverride.js b/demos/use_cases/chatgpt-preference-model-selector/src/scripts/pageFetchOverride.js index f27a0ee2..89230cb1 100644 --- a/demos/use_cases/chatgpt-preference-model-selector/src/scripts/pageFetchOverride.js +++ b/demos/use_cases/chatgpt-preference-model-selector/src/scripts/pageFetchOverride.js @@ -1,82 +1,39 @@ -// pageFetchOverride.js (function() { const TAG = '[ModelSelector][Page]'; console.log(`${TAG} installing fetch override`); - window.archgwSettings = window.archgwSettings || { preferences: [], defaultModel: null }; - - window.addEventListener('message', ev => { - if (ev.source === window && ev.data?.type === 'PBMS_SETTINGS') { - window.archgwSettings = ev.data.settings; - console.log(`${TAG} got updated settings`, window.archgwSettings); - } - }); - - // New function: scrape current messages from the DOM - function get_messages() { - const bubbles = [...document.querySelectorAll('[data-message-author-role]')]; - - const messages = bubbles - .map(b => { - const role = b.getAttribute('data-message-author-role'); // "user" | "assistant" - const content = - role === 'assistant' - ? (b.querySelector('.markdown')?.innerText ?? b.innerText ?? '').trim() - : (b.innerText ?? '').trim(); - return content ? { role, content } : null; - }) - .filter(Boolean); - - return { messages }; - } - const origFetch = window.fetch; window.fetch = async function(input, init = {}) { - const urlString = typeof input === 'string' ? input : input.url; - console.log(`${TAG} fetch →`, urlString); - let pathname; - try { - pathname = new URL(urlString).pathname; - } catch { - pathname = urlString; + const urlString = typeof input === 'string' ? input : input.url; + const urlObj = new URL(urlString, window.location.origin); + const pathname = urlObj.pathname; + console.log(`${TAG} fetch →`, pathname); + + const method = (init.method || 'GET').toUpperCase(); + if (method === 'OPTIONS') { + console.log(`${TAG} OPTIONS request → bypassing completely`); + return origFetch(input, init); } + // Only intercept conversation fetches if (pathname === '/backend-api/conversation') { - console.log(`${TAG} matched conversation → proxy via content script`); - - let body = {}; - try { body = JSON.parse(init.body); } catch {} - - const currentMessages = get_messages(); - console.log(`${TAG} scraped messages →`, currentMessages); - - // Patch metadata with current preferences - body.metadata = { - archgw_preference_config: window.archgwSettings.preferences - .map(p => `- name: ${p.name}\n model: ${p.model}\n usage: ${p.usage}`) - .join('\n'), - - // Add current messages dynamically - archgw_current_messages: JSON.stringify(currentMessages) - }; - - init.body = JSON.stringify(body); - - const safeInit = { - method: init.method, - headers: init.headers, - body: init.body, - credentials: init.credentials - }; + console.log(`${TAG} matched → proxy via content script`); const { port1, port2 } = new MessageChannel(); + + // ✅ Remove non-cloneable properties like 'signal' + const safeInit = { ...init }; + delete safeInit.signal; + + // Forward the fetch details to the content script window.postMessage({ type: 'ARCHGW_FETCH', - url: 'http://localhost:12000/v1/chat/completions', + url: urlString, init: safeInit }, '*', [port2]); + // Return a stream response that the content script will fulfill return new Response(new ReadableStream({ start(controller) { port1.onmessage = ({ data }) => { @@ -96,6 +53,7 @@ }); } + // Otherwise, pass through to the original fetch return origFetch(input, init); };