From a971ff870ff8f0bda03bfa11b117393a108b577c Mon Sep 17 00:00:00 2001 From: Twest2 Date: Fri, 26 Jun 2026 08:05:18 -0500 Subject: [PATCH 1/3] Added feature to where local ai models can see the image in each step. --- TODO.md | 2 +- app/renderer/dialogs.js | 6 +- app/text-intel.js | 76 +++++++++++++++++++++- core/text-intel.js | 5 ++ docs/getting_started_with_ai.md | 11 +++- tests/unit/text-intel.test.js | 111 +++++++++++++++++++++++++++++--- 6 files changed, 196 insertions(+), 15 deletions(-) diff --git a/TODO.md b/TODO.md index 97ce594..b6feb53 100644 --- a/TODO.md +++ b/TODO.md @@ -1 +1 @@ -In the future, planned updates include a re-vamp of the settings panel, text box's (important, tip, warning, and note) are not properly moving to the intended areas and I will look into locking the size (or being able to control the size) of images throughout varias export formats. +In the future, planned updates include a re-vamp of the settings panel, text box's (important, tip, warning, and note) are not properly moving to the intended areas and I will look into locking the size (or being able to control the size) of images throughout varias export formats. OCR or local ai would be pretty cool.... diff --git a/app/renderer/dialogs.js b/app/renderer/dialogs.js index e92181a..37a4972 100644 --- a/app/renderer/dialogs.js +++ b/app/renderer/dialogs.js @@ -318,7 +318,7 @@ function showSettingsDialog({ const aiAutoDoc = el('input', { type: 'checkbox', checked: Boolean(settings.ai?.autoDoc) }); const ollamaHost = makeInput(settings.ai?.ollama?.host || 'http://127.0.0.1:11434'); const ollamaModel = makeInput(settings.ai?.ollama?.model || 'llama3.2:1b'); - const aiStatus = el('div', { className: 'muted ai-status' }, 'AI stays local through Ollama. Nothing is sent to the cloud.'); + const aiStatus = el('div', { className: 'muted ai-status' }, 'AI stays local through Ollama. Vision-capable models can also inspect the screenshot attached to each step.'); const testAiBtn = el('button', { type: 'button' }, 'Test connection'); const updateAiStatus = (message, { error = false } = {}) => { @@ -341,7 +341,9 @@ function showSettingsDialog({ return; } if (result.installed) { - updateAiStatus(`Connected to ${result.host} with ${result.model}.`); + updateAiStatus(result.vision + ? `Connected to ${result.host} with ${result.model}. It can inspect screenshots.` + : `Connected to ${result.host} with ${result.model}. This model is text-only, so StepForge will use OCR and metadata only.`); } else { updateAiStatus(`Connected to ${result.host}. Model ${result.model} is not installed yet.`, { error: true }); } diff --git a/app/text-intel.js b/app/text-intel.js index 961db39..0388642 100644 --- a/app/text-intel.js +++ b/app/text-intel.js @@ -35,6 +35,15 @@ function clamp(v, min, max) { return Math.min(max, Math.max(min, v)); } +function modelLooksVisionCapable(model) { + const clean = normalizeWhitespace(model).toLowerCase(); + if (!clean) return false; + return clean.includes('vision') + || clean.includes('llava') + || clean.includes('gemma4') + || /(^|[^a-z0-9])qwen[23](?:\.[0-9]+)?vl([^a-z0-9]|$)/.test(clean); +} + let createWorkerImpl = null; function loadCreateWorker() { if (createWorkerImpl) return createWorkerImpl; @@ -64,6 +73,7 @@ class TextIntelService { this.workerPromise = null; this.workerQueue = Promise.resolve(); this.ocrDataDir = path.join(dataDir, 'ocr', 'eng'); + this.modelCapabilityCache = new Map(); } async shutdown() { @@ -372,15 +382,63 @@ public static class Win32 { const data = await res.json(); const models = Array.isArray(data?.models) ? data.models.map((model) => model.name).filter(Boolean) : []; const installed = config.ollama.model ? models.includes(config.ollama.model) : false; + const vision = installed ? await this.modelSupportsVision({ + host: config.ollama.host, + model: config.ollama.model, + }) : false; return { ok: true, installed, + vision, models, host: config.ollama.host, model: config.ollama.model, }; } + async modelCapabilities({ host, model }) { + const normalizedHost = normalizeOllamaHost(host); + const normalizedModel = normalizeWhitespace(model); + if (!normalizedHost || !normalizedModel) return []; + const cacheKey = `${normalizedHost}::${normalizedModel}`; + if (this.modelCapabilityCache.has(cacheKey)) { + return this.modelCapabilityCache.get(cacheKey); + } + const url = new URL('/api/show', `${normalizedHost.replace(/\/+$/, '')}/`); + let capabilities = []; + try { + const response = await this.fetch(url, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ model: normalizedModel }), + }); + if (response.ok) { + const payload = await response.json(); + capabilities = Array.isArray(payload?.capabilities) + ? payload.capabilities.map((cap) => normalizeWhitespace(cap).toLowerCase()).filter(Boolean) + : []; + } + } catch { + capabilities = []; + } + if (!capabilities.includes('vision') && modelLooksVisionCapable(normalizedModel)) { + capabilities = [...capabilities, 'vision']; + } + this.modelCapabilityCache.set(cacheKey, capabilities); + return capabilities; + } + + async modelSupportsVision({ host, model }) { + const capabilities = await this.modelCapabilities({ host, model }); + return capabilities.includes('vision'); + } + + readStepImageBase64(guideId, stepId) { + const imagePath = this.store.stepImagePath(guideId, stepId, 'working') || this.store.stepImagePath(guideId, stepId, 'original'); + if (!imagePath || !fs.existsSync(imagePath)) return ''; + return fs.readFileSync(imagePath).toString('base64'); + } + async callOllamaText({ host, model, prompt, systemPrompt }) { const url = new URL('/api/chat', `${host.replace(/\/+$/, '')}/`); const response = await this.fetch(url, { @@ -403,8 +461,12 @@ public static class Win32 { return content.trim(); } - async callOllama({ host, model, prompt, systemPrompt }) { + async callOllama({ host, model, prompt, systemPrompt, images = [] }) { const url = new URL('/api/chat', `${host.replace(/\/+$/, '')}/`); + const userMessage = { role: 'user', content: prompt }; + if (Array.isArray(images) && images.length) { + userMessage.images = images; + } const response = await this.fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json' }, @@ -414,7 +476,7 @@ public static class Win32 { format: 'json', messages: [ { role: 'system', content: systemPrompt }, - { role: 'user', content: prompt }, + userMessage, ], options: { temperature: 0.2, @@ -460,6 +522,14 @@ public static class Win32 { return { ok: false, reason: 'Block not found.' }; } + const screenshotBase64 = step.image ? this.readStepImageBase64(guideId, stepId) : ''; + const screenshotAttached = Boolean(screenshotBase64) + ? await this.modelSupportsVision({ + host: config.ollama.host, + model: config.ollama.model, + }) + : false; + let captureContext = null; // Use stored capture metadata when available (best context, from capture time). // Fall back to re-running OCR on the stored image only when metadata is absent. @@ -514,6 +584,7 @@ public static class Win32 { step, captureContext, block: currentBlock, + screenshotAttached, }); const raw = await this.callOllama({ @@ -521,6 +592,7 @@ public static class Win32 { model: config.ollama.model, prompt, systemPrompt, + images: screenshotAttached ? [screenshotBase64] : [], }); const patch = normalizeAiPatch(raw); const updated = applyAiPatchToStep(step, patch, { target, blockId }); diff --git a/core/text-intel.js b/core/text-intel.js index 90ac831..e7cf21d 100644 --- a/core/text-intel.js +++ b/core/text-intel.js @@ -673,6 +673,7 @@ function buildAiPrompt({ step = null, captureContext = null, block = null, + screenshotAttached = false, } = {}) { const hasDraftTitle = step && !isPlaceholderTitle(step.title); const hasDraftDesc = step && Boolean(htmlToText(step.descriptionHtml || '')); @@ -717,6 +718,7 @@ function buildAiPrompt({ (!hasDraftTitle || target === 'description') && captureContext.titleCandidate ? `Suggested title: ${captureContext.titleCandidate}` : null, ] : []), + screenshotAttached ? 'Screenshot: attached to this request.' : null, draftTitleLine, draftDescLine, ].filter(Boolean); @@ -770,6 +772,9 @@ function buildAiPrompt({ richContext ? '- Use the OCR text, window title, app name, and element info to make the documentation specific.' : '- Context is limited. Use the app name or window title if available; generate a reasonable action title.', + screenshotAttached + ? '- A screenshot is attached. Use it together with the OCR and metadata to resolve visual details, but do not mention the screenshot in the output.' + : '- No screenshot is attached. Rely on OCR, the window title, app name, and element info.', '- Do NOT generate blocks that describe the technical capture process or mention OCR.', '- Do NOT invent details not supported by the capture context.', '- If the target is one block, only rewrite that block.', diff --git a/docs/getting_started_with_ai.md b/docs/getting_started_with_ai.md index 193218e..fe55b67 100644 --- a/docs/getting_started_with_ai.md +++ b/docs/getting_started_with_ai.md @@ -22,6 +22,14 @@ ollama pull llama3.2:1b That model is small enough to feel responsive on modest hardware, but still good enough for human-sounding titles and short text blocks. +If you want StepForge to send the screenshot itself to the model, pull a vision-capable model instead: + +```bash +ollama pull llama3.2-vision +``` + +That model can inspect pictures as well as text, so it is better when you want the AI to read the UI directly from the screenshot. + If you need something even smaller, try: ```bash @@ -44,7 +52,7 @@ Set: * `Enable AI text filling` to on * `Ollama host` to your local Ollama server -* `Ollama model` to `llama3.2:1b` or the smaller model you pulled +* `Ollama model` to `llama3.2:1b` for text-only mode, or `llama3.2-vision` if you want screenshot-aware AI The default host is: @@ -73,3 +81,4 @@ You can also use `More -> Generate all text fields with AI` to fill the whole st * Capture titles are still generated automatically without AI. * AI generation only works when `Enable AI text filling` is turned on. * The app always uses local OCR around the click area first, then local AI only when you ask for it. +* When the selected Ollama model supports vision, StepForge also sends the screenshot to the model so it can cross-check OCR and visual context. diff --git a/tests/unit/text-intel.test.js b/tests/unit/text-intel.test.js index 86cd3ff..379799b 100644 --- a/tests/unit/text-intel.test.js +++ b/tests/unit/text-intel.test.js @@ -1,5 +1,7 @@ 'use strict'; +const fs = require('node:fs'); +const path = require('node:path'); const test = require('node:test'); const assert = require('node:assert/strict'); @@ -309,21 +311,112 @@ test('ollama connection test reports installed models', async (t) => { settings: makeSettings(), getWindow: () => null, dataDir: root, - fetchImpl: async () => ({ - ok: true, - json: async () => ({ - models: [ - { name: 'llama3.2:1b' }, - { name: 'qwen3:0.6b' }, - ], - }), - }), + fetchImpl: async (url) => { + const pathname = new URL(url).pathname; + if (pathname === '/api/tags') { + return { + ok: true, + json: async () => ({ + models: [ + { name: 'llama3.2:1b' }, + { name: 'qwen3:0.6b' }, + ], + }), + }; + } + if (pathname === '/api/show') { + return { + ok: true, + json: async () => ({ + capabilities: ['completion', 'vision'], + }), + }; + } + throw new Error(`unexpected fetch: ${pathname}`); + }, }); const result = await service.testAiConnection(); assert.equal(result.ok, true); assert.equal(result.installed, true); assert.equal(result.model, 'llama3.2:1b'); + assert.equal(result.vision, true); +}); + +test('vision-capable models receive the screenshot in the chat request', async (t) => { + const root = makeTmpDir('text-intel-ai-vision'); + t.after(() => rmrf(root)); + const imagePath = path.join(root, 'step.png'); + fs.writeFileSync(imagePath, Buffer.from('fake screenshot bytes')); + + const step = createStep({ + title: 'Old title', + descriptionHtml: '

Old text

', + image: { + originalPath: 'original.png', + workingPath: 'working.png', + size: { width: 10, height: 10 }, + }, + captureMetadata: { + windowTitle: 'Settings', + appName: 'chrome', + ocrText: 'Open settings', + titleCandidate: 'Open settings', + mode: 'fullscreen', + }, + }); + const fetchCalls = []; + const service = new TextIntelService({ + store: { + settingsDir: root, + getGuide: () => ({ guideId: 'g1', title: 'Guide', descriptionHtml: '', stepsOrder: ['s1'] }), + getStep: () => step, + stepImagePath: () => imagePath, + saveStep: (_, next) => next, + }, + settings: makeSettings(), + getWindow: () => null, + dataDir: root, + fetchImpl: async (url, init = {}) => { + const pathname = new URL(url).pathname; + fetchCalls.push({ pathname, init }); + if (pathname === '/api/show') { + return { + ok: true, + json: async () => ({ + capabilities: ['completion', 'vision'], + }), + }; + } + if (pathname === '/api/chat') { + return { + ok: true, + json: async () => ({ + message: { + content: JSON.stringify({ + title: 'Open settings', + description: 'Use the AI tab.', + }), + }, + }), + }; + } + throw new Error(`unexpected fetch: ${pathname}`); + }, + }); + + const result = await service.generateStepPatch({ + guideId: 'g1', + stepId: 's1', + target: 'all', + }); + + assert.equal(result.ok, true); + const chatCall = fetchCalls.find((call) => call.pathname === '/api/chat'); + assert.ok(chatCall, 'expected an Ollama chat request'); + const body = JSON.parse(chatCall.init.body); + assert.deepEqual(body.messages[1].images, [fs.readFileSync(imagePath).toString('base64')]); + assert.match(body.messages[1].content, /Screenshot: attached/i); }); test('invalid ollama output fails safely without saving the step', async (t) => { From 79edb6bb56dd0ef59e5d148b28435bedd9b17197 Mon Sep 17 00:00:00 2001 From: Twest2 Date: Fri, 26 Jun 2026 08:36:45 -0500 Subject: [PATCH 2/3] Change model Ran into an issue with the previous model for having mllama archetecture so I changed it --- docs/getting_started_with_ai.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/getting_started_with_ai.md b/docs/getting_started_with_ai.md index fe55b67..817d38d 100644 --- a/docs/getting_started_with_ai.md +++ b/docs/getting_started_with_ai.md @@ -25,7 +25,7 @@ That model is small enough to feel responsive on modest hardware, but still good If you want StepForge to send the screenshot itself to the model, pull a vision-capable model instead: ```bash -ollama pull llama3.2-vision +ollama pull gemma3 ``` That model can inspect pictures as well as text, so it is better when you want the AI to read the UI directly from the screenshot. From 3915865029e6b32293e5f7caf6ddaffc778eefc3 Mon Sep 17 00:00:00 2001 From: Twest2 Date: Fri, 26 Jun 2026 08:48:38 -0500 Subject: [PATCH 3/3] Keep ai model saved when app closes. --- app/renderer/dialogs.js | 7 +++++++ tests/unit/placeholders.test.js | 2 ++ 2 files changed, 9 insertions(+) diff --git a/app/renderer/dialogs.js b/app/renderer/dialogs.js index 37a4972..ababa68 100644 --- a/app/renderer/dialogs.js +++ b/app/renderer/dialogs.js @@ -320,6 +320,11 @@ function showSettingsDialog({ const ollamaModel = makeInput(settings.ai?.ollama?.model || 'llama3.2:1b'); const aiStatus = el('div', { className: 'muted ai-status' }, 'AI stays local through Ollama. Vision-capable models can also inspect the screenshot attached to each step.'); const testAiBtn = el('button', { type: 'button' }, 'Test connection'); + const persistOllamaModel = debounce(() => { + const model = ollamaModel.value.trim(); + // Keep the last model choice even if the dialog is dismissed without a full save. + void api.settings.set({ keyPath: 'ai.ollama.model', value: model }).catch(() => {}); + }, 250); const updateAiStatus = (message, { error = false } = {}) => { aiStatus.textContent = message; @@ -353,6 +358,8 @@ function showSettingsDialog({ setButtonLoading(testAiBtn, false); } }; + ollamaModel.addEventListener('input', () => persistOllamaModel()); + ollamaModel.addEventListener('blur', () => persistOllamaModel.flush()); const placeholderRows = el('div', { className: 'placeholder-rows' }); const rows = []; diff --git a/tests/unit/placeholders.test.js b/tests/unit/placeholders.test.js index 2b15ac3..62c93d6 100644 --- a/tests/unit/placeholders.test.js +++ b/tests/unit/placeholders.test.js @@ -56,12 +56,14 @@ test('settings persist, deep-merge with defaults, and store global placeholders' assert.equal(s1.get('appearance'), DEFAULT_SETTINGS.appearance); s1.set('appearance', 'dark'); s1.set('capture.delayMs', 1500); + s1.set('ai.ollama.model', 'qwen3:0.6b'); s1.setGlobalPlaceholders({ Company: 'Acme', Author: 'Tyler' }); // A fresh instance reads back the changed values merged over defaults. const s2 = new Settings(dir); assert.equal(s2.get('appearance'), 'dark'); assert.equal(s2.get('capture.delayMs'), 1500); + assert.equal(s2.get('ai.ollama.model'), 'qwen3:0.6b'); assert.equal(s2.get('capture.clickMarker'), DEFAULT_SETTINGS.capture.clickMarker); assert.deepEqual(s2.getGlobalPlaceholders(), { Company: 'Acme', Author: 'Tyler' }); });