Implement Google Vertex AI adapter

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<string, string> = { 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<typeof vi.fn>).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<typeof vi.fn>).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/,
+    );
+  });
+});

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<Config>({
   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<string, unknown> = {};
+    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<Config>({
     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;
+  };
+}