From 75bf0c5a45704fc3f8278fbd654384955c200858 Mon Sep 17 00:00:00 2001 From: op-simoneromeo <284916543+op-simoneromeo@users.noreply.github.com> Date: Thu, 21 May 2026 02:19:52 +0800 Subject: [PATCH] Implement Google Vertex AI adapter --- packages/ai/google-vertex/src/index.test.ts | 150 ++++++++++++++++++++ packages/ai/google-vertex/src/index.ts | 95 +++++++++++-- 2 files changed, 234 insertions(+), 11 deletions(-) diff --git a/packages/ai/google-vertex/src/index.test.ts b/packages/ai/google-vertex/src/index.test.ts index f43ad207..7b0c20a4 100644 --- a/packages/ai/google-vertex/src/index.test.ts +++ b/packages/ai/google-vertex/src/index.test.ts @@ -1,4 +1,154 @@ import { smokeTest } from '@profullstack/sh1pt-core/testing'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import adapter from './index.js'; smokeTest(adapter, { idPrefix: 'ai' }); + +const ctx = ( + secrets: Record = { GOOGLE_VERTEX_API_KEY: 'test-key' }, + dryRun = false, +) => ({ + secret: (key: string) => secrets[key], + log: () => {}, + dryRun, +}); + +describe('Google Vertex AI generation', () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it('short-circuits dry-run before network calls', async () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const result = await adapter.generate( + ctx({ GOOGLE_VERTEX_API_KEY: 'test-key' }, true), + 'hello', + {}, + {}, + ); + + expect(result).toEqual({ + text: '[dry-run]', + model: 'gemini-1.5-pro', + }); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('posts express-mode generateContent requests and maps usage tokens', async () => { + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + candidates: [ + { + content: { + parts: [{ text: 'hello ' }, { text: 'from vertex' }], + }, + }, + ], + usageMetadata: { + promptTokenCount: 8, + candidatesTokenCount: 5, + totalTokenCount: 13, + }, + }), + }); + vi.stubGlobal('fetch', fetchMock); + + const result = await adapter.generate( + ctx(), + 'hello', + { + model: 'gemini-2.5-flash', + system: 'be concise', + maxTokens: 80, + temperature: 0.2, + extra: { safetySettings: [{ category: 'HARM_CATEGORY_DANGEROUS_CONTENT', threshold: 'BLOCK_ONLY_HIGH' }] }, + }, + {}, + ); + + expect(fetchMock).toHaveBeenCalledOnce(); + const call = fetchMock.mock.calls[0]; + expect(call).toBeDefined(); + const [url, request] = call!; + expect(url).toBe( + 'https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-flash:generateContent?key=test-key', + ); + expect(request.headers['content-type']).toBe('application/json'); + expect(JSON.parse(request.body)).toEqual({ + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + systemInstruction: { parts: [{ text: 'be concise' }] }, + generationConfig: { + maxOutputTokens: 80, + temperature: 0.2, + }, + safetySettings: [{ category: 'HARM_CATEGORY_DANGEROUS_CONTENT', threshold: 'BLOCK_ONLY_HIGH' }], + }); + expect(result).toEqual({ + text: 'hello from vertex', + model: 'gemini-2.5-flash', + inputTokens: 8, + outputTokens: 5, + }); + }); + + it('supports standard project/location model paths and custom base URLs', async () => { + vi.stubGlobal('fetch', vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ candidates: [{ content: { parts: [{ text: 'standard path' }] } }] }), + })); + + const result = await adapter.generate( + ctx({ GOOGLE_VERTEX_API_KEY: 'a key with spaces' }), + 'hello', + { model: 'gemini-2.0-flash' }, + { + baseUrl: 'https://europe-west4-aiplatform.googleapis.com/v1/', + project: 'test-project', + location: 'europe-west4', + }, + ); + + const [url] = (fetch as unknown as ReturnType).mock.calls[0]!; + expect(url).toBe( + 'https://europe-west4-aiplatform.googleapis.com/v1/projects/test-project/locations/europe-west4/publishers/google/models/gemini-2.0-flash:generateContent?key=a%20key%20with%20spaces', + ); + expect(result).toEqual({ + text: 'standard path', + model: 'gemini-2.0-flash', + }); + }); + + it('passes through fully qualified model paths', async () => { + vi.stubGlobal('fetch', vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ candidates: [{ content: { parts: [{ text: 'qualified' }] } }] }), + })); + + await adapter.generate( + ctx(), + 'hello', + { model: 'publishers/google/models/gemini-1.5-pro' }, + {}, + ); + + const [url] = (fetch as unknown as ReturnType).mock.calls[0]!; + expect(url).toBe( + 'https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-1.5-pro:generateContent?key=test-key', + ); + }); + + it('includes status and response body excerpt on errors', async () => { + vi.stubGlobal('fetch', vi.fn().mockResolvedValue({ + ok: false, + status: 403, + text: async () => 'api key rejected'.repeat(30), + })); + + await expect(adapter.generate(ctx(), 'hello', {}, {})).rejects.toThrow( + /Google Vertex AI 403: api key rejected/, + ); + }); +}); diff --git a/packages/ai/google-vertex/src/index.ts b/packages/ai/google-vertex/src/index.ts index ca5fc8ef..960d8fde 100644 --- a/packages/ai/google-vertex/src/index.ts +++ b/packages/ai/google-vertex/src/index.ts @@ -2,29 +2,102 @@ import { defineAi, tokenSetup } from '@profullstack/sh1pt-core'; interface Config { baseUrl?: string; + project?: string; + location?: string; + publisher?: string; } +const DEFAULT_BASE = 'https://aiplatform.googleapis.com/v1'; +const DEFAULT_LOCATION = 'us-central1'; +const DEFAULT_MODEL = 'gemini-1.5-pro'; +const DEFAULT_PUBLISHER = 'google'; + +const trimTrailingSlash = (value: string) => value.replace(/\/+$/, ''); + export default defineAi({ id: 'ai-google-vertex', - label: 'Google Vertex', - defaultModel: 'gemini-1.5-pro', - models: ['gemini-1.5-pro'], + label: 'Google Vertex AI', + defaultModel: DEFAULT_MODEL, + models: [ + DEFAULT_MODEL, + 'gemini-1.5-flash', + 'gemini-2.0-flash', + 'gemini-2.5-flash', + 'gemini-2.5-pro', + ], - async generate(ctx, prompt, _opts, _config) { + async generate(ctx, prompt, opts, config) { const apiKey = ctx.secret('GOOGLE_VERTEX_API_KEY'); - if (!apiKey) throw new Error('GOOGLE_VERTEX_API_KEY not in vault — run `sh1pt promote ai setup`'); - ctx.log(`[stub] ai-google-vertex · ${prompt.length} chars in — integration pending`); - return { text: '[stub — ai-google-vertex integration not yet implemented]', model: 'gemini-1.5-pro' }; + if (!apiKey) throw new Error('GOOGLE_VERTEX_API_KEY not in vault'); + const model = opts.model ?? DEFAULT_MODEL; + ctx.log(`google vertex · model=${model} · ${prompt.length} chars in`); + if (ctx.dryRun) return { text: '[dry-run]', model }; + + const generationConfig: Record = {}; + if (opts.maxTokens !== undefined) generationConfig.maxOutputTokens = opts.maxTokens; + if (opts.temperature !== undefined) generationConfig.temperature = opts.temperature; + + const res = await fetch(buildGenerateContentUrl(apiKey, model, config), { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + contents: [{ role: 'user', parts: [{ text: prompt }] }], + ...(opts.system ? { systemInstruction: { parts: [{ text: opts.system }] } } : {}), + ...(Object.keys(generationConfig).length > 0 ? { generationConfig } : {}), + ...opts.extra, + }), + }); + if (!res.ok) throw new Error(`Google Vertex AI ${res.status}: ${(await res.text()).slice(0, 200)}`); + + const data = await res.json() as VertexGenerateContentResponse; + const text = data.candidates?.[0]?.content?.parts?.map((part) => part.text ?? '').join('') ?? ''; + return { + text, + model, + inputTokens: data.usageMetadata?.promptTokenCount, + outputTokens: data.usageMetadata?.candidatesTokenCount, + }; }, setup: tokenSetup({ secretKey: 'GOOGLE_VERTEX_API_KEY', - label: 'Google Vertex', - vendorDocUrl: 'https://console.cloud.google.com/vertex-ai', + label: 'Google Vertex AI', + vendorDocUrl: 'https://docs.cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/vertex-ai-express-mode-api-reference', steps: [ - 'Sign in at https://console.cloud.google.com/vertex-ai and create an API key', - 'Copy the key — usually shown once', + 'Create a Vertex AI express mode API key', + 'Copy the key - usually shown once', 'Paste below; sh1pt encrypts it in the vault', ], + fields: [ + { key: 'baseUrl', message: 'Vertex AI REST base URL (default: https://aiplatform.googleapis.com/v1):' }, + { key: 'project', message: 'Optional Google Cloud project id for standard Vertex endpoints:' }, + { key: 'location', message: 'Optional Vertex location for standard endpoints (default: us-central1):' }, + { key: 'publisher', message: 'Optional publisher id (default: google):' }, + ], }), }); + +function buildGenerateContentUrl(apiKey: string, model: string, config: Config): string { + const baseUrl = trimTrailingSlash(config.baseUrl ?? DEFAULT_BASE); + const modelPath = model.includes('/') + ? model + : config.project + ? `projects/${config.project}/locations/${config.location ?? DEFAULT_LOCATION}/publishers/${config.publisher ?? DEFAULT_PUBLISHER}/models/${model}` + : `publishers/${config.publisher ?? DEFAULT_PUBLISHER}/models/${model}`; + const separator = baseUrl.includes('?') ? '&' : '?'; + return `${baseUrl}/${modelPath}:generateContent${separator}key=${encodeURIComponent(apiKey)}`; +} + +interface VertexGenerateContentResponse { + candidates?: Array<{ + content?: { + parts?: Array<{ + text?: string; + }>; + }; + }>; + usageMetadata?: { + promptTokenCount?: number; + candidatesTokenCount?: number; + }; +}