@ai-sdk/google-vertex 4.0.23 → 4.0.25

package/src/google-vertex-provider.test.ts

@@ -0,0 +1,318 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { createVertex } from './google-vertex-provider';
+import { GoogleGenerativeAILanguageModel } from '@ai-sdk/google/internal';
+import { GoogleVertexEmbeddingModel } from './google-vertex-embedding-model';
+import { GoogleVertexImageModel } from './google-vertex-image-model';
+
+// Mock the imported modules
+vi.mock('@ai-sdk/provider-utils', () => ({
+  loadSetting: vi.fn().mockImplementation(({ settingValue }) => settingValue),
+  loadOptionalSetting: vi
+    .fn()
+    .mockImplementation(({ settingValue, environmentVariableName }) => {
+      if (settingValue) return settingValue;
+      if (
+        environmentVariableName === 'GOOGLE_VERTEX_API_KEY' &&
+        process.env.GOOGLE_VERTEX_API_KEY
+      ) {
+        return process.env.GOOGLE_VERTEX_API_KEY;
+      }
+      return undefined;
+    }),
+  generateId: vi.fn().mockReturnValue('mock-id'),
+  withoutTrailingSlash: vi.fn().mockImplementation(url => url),
+  resolve: vi.fn().mockImplementation(async value => {
+    if (typeof value === 'function') return value();
+    return value;
+  }),
+  withUserAgentSuffix: vi.fn().mockImplementation(headers => headers),
+}));
+
+vi.mock('@ai-sdk/google/internal', () => ({
+  GoogleGenerativeAILanguageModel: vi.fn(),
+  googleTools: {
+    googleSearch: vi.fn(),
+    urlContext: vi.fn(),
+    fileSearch: vi.fn(),
+    codeExecution: vi.fn(),
+  },
+}));
+
+vi.mock('./google-vertex-embedding-model', () => ({
+  GoogleVertexEmbeddingModel: vi.fn(),
+}));
+
+vi.mock('./google-vertex-image-model', () => ({
+  GoogleVertexImageModel: vi.fn(),
+}));
+
+describe('google-vertex-provider', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    delete process.env.GOOGLE_VERTEX_API_KEY;
+  });
+
+  afterEach(() => {
+    delete process.env.GOOGLE_VERTEX_API_KEY;
+  });
+
+  it('should create a language model with default settings', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      'test-model-id',
+      expect.objectContaining({
+        provider: 'google.vertex.chat',
+        baseURL:
+          'https://test-location-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/test-location/publishers/google',
+        headers: expect.any(Function),
+        generateId: expect.any(Function),
+      }),
+    );
+  });
+
+  it('should throw an error when using new keyword', () => {
+    const provider = createVertex({ project: 'test-project' });
+
+    expect(() => new (provider as any)('test-model-id')).toThrow(
+      'The Google Vertex AI model function cannot be called with the new keyword.',
+    );
+  });
+
+  it('should create an embedding model with correct settings', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+    });
+    provider.embeddingModel('test-embedding-model');
+
+    expect(GoogleVertexEmbeddingModel).toHaveBeenCalledWith(
+      'test-embedding-model',
+      expect.objectContaining({
+        provider: 'google.vertex.embedding',
+        headers: expect.any(Function),
+        baseURL:
+          'https://test-location-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/test-location/publishers/google',
+      }),
+    );
+  });
+
+  it('should pass custom headers to the model constructor', () => {
+    const customHeaders = { 'Custom-Header': 'custom-value' };
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+      headers: customHeaders,
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({
+        headers: expect.any(Function),
+      }),
+    );
+  });
+
+  it('should pass custom generateId function to the model constructor', () => {
+    const customGenerateId = () => 'custom-id';
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+      generateId: customGenerateId,
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({
+        generateId: customGenerateId,
+      }),
+    );
+  });
+
+  it('should use languageModel method to create a model', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+    });
+    provider.languageModel('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      'test-model-id',
+      expect.any(Object),
+    );
+  });
+
+  it('should use custom baseURL when provided', () => {
+    const customBaseURL = 'https://custom-endpoint.example.com';
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+      baseURL: customBaseURL,
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      'test-model-id',
+      expect.objectContaining({
+        baseURL: customBaseURL,
+      }),
+    );
+  });
+
+  it('should create an image model with default settings', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+    });
+    provider.image('imagen-3.0-generate-002');
+
+    expect(GoogleVertexImageModel).toHaveBeenCalledWith(
+      'imagen-3.0-generate-002',
+      expect.objectContaining({
+        provider: 'google.vertex.image',
+        baseURL:
+          'https://test-location-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/test-location/publishers/google',
+        headers: expect.any(Function),
+      }),
+    );
+  });
+
+  it('should use correct URL for global region', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'global',
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      'test-model-id',
+      expect.objectContaining({
+        provider: 'google.vertex.chat',
+        baseURL:
+          'https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/publishers/google',
+        headers: expect.any(Function),
+        generateId: expect.any(Function),
+      }),
+    );
+  });
+
+  it('should use correct URL for global region with embedding model', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'global',
+    });
+    provider.embeddingModel('test-embedding-model');
+
+    expect(GoogleVertexEmbeddingModel).toHaveBeenCalledWith(
+      'test-embedding-model',
+      expect.objectContaining({
+        provider: 'google.vertex.embedding',
+        headers: expect.any(Function),
+        baseURL:
+          'https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/publishers/google',
+      }),
+    );
+  });
+
+  it('should use correct URL for global region with image model', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'global',
+    });
+    provider.image('imagen-3.0-generate-002');
+
+    expect(GoogleVertexImageModel).toHaveBeenCalledWith(
+      'imagen-3.0-generate-002',
+      expect.objectContaining({
+        provider: 'google.vertex.image',
+        baseURL:
+          'https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/publishers/google',
+        headers: expect.any(Function),
+      }),
+    );
+  });
+
+  it('should expose tools', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'test-location',
+    });
+
+    expect(provider.tools).toBeDefined();
+    expect(provider.tools.googleSearch).toBeDefined();
+    expect(provider.tools.urlContext).toBeDefined();
+    expect(provider.tools.codeExecution).toBeDefined();
+  });
+
+  it('should use region-prefixed URL for non-global regions', () => {
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'us-central1',
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      'test-model-id',
+      expect.objectContaining({
+        provider: 'google.vertex.chat',
+        baseURL:
+          'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google',
+        headers: expect.any(Function),
+        generateId: expect.any(Function),
+      }),
+    );
+  });
+
+  it('should use express mode base URL when apiKey is provided', () => {
+    const provider = createVertex({
+      apiKey: 'test-api-key',
+    });
+    provider('test-model-id');
+
+    expect(GoogleGenerativeAILanguageModel).toHaveBeenCalledWith(
+      'test-model-id',
+      expect.objectContaining({
+        baseURL: 'https://aiplatform.googleapis.com/v1/publishers/google',
+      }),
+    );
+  });
+
+  it('should add API key as query parameter via custom fetch', async () => {
+    const provider = createVertex({
+      apiKey: 'test-api-key',
+    });
+    provider('test-model-id');
+
+    const calledConfig = vi.mocked(GoogleGenerativeAILanguageModel).mock
+      .calls[0][1];
+    const customFetch = calledConfig.fetch;
+
+    expect(customFetch).toBeDefined();
+
+    const mockResponse = new Response('{}');
+    const originalFetch = vi.fn().mockResolvedValue(mockResponse);
+    vi.stubGlobal('fetch', originalFetch);
+
+    await customFetch!(
+      'https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-pro:streamGenerateContent',
+      {},
+    );
+
+    expect(originalFetch).toHaveBeenCalledWith(
+      'https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-pro:streamGenerateContent',
+      {
+        headers: {
+          'x-goog-api-key': 'test-api-key',
+        },
+      },
+    );
+
+    vi.unstubAllGlobals();
+  });
+});

package/src/google-vertex-provider.ts

@@ -0,0 +1,217 @@
+import { GoogleGenerativeAILanguageModel } from '@ai-sdk/google/internal';
+import { ImageModelV3, LanguageModelV3, ProviderV3 } from '@ai-sdk/provider';
+import {
+  FetchFunction,
+  generateId,
+  loadOptionalSetting,
+  loadSetting,
+  normalizeHeaders,
+  resolve,
+  Resolvable,
+  withoutTrailingSlash,
+  withUserAgentSuffix,
+} from '@ai-sdk/provider-utils';
+import { VERSION } from './version';
+import { GoogleVertexConfig } from './google-vertex-config';
+import { GoogleVertexEmbeddingModel } from './google-vertex-embedding-model';
+import { GoogleVertexEmbeddingModelId } from './google-vertex-embedding-options';
+import { GoogleVertexImageModel } from './google-vertex-image-model';
+import { GoogleVertexImageModelId } from './google-vertex-image-settings';
+import { GoogleVertexModelId } from './google-vertex-options';
+import { googleVertexTools } from './google-vertex-tools';
+
+const EXPRESS_MODE_BASE_URL =
+  'https://aiplatform.googleapis.com/v1/publishers/google';
+
+// set `x-goog-api-key` header to API key for express mode
+function createExpressModeFetch(
+  apiKey: string,
+  customFetch?: FetchFunction,
+): FetchFunction {
+  return async (url, init) => {
+    const modifiedInit: RequestInit = {
+      ...init,
+      headers: {
+        ...(init?.headers ? normalizeHeaders(init.headers) : {}),
+        'x-goog-api-key': apiKey,
+      },
+    };
+    return (customFetch ?? fetch)(url.toString(), modifiedInit);
+  };
+}
+
+export interface GoogleVertexProvider extends ProviderV3 {
+  /**
+Creates a model for text generation.
+   */
+  (modelId: GoogleVertexModelId): LanguageModelV3;
+
+  languageModel: (modelId: GoogleVertexModelId) => LanguageModelV3;
+
+  /**
+   * Creates a model for image generation.
+   */
+  image(modelId: GoogleVertexImageModelId): ImageModelV3;
+
+  /**
+Creates a model for image generation.
+   */
+  imageModel(modelId: GoogleVertexImageModelId): ImageModelV3;
+
+  tools: typeof googleVertexTools;
+
+  /**
+   * @deprecated Use `embeddingModel` instead.
+   */
+  textEmbeddingModel(
+    modelId: GoogleVertexEmbeddingModelId,
+  ): GoogleVertexEmbeddingModel;
+}
+
+export interface GoogleVertexProviderSettings {
+  /**
+   * Optional. The API key for the Google Cloud project. If provided, the
+   * provider will use express mode with API key authentication. Defaults to
+   * the value of the `GOOGLE_VERTEX_API_KEY` environment variable.
+   */
+  apiKey?: string;
+
+  /**
+Your Google Vertex location. Defaults to the environment variable `GOOGLE_VERTEX_LOCATION`.
+   */
+  location?: string;
+
+  /**
+Your Google Vertex project. Defaults to the environment variable `GOOGLE_VERTEX_PROJECT`.
+  */
+  project?: string;
+
+  /**
+   * Headers to use for requests. Can be:
+   * - A headers object
+   * - A Promise that resolves to a headers object
+   * - A function that returns a headers object
+   * - A function that returns a Promise of a headers object
+   */
+  headers?: Resolvable<Record<string, string | undefined>>;
+
+  /**
+Custom fetch implementation. You can use it as a middleware to intercept requests,
+or to provide a custom fetch implementation for e.g. testing.
+    */
+  fetch?: FetchFunction;
+
+  // for testing
+  generateId?: () => string;
+
+  /**
+Base URL for the Google Vertex API calls.
+     */
+  baseURL?: string;
+}
+
+/**
+Create a Google Vertex AI provider instance.
+ */
+export function createVertex(
+  options: GoogleVertexProviderSettings = {},
+): GoogleVertexProvider {
+  const apiKey = loadOptionalSetting({
+    settingValue: options.apiKey,
+    environmentVariableName: 'GOOGLE_VERTEX_API_KEY',
+  });
+
+  const loadVertexProject = () =>
+    loadSetting({
+      settingValue: options.project,
+      settingName: 'project',
+      environmentVariableName: 'GOOGLE_VERTEX_PROJECT',
+      description: 'Google Vertex project',
+    });
+
+  const loadVertexLocation = () =>
+    loadSetting({
+      settingValue: options.location,
+      settingName: 'location',
+      environmentVariableName: 'GOOGLE_VERTEX_LOCATION',
+      description: 'Google Vertex location',
+    });
+
+  const loadBaseURL = () => {
+    if (apiKey) {
+      return withoutTrailingSlash(options.baseURL) ?? EXPRESS_MODE_BASE_URL;
+    }
+
+    const region = loadVertexLocation();
+    const project = loadVertexProject();
+
+    // For global region, use aiplatform.googleapis.com directly
+    // For other regions, use region-aiplatform.googleapis.com
+    const baseHost = `${region === 'global' ? '' : region + '-'}aiplatform.googleapis.com`;
+
+    return (
+      withoutTrailingSlash(options.baseURL) ??
+      `https://${baseHost}/v1beta1/projects/${project}/locations/${region}/publishers/google`
+    );
+  };
+
+  const createConfig = (name: string): GoogleVertexConfig => {
+    const getHeaders = async () => {
+      const originalHeaders = await resolve(options.headers ?? {});
+      return withUserAgentSuffix(
+        originalHeaders,
+        `ai-sdk/google-vertex/${VERSION}`,
+      );
+    };
+
+    return {
+      provider: `google.vertex.${name}`,
+      headers: getHeaders,
+      fetch: apiKey
+        ? createExpressModeFetch(apiKey, options.fetch)
+        : options.fetch,
+      baseURL: loadBaseURL(),
+    };
+  };
+
+  const createChatModel = (modelId: GoogleVertexModelId) => {
+    return new GoogleGenerativeAILanguageModel(modelId, {
+      ...createConfig('chat'),
+      generateId: options.generateId ?? generateId,
+      supportedUrls: () => ({
+        '*': [
+          // HTTP URLs:
+          /^https?:\/\/.*$/,
+          // Google Cloud Storage URLs:
+          /^gs:\/\/.*$/,
+        ],
+      }),
+    });
+  };
+
+  const createEmbeddingModel = (modelId: GoogleVertexEmbeddingModelId) =>
+    new GoogleVertexEmbeddingModel(modelId, createConfig('embedding'));
+
+  const createImageModel = (modelId: GoogleVertexImageModelId) =>
+    new GoogleVertexImageModel(modelId, createConfig('image'));
+
+  const provider = function (modelId: GoogleVertexModelId) {
+    if (new.target) {
+      throw new Error(
+        'The Google Vertex AI model function cannot be called with the new keyword.',
+      );
+    }
+
+    return createChatModel(modelId);
+  };
+
+  provider.specificationVersion = 'v3' as const;
+  provider.languageModel = createChatModel;
+  provider.embeddingModel = createEmbeddingModel;
+  provider.textEmbeddingModel = createEmbeddingModel;
+  provider.image = createImageModel;
+  provider.imageModel = createImageModel;
+  provider.tools = googleVertexTools;
+
+  return provider;
+}

package/src/google-vertex-tools.ts

@@ -0,0 +1,11 @@
+import { googleTools } from '@ai-sdk/google/internal';
+
+export const googleVertexTools = {
+  googleSearch: googleTools.googleSearch,
+  enterpriseWebSearch: googleTools.enterpriseWebSearch,
+  googleMaps: googleTools.googleMaps,
+  urlContext: googleTools.urlContext,
+  fileSearch: googleTools.fileSearch,
+  codeExecution: googleTools.codeExecution,
+  vertexRagStore: googleTools.vertexRagStore,
+};

package/src/index.ts

@@ -0,0 +1,7 @@
+export type { GoogleVertexImageProviderOptions } from './google-vertex-image-model';
+export { createVertex, vertex } from './google-vertex-provider-node';
+export type {
+  GoogleVertexProvider,
+  GoogleVertexProviderSettings,
+} from './google-vertex-provider-node';
+export { VERSION } from './version';

package/src/version.ts

@@ -0,0 +1,6 @@
+// Version string of this package injected at build time.
+declare const __PACKAGE_VERSION__: string | undefined;
+export const VERSION: string =
+  typeof __PACKAGE_VERSION__ !== 'undefined'
+    ? __PACKAGE_VERSION__
+    : '0.0.0-test';