Merge pull request #24 from verisoft-ai/feat/extra_llm_support_vision

y-schwab · web-flow · commit 361e36fa79d9 · 2026-04-15T11:54:50.000-03:00
Feat/extra llm support vision
diff --git a/lib/commands/vision.ts b/lib/commands/vision.ts
@@ -7,6 +7,8 @@ import {
     buildVisionPrompt,
     callVisionLLM,
     computeCoordMapping,
+    getApiKeyEnvVar,
+    getProviderForModel,
     parseVisionCoords,
 } from '../vision-utils';
 
@@ -44,17 +46,18 @@ export async function executeFindByVision(
     this: AppiumDesktopDriver,
     args: { prompt: string; model?: string },
 ): Promise<{ x: number; y: number; label: string }> {
-    const apiKey = process.env.ANTHROPIC_API_KEY;
+    const model = args.model ?? 'claude-opus-4-6';
+    const envVar = getApiKeyEnvVar(getProviderForModel(model));
+    const apiKey = process.env[envVar];
     if (!apiKey) {
         throw new Error(
-            'ANTHROPIC_API_KEY environment variable is required for windows: findByVision'
+            `${envVar} environment variable is required for windows: findByVision (model: ${model})`
         );
     }
 
     const base64 = await this.getScreenshot();
     const { width: ssW, height: ssH } = getPngDimensions(base64);
 
-    const model = args.model ?? 'claude-opus-4-6';
     const raw = await callVisionLLM(base64, buildVisionPrompt(args.prompt, ssW, ssH), model, apiKey);
     const parsed = parseVisionCoords(raw, args.prompt);
 
diff --git a/lib/mcp/tools/vision.ts b/lib/mcp/tools/vision.ts
@@ -10,6 +10,8 @@ import {
     buildVisionPrompt,
     callVisionLLM,
     computeCoordMapping,
+    getApiKeyEnvVar,
+    getProviderForModel,
     parseVisionCoords,
 } from '../../vision-utils';
 
@@ -52,7 +54,8 @@ export function registerVisionTools(server: McpServer, session: AppiumSession):
                 'For "coordinates" format, locates a UI element and returns {x,y,label} with actual screen ' +
                 'coordinates (DPI-corrected) ready to pass to click tools. ' +
                 'For "text" format, answers a general question about the screen in plain text. ' +
-                'Requires ANTHROPIC_API_KEY environment variable.',
+                'Requires ANTHROPIC_API_KEY (Claude), OPENAI_API_KEY (GPT-4o / o-series), or ' +
+                'GEMINI_API_KEY (Gemini) depending on the chosen model.',
             inputSchema: {
                 prompt: z.string().min(1).describe('Question or instruction about the screenshot'),
                 responseFormat: z.enum(['coordinates', 'text']).default('coordinates').describe(
@@ -65,15 +68,18 @@ export function registerVisionTools(server: McpServer, session: AppiumSession):
         },
         async ({ prompt, responseFormat, model }) => {
             try {
-                const apiKey = process.env.ANTHROPIC_API_KEY;
+                const visionModel = model ?? DEFAULT_MODEL;
+                const envVar = getApiKeyEnvVar(getProviderForModel(visionModel));
+                const apiKey = process.env[envVar];
                 if (!apiKey) {
-                    throw new Error('ANTHROPIC_API_KEY environment variable is required for find_by_vision');
+                    throw new Error(
+                        `${envVar} environment variable is required for find_by_vision (model: ${visionModel})`
+                    );
                 }
 
                 const driver = session.getDriver();
                 const base64 = await driver.takeScreenshot() as string;
                 const { width: ssW, height: ssH } = getPngDimensions(base64);
-                const visionModel = model ?? DEFAULT_MODEL;
 
                 if (responseFormat === 'text') {
                     const textPrompt = `Answer the following about this screenshot: "${prompt}"\nRespond with plain text.`;
diff --git a/lib/vision-utils.ts b/lib/vision-utils.ts
@@ -1,5 +1,41 @@
 import Anthropic from '@anthropic-ai/sdk';
 
+export type LLMProvider = 'anthropic' | 'openai' | 'google';
+
+/** Infers the LLM provider from the model identifier. */
+const SUPPORTED_MODELS = [
+    'claude-*  (e.g. claude-sonnet-4-6)',
+    'gpt-*     (e.g. gpt-4o)',
+    'o1, o3, o4, o1-mini, o3-pro, …',
+    'gemini-*  (e.g. gemini-1.5-pro)',
+];
+
+export function getProviderForModel(model: string): LLMProvider {
+    const lower = model.toLowerCase();
+    if (lower.startsWith('gpt-') || /^o\d/.test(lower)) {
+        return 'openai';
+    }
+    if (lower.startsWith('gemini-')) {
+        return 'google';
+    }
+    if (lower.startsWith('claude-')) {
+        return 'anthropic';
+    }
+    throw new Error(
+        `Unsupported model: "${model}". ` +
+        `Supported model prefixes are:\n  ${SUPPORTED_MODELS.join('\n  ')}`,
+    );
+}
+
+/** Returns the environment variable name that holds the API key for the given provider. */
+export function getApiKeyEnvVar(provider: LLMProvider): string {
+    switch (provider) {
+        case 'openai': return 'OPENAI_API_KEY';
+        case 'google': return 'GEMINI_API_KEY';
+        default: return 'ANTHROPIC_API_KEY';
+    }
+}
+
 export interface CoordMapping {
     offsetX: number;
     offsetY: number;
@@ -89,17 +125,12 @@ export function parseVisionCoords(
     return parsed;
 }
 
-/**
- * Sends a base64 screenshot + text prompt to a Claude vision model and returns
- * the raw text response. The caller is responsible for building the prompt and
- * parsing the result.
- */
-export async function callVisionLLM(
+async function callAnthropicVision(
     base64: string,
     textPrompt: string,
     model: string,
     apiKey: string,
-    maxTokens = 256,
+    maxTokens: number,
 ): Promise<string> {
     const client = new Anthropic({ apiKey });
     const response = await client.messages.create({
@@ -118,3 +149,108 @@ export async function callVisionLLM(
     });
     return response.content.find((b) => b.type === 'text')?.text ?? '';
 }
+
+async function callOpenAIVision(
+    base64: string,
+    textPrompt: string,
+    model: string,
+    apiKey: string,
+    maxTokens: number,
+): Promise<string> {
+    const res = await fetch('https://api.openai.com/v1/chat/completions', {
+        method: 'POST',
+        headers: {
+            'Content-Type': 'application/json',
+            'Authorization': `Bearer ${apiKey}`,
+        },
+        body: JSON.stringify({
+            model,
+            max_tokens: maxTokens,
+            messages: [{
+                role: 'user',
+                content: [
+                    { type: 'image_url', image_url: { url: `data:image/png;base64,${base64}` } },
+                    { type: 'text', text: textPrompt },
+                ],
+            }],
+        }),
+    });
+    if (!res.ok) {
+        const body = await res.text();
+        let message: string;
+        try {
+            message = (JSON.parse(body) as { error?: { message: string } }).error?.message ?? body;
+        } catch {
+            message = body || res.statusText;
+        }
+        throw new Error(`OpenAI API error: ${message}`);
+    }
+    const data = await res.json() as { choices?: Array<{ message: { content: string } }> };
+    const content = data.choices?.[0]?.message?.content;
+    if (typeof content !== 'string') {
+        throw new Error(`Unexpected response from OpenAI model "${model}": no text content in choices[0].message.content`);
+    }
+    return content;
+}
+
+async function callGoogleVision(
+    base64: string,
+    textPrompt: string,
+    model: string,
+    apiKey: string,
+    maxTokens: number,
+): Promise<string> {
+    const url = `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent`;
+    const res = await fetch(url, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json', 'x-goog-api-key': apiKey },
+        body: JSON.stringify({
+            contents: [{
+                parts: [
+                    { inline_data: { mime_type: 'image/png', data: base64 } },
+                    { text: textPrompt },
+                ],
+            }],
+            generationConfig: { maxOutputTokens: maxTokens },
+        }),
+    });
+    if (!res.ok) {
+        const body = await res.text();
+        let message: string;
+        try {
+            message = (JSON.parse(body) as { error?: { message: string } }).error?.message ?? body;
+        } catch {
+            message = body || res.statusText;
+        }
+        throw new Error(`Gemini API error: ${message}`);
+    }
+    const data = await res.json() as {
+        candidates?: Array<{ content: { parts: Array<{ text: string }> } }>;
+    };
+    const text = data.candidates?.[0]?.content?.parts?.[0]?.text;
+    if (typeof text !== 'string') {
+        throw new Error(`Unexpected response from Gemini model "${model}": no text in candidates[0].content.parts[0].text`);
+    }
+    return text;
+}
+
+/**
+ * Sends a base64 screenshot + text prompt to a vision model and returns the raw
+ * text response. Dispatches to Anthropic, OpenAI, or Google Gemini based on the
+ * model name prefix. The caller is responsible for building the prompt and
+ * parsing the result.
+ */
+export async function callVisionLLM(
+    base64: string,
+    textPrompt: string,
+    model: string,
+    apiKey: string,
+    maxTokens = 256,
+): Promise<string> {
+    const provider = getProviderForModel(model);
+    switch (provider) {
+        case 'openai': return callOpenAIVision(base64, textPrompt, model, apiKey, maxTokens);
+        case 'google': return callGoogleVision(base64, textPrompt, model, apiKey, maxTokens);
+        default: return callAnthropicVision(base64, textPrompt, model, apiKey, maxTokens);
+    }
+}
diff --git a/test/commands/vision.test.ts b/test/commands/vision.test.ts
@@ -3,8 +3,9 @@
  */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 
-const { mockCreate } = vi.hoisted(() => ({
+const { mockCreate, mockFetch } = vi.hoisted(() => ({
     mockCreate: vi.fn(),
+    mockFetch: vi.fn(),
 }));
 
 vi.mock('@anthropic-ai/sdk', () => ({
@@ -13,6 +14,8 @@ vi.mock('@anthropic-ai/sdk', () => ({
     })),
 }));
 
+vi.stubGlobal('fetch', mockFetch);
+
 vi.mock('../../lib/winapi/user32', () => ({
     getResolutionScalingFactor: vi.fn().mockReturnValue(1.0),
 }));
@@ -55,7 +58,7 @@ describe('executeFindByVision', () => {
         process.env = { ...savedEnv };
     });
 
-    it('throws when ANTHROPIC_API_KEY is not set', async () => {
+    it('throws when ANTHROPIC_API_KEY is not set for default model', async () => {
         delete process.env.ANTHROPIC_API_KEY;
         const driver = makeMockDriver();
 
@@ -64,6 +67,24 @@ describe('executeFindByVision', () => {
         ).rejects.toThrow('ANTHROPIC_API_KEY');
     });
 
+    it('throws when OPENAI_API_KEY is not set for GPT model', async () => {
+        delete process.env.OPENAI_API_KEY;
+        const driver = makeMockDriver();
+
+        await expect(
+            executeFindByVision.call(driver as any, { prompt: 'OK button', model: 'gpt-4o' })
+        ).rejects.toThrow('OPENAI_API_KEY');
+    });
+
+    it('throws when GEMINI_API_KEY is not set for Gemini model', async () => {
+        delete process.env.GEMINI_API_KEY;
+        const driver = makeMockDriver();
+
+        await expect(
+            executeFindByVision.call(driver as any, { prompt: 'OK button', model: 'gemini-2.0-flash' })
+        ).rejects.toThrow('GEMINI_API_KEY');
+    });
+
     it('returns screen coordinates for app session at 100% DPI', async () => {
         // At 100% DPI: rect.width === ssW, dpiScale = 1.0, isLogical = false
         // scaleX = ssW / rect.width = 1920 / 1920 = 1.0
@@ -207,4 +228,90 @@ describe('executeFindByVision', () => {
         expect(result.x).toBe(Math.round(0 + 100 * (2560 / 1920)));
         expect(result.y).toBe(Math.round(0 + 100 * (1440 / 1080)));
     });
+
+    describe('OpenAI provider', () => {
+        beforeEach(() => {
+            process.env.OPENAI_API_KEY = 'openai-test-key';
+        });
+
+        it('calls OpenAI API for gpt-4o model', async () => {
+            const driver = makeMockDriver();
+            mockFetch.mockResolvedValue({
+                ok: true,
+                json: () => Promise.resolve({
+                    choices: [{ message: { content: JSON.stringify({ x: 300, y: 400, label: 'button' }) } }],
+                }),
+            });
+
+            const result = await executeFindByVision.call(driver as any, {
+                prompt: 'button',
+                model: 'gpt-4o',
+            });
+
+            expect(mockFetch).toHaveBeenCalledWith(
+                'https://api.openai.com/v1/chat/completions',
+                expect.objectContaining({
+                    method: 'POST',
+                    headers: expect.objectContaining({ 'Authorization': 'Bearer openai-test-key' }),
+                })
+            );
+            expect(result.label).toBe('button');
+        });
+
+        it('throws on OpenAI API error response', async () => {
+            const driver = makeMockDriver();
+            mockFetch.mockResolvedValue({
+                ok: false,
+                statusText: 'Unauthorized',
+                text: () => Promise.resolve(JSON.stringify({ error: { message: 'Invalid API key' } })),
+            });
+
+            await expect(
+                executeFindByVision.call(driver as any, { prompt: 'button', model: 'gpt-4o-mini' })
+            ).rejects.toThrow('OpenAI API error: Invalid API key');
+        });
+    });
+
+    describe('Google Gemini provider', () => {
+        beforeEach(() => {
+            process.env.GEMINI_API_KEY = 'gemini-test-key';
+        });
+
+        it('calls Gemini API for gemini- model', async () => {
+            const driver = makeMockDriver();
+            mockFetch.mockResolvedValue({
+                ok: true,
+                json: () => Promise.resolve({
+                    candidates: [{ content: { parts: [{ text: JSON.stringify({ x: 200, y: 300, label: 'icon' }) }] } }],
+                }),
+            });
+
+            const result = await executeFindByVision.call(driver as any, {
+                prompt: 'icon',
+                model: 'gemini-2.0-flash',
+            });
+
+            expect(mockFetch).toHaveBeenCalledWith(
+                expect.stringContaining('generativelanguage.googleapis.com'),
+                expect.objectContaining({
+                    method: 'POST',
+                    headers: expect.objectContaining({ 'x-goog-api-key': 'gemini-test-key' }),
+                })
+            );
+            expect(result.label).toBe('icon');
+        });
+
+        it('throws on Gemini API error response', async () => {
+            const driver = makeMockDriver();
+            mockFetch.mockResolvedValue({
+                ok: false,
+                statusText: 'Bad Request',
+                text: () => Promise.resolve(JSON.stringify({ error: { message: 'API key not valid' } })),
+            });
+
+            await expect(
+                executeFindByVision.call(driver as any, { prompt: 'icon', model: 'gemini-1.5-pro' })
+            ).rejects.toThrow('Gemini API error: API key not valid');
+        });
+    });
 });
diff --git a/test/mcp/tools/vision.test.ts b/test/mcp/tools/vision.test.ts
