@ai-sdk/google-vertex 4.0.22 → 4.0.24

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

@@ -0,0 +1,105 @@
+import { resolve } from '@ai-sdk/provider-utils';
+import { createVertex as createVertexEdge } from './google-vertex-provider-edge';
+import { createVertex as createVertexOriginal } from '../google-vertex-provider';
+import * as edgeAuth from './google-vertex-auth-edge';
+import { describe, beforeEach, afterEach, expect, it, vi } from 'vitest';
+
+// Mock the imported modules
+vi.mock('./google-vertex-auth-edge', () => ({
+  generateAuthToken: vi.fn().mockResolvedValue('mock-auth-token'),
+}));
+
+vi.mock('../google-vertex-provider', () => ({
+  createVertex: vi.fn().mockImplementation(options => ({
+    ...options,
+  })),
+}));
+
+describe('google-vertex-provider-edge', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    delete process.env.GOOGLE_VERTEX_API_KEY;
+  });
+
+  afterEach(() => {
+    delete process.env.GOOGLE_VERTEX_API_KEY;
+  });
+
+  it('default headers function should return auth token', async () => {
+    createVertexEdge({ project: 'test-project' });
+
+    const mockCreateVertex = vi.mocked(createVertexOriginal);
+    const passedOptions = mockCreateVertex.mock.calls[0][0];
+
+    expect(mockCreateVertex).toHaveBeenCalledTimes(1);
+    expect(typeof passedOptions?.headers).toBe('function');
+
+    expect(await resolve(passedOptions?.headers)).toStrictEqual({
+      Authorization: 'Bearer mock-auth-token',
+    });
+  });
+
+  it('should use custom headers in addition to auth token when provided', async () => {
+    createVertexEdge({
+      project: 'test-project',
+      headers: async () => ({
+        'Custom-Header': 'custom-value',
+      }),
+    });
+
+    const mockCreateVertex = vi.mocked(createVertexOriginal);
+    const passedOptions = mockCreateVertex.mock.calls[0][0];
+
+    expect(mockCreateVertex).toHaveBeenCalledTimes(1);
+    expect(typeof passedOptions?.headers).toBe('function');
+    expect(await resolve(passedOptions?.headers)).toEqual({
+      Authorization: 'Bearer mock-auth-token',
+      'Custom-Header': 'custom-value',
+    });
+  });
+
+  it('should use edge auth token generator', async () => {
+    createVertexEdge({ project: 'test-project' });
+
+    const mockCreateVertex = vi.mocked(createVertexOriginal);
+    const passedOptions = mockCreateVertex.mock.calls[0][0];
+
+    // Verify the headers function actually calls generateAuthToken by checking its result
+    expect(passedOptions?.headers).toBeDefined();
+    await resolve(passedOptions?.headers);
+    expect(edgeAuth.generateAuthToken).toHaveBeenCalled();
+  });
+
+  it('passes googleCredentials to generateAuthToken', async () => {
+    createVertexEdge({
+      project: 'test-project',
+      googleCredentials: {
+        clientEmail: 'test@example.com',
+        privateKey: 'test-key',
+      },
+    });
+
+    const mockCreateVertex = vi.mocked(createVertexOriginal);
+    const passedOptions = mockCreateVertex.mock.calls[0][0];
+
+    await resolve(passedOptions?.headers); // call the headers function
+
+    expect(edgeAuth.generateAuthToken).toHaveBeenCalledWith({
+      clientEmail: 'test@example.com',
+      privateKey: 'test-key',
+    });
+  });
+
+  it('should pass options through to base provider when apiKey is provided', async () => {
+    createVertexEdge({
+      apiKey: 'test-api-key',
+    });
+
+    const mockCreateVertex = vi.mocked(createVertexOriginal);
+    const passedOptions = mockCreateVertex.mock.calls[0][0];
+
+    expect(passedOptions?.apiKey).toBe('test-api-key');
+    expect(passedOptions?.headers).toBeUndefined();
+    expect(edgeAuth.generateAuthToken).not.toHaveBeenCalled();
+  });
+});

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

@@ -0,0 +1,50 @@
+import { loadOptionalSetting, resolve } from '@ai-sdk/provider-utils';
+import {
+  createVertex as createVertexOriginal,
+  GoogleVertexProvider,
+  GoogleVertexProviderSettings as GoogleVertexProviderSettingsOriginal,
+} from '../google-vertex-provider';
+import {
+  generateAuthToken,
+  GoogleCredentials,
+} from './google-vertex-auth-edge';
+
+export type { GoogleVertexProvider };
+
+export interface GoogleVertexProviderSettings
+  extends GoogleVertexProviderSettingsOriginal {
+  /**
+   * Optional. The Google credentials for the Google Cloud service account. If
+   * not provided, the Google Vertex provider will use environment variables to
+   * load the credentials.
+   */
+  googleCredentials?: GoogleCredentials;
+}
+
+export function createVertex(
+  options: GoogleVertexProviderSettings = {},
+): GoogleVertexProvider {
+  const apiKey = loadOptionalSetting({
+    settingValue: options.apiKey,
+    environmentVariableName: 'GOOGLE_VERTEX_API_KEY',
+  });
+
+  if (apiKey) {
+    return createVertexOriginal(options);
+  }
+
+  return createVertexOriginal({
+    ...options,
+    headers: async () => ({
+      Authorization: `Bearer ${await generateAuthToken(
+        options.googleCredentials,
+      )}`,
+      ...(await resolve(options.headers)),
+    }),
+  });
+}
+
+/**
+Default Google Vertex AI provider instance.
+ */
+export const vertex = createVertex();

package/src/edge/index.ts

@@ -0,0 +1,5 @@
+export { createVertex, vertex } from './google-vertex-provider-edge';
+export type {
+  GoogleVertexProviderSettings,
+  GoogleVertexProvider,
+} from './google-vertex-provider-edge';

package/src/google-vertex-auth-google-auth-library.test.ts

@@ -0,0 +1,59 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import {
+  generateAuthToken,
+  _resetAuthInstance,
+} from './google-vertex-auth-google-auth-library';
+import { GoogleAuth } from 'google-auth-library';
+
+vi.mock('google-auth-library', () => {
+  return {
+    GoogleAuth: vi.fn().mockImplementation(() => {
+      return {
+        getClient: vi.fn().mockResolvedValue({
+          getAccessToken: vi.fn().mockResolvedValue({ token: 'mocked-token' }),
+        }),
+      };
+    }),
+  };
+});
+
+describe('generateAuthToken', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    _resetAuthInstance();
+  });
+
+  it('should generate a valid auth token', async () => {
+    const token = await generateAuthToken();
+    expect(token).toBe('mocked-token');
+  });
+
+  it('should return null if no token is received', async () => {
+    // Reset the mock completely
+    vi.mocked(GoogleAuth).mockReset();
+
+    // Create a new mock implementation
+    vi.mocked(GoogleAuth).mockImplementation(
+      () =>
+        ({
+          getClient: vi.fn().mockResolvedValue({
+            getAccessToken: vi.fn().mockResolvedValue({ token: null }),
+          }),
+          isGCE: vi.fn(),
+        }) as unknown as GoogleAuth,
+    );
+
+    const token = await generateAuthToken();
+    expect(token).toBeNull();
+  });
+
+  it('should create new auth instance with provided options', async () => {
+    const options = { keyFile: 'test-key.json' };
+    await generateAuthToken(options);
+
+    expect(GoogleAuth).toHaveBeenCalledWith({
+      scopes: ['https://www.googleapis.com/auth/cloud-platform'],
+      keyFile: 'test-key.json',
+    });
+  });
+});

package/src/google-vertex-auth-google-auth-library.ts

@@ -0,0 +1,27 @@
+import { GoogleAuth, GoogleAuthOptions } from 'google-auth-library';
+
+let authInstance: GoogleAuth | null = null;
+let authOptions: GoogleAuthOptions | null = null;
+
+function getAuth(options: GoogleAuthOptions) {
+  if (!authInstance || options !== authOptions) {
+    authInstance = new GoogleAuth({
+      scopes: ['https://www.googleapis.com/auth/cloud-platform'],
+      ...options,
+    });
+    authOptions = options;
+  }
+  return authInstance;
+}
+
+export async function generateAuthToken(options?: GoogleAuthOptions) {
+  const auth = getAuth(options || {});
+  const client = await auth.getClient();
+  const token = await client.getAccessToken();
+  return token?.token || null;
+}
+
+// For testing purposes only
+export function _resetAuthInstance() {
+  authInstance = null;
+}

package/src/google-vertex-config.ts

@@ -0,0 +1,8 @@
+import { FetchFunction, Resolvable } from '@ai-sdk/provider-utils';
+
+export interface GoogleVertexConfig {
+  provider: string;
+  baseURL: string;
+  headers: Resolvable<Record<string, string | undefined>>;
+  fetch?: FetchFunction;
+}

package/src/google-vertex-embedding-model.test.ts

@@ -0,0 +1,315 @@
+import {
+  EmbeddingModelV3Embedding,
+  TooManyEmbeddingValuesForCallError,
+} from '@ai-sdk/provider';
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import { GoogleVertexEmbeddingModel } from './google-vertex-embedding-model';
+import { describe, it, expect, vi } from 'vitest';
+import { createVertex } from './google-vertex-provider';
+
+vi.mock('./version', () => ({
+  VERSION: '0.0.0-test',
+}));
+
+const dummyEmbeddings = [
+  [0.1, 0.2, 0.3],
+  [0.4, 0.5, 0.6],
+];
+const testValues = ['test text one', 'test text two'];
+
+const DEFAULT_URL =
+  'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/textembedding-gecko@001:predict';
+
+const CUSTOM_URL =
+  'https://custom-endpoint.com/models/textembedding-gecko@001:predict';
+
+const server = createTestServer({
+  [DEFAULT_URL]: {},
+  [CUSTOM_URL]: {},
+});
+
+describe('GoogleVertexEmbeddingModel', () => {
+  const mockModelId = 'textembedding-gecko@001';
+  const mockProviderOptions = {
+    outputDimensionality: 768,
+    taskType: 'SEMANTIC_SIMILARITY',
+    title: 'test title',
+    autoTruncate: false,
+  };
+
+  const mockConfig = {
+    provider: 'google-vertex',
+    region: 'us-central1',
+    project: 'test-project',
+    headers: () => ({}),
+    baseURL:
+      'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google',
+  };
+
+  const model = new GoogleVertexEmbeddingModel(mockModelId, mockConfig);
+
+  function prepareJsonResponse({
+    embeddings = dummyEmbeddings,
+    tokenCounts = [1, 1],
+    headers,
+  }: {
+    embeddings?: EmbeddingModelV3Embedding[];
+    tokenCounts?: number[];
+    headers?: Record<string, string>;
+  } = {}) {
+    server.urls[DEFAULT_URL].response = {
+      type: 'json-value',
+      headers,
+      body: {
+        predictions: embeddings.map((values, i) => ({
+          embeddings: {
+            values,
+            statistics: { token_count: tokenCounts[i] },
+          },
+        })),
+      },
+    };
+  }
+
+  it('should extract embeddings', async () => {
+    prepareJsonResponse();
+
+    const { embeddings } = await model.doEmbed({
+      values: testValues,
+      providerOptions: { google: mockProviderOptions },
+    });
+
+    expect(embeddings).toStrictEqual(dummyEmbeddings);
+  });
+
+  it('should expose the raw response', async () => {
+    prepareJsonResponse({
+      headers: {
+        'test-header': 'test-value',
+      },
+    });
+
+    const { response } = await model.doEmbed({
+      values: testValues,
+      providerOptions: { google: mockProviderOptions },
+    });
+
+    expect(response?.headers).toStrictEqual({
+      // default headers:
+      'content-length': '159',
+      'content-type': 'application/json',
+      // custom header
+      'test-header': 'test-value',
+    });
+    expect(response).toMatchSnapshot();
+  });
+
+  it('should extract usage', async () => {
+    prepareJsonResponse({
+      tokenCounts: [10, 15],
+    });
+
+    const { usage } = await model.doEmbed({
+      values: testValues,
+      providerOptions: { google: mockProviderOptions },
+    });
+
+    expect(usage).toStrictEqual({ tokens: 25 });
+  });
+
+  it('should pass the model parameters correctly', async () => {
+    prepareJsonResponse();
+
+    await model.doEmbed({
+      values: testValues,
+      providerOptions: { google: mockProviderOptions },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      instances: testValues.map(value => ({
+        content: value,
+        task_type: mockProviderOptions.taskType,
+        title: mockProviderOptions.title,
+      })),
+      parameters: {
+        outputDimensionality: mockProviderOptions.outputDimensionality,
+        autoTruncate: mockProviderOptions.autoTruncate,
+      },
+    });
+  });
+
+  it('should accept vertex as provider options key', async () => {
+    prepareJsonResponse();
+
+    await model.doEmbed({
+      values: testValues,
+      providerOptions: { vertex: mockProviderOptions },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      instances: testValues.map(value => ({
+        content: value,
+        task_type: mockProviderOptions.taskType,
+        title: mockProviderOptions.title,
+      })),
+      parameters: {
+        outputDimensionality: mockProviderOptions.outputDimensionality,
+        autoTruncate: mockProviderOptions.autoTruncate,
+      },
+    });
+  });
+
+  it('should pass the taskType setting in instances', async () => {
+    prepareJsonResponse();
+
+    await model.doEmbed({
+      values: testValues,
+      providerOptions: { google: { taskType: mockProviderOptions.taskType } },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      instances: testValues.map(value => ({
+        content: value,
+        task_type: mockProviderOptions.taskType,
+      })),
+      parameters: {},
+    });
+  });
+
+  it('should pass the title setting in instances', async () => {
+    prepareJsonResponse();
+
+    await model.doEmbed({
+      values: testValues,
+      providerOptions: { google: { title: mockProviderOptions.title } },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      instances: testValues.map(value => ({
+        content: value,
+        title: mockProviderOptions.title,
+      })),
+      parameters: {},
+    });
+  });
+
+  // changed test to go through the provider `createVertex`
+  it('should pass headers correctly', async () => {
+    prepareJsonResponse();
+
+    const provider = createVertex({
+      project: 'test-project',
+      location: 'us-central1',
+      headers: { 'X-Custom-Header': 'custom-value' },
+    });
+
+    await provider.embeddingModel(mockModelId).doEmbed({
+      values: testValues,
+      headers: { 'X-Request-Header': 'request-value' },
+      providerOptions: { google: mockProviderOptions },
+    });
+
+    expect(server.calls[0].requestHeaders).toStrictEqual({
+      'content-type': 'application/json',
+      'x-custom-header': 'custom-value',
+      'x-request-header': 'request-value',
+    });
+    expect(server.calls[0].requestUserAgent).toContain(
+      `ai-sdk/google-vertex/0.0.0-test`,
+    );
+  });
+
+  it('should throw TooManyEmbeddingValuesForCallError when too many values provided', async () => {
+    const tooManyValues = Array(2049).fill('test');
+
+    await expect(
+      model.doEmbed({
+        values: tooManyValues,
+        providerOptions: { google: mockProviderOptions },
+      }),
+    ).rejects.toThrow(TooManyEmbeddingValuesForCallError);
+  });
+
+  it('should use custom baseURL when provided', async () => {
+    server.urls[CUSTOM_URL].response = {
+      type: 'json-value',
+      body: {
+        predictions: dummyEmbeddings.map(values => ({
+          embeddings: {
+            values,
+            statistics: { token_count: 1 },
+          },
+        })),
+      },
+    };
+
+    const modelWithCustomUrl = new GoogleVertexEmbeddingModel(
+      'textembedding-gecko@001',
+      {
+        headers: () => ({}),
+        baseURL: 'https://custom-endpoint.com',
+        provider: 'google-vertex',
+      },
+    );
+
+    const response = await modelWithCustomUrl.doEmbed({
+      values: testValues,
+      providerOptions: {
+        google: { outputDimensionality: 768 },
+      },
+    });
+
+    expect(response.embeddings).toStrictEqual(dummyEmbeddings);
+
+    expect(server.calls[0].requestUrl).toBe(
+      'https://custom-endpoint.com/models/textembedding-gecko@001:predict',
+    );
+  });
+
+  it('should use custom fetch when provided and include proper request content', async () => {
+    const customFetch = vi.fn().mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          predictions: dummyEmbeddings.map(values => ({
+            embeddings: {
+              values,
+              statistics: { token_count: 1 },
+            },
+          })),
+        }),
+      ),
+    );
+
+    const modelWithCustomFetch = new GoogleVertexEmbeddingModel(
+      'textembedding-gecko@001',
+
+      {
+        headers: () => ({}),
+        baseURL: 'https://custom-endpoint.com',
+        provider: 'google-vertex',
+        fetch: customFetch,
+      },
+    );
+
+    const response = await modelWithCustomFetch.doEmbed({
+      values: testValues,
+      providerOptions: {
+        google: { outputDimensionality: 768 },
+      },
+    });
+
+    expect(response.embeddings).toStrictEqual(dummyEmbeddings);
+
+    expect(customFetch).toHaveBeenCalledWith(CUSTOM_URL, expect.any(Object));
+
+    const [_, secondArgument] = customFetch.mock.calls[0];
+    const requestBody = JSON.parse(secondArgument.body);
+
+    expect(requestBody).toStrictEqual({
+      instances: testValues.map(value => ({ content: value })),
+      parameters: {
+        outputDimensionality: 768,
+      },
+    });
+  });
+});

package/src/google-vertex-embedding-model.ts

@@ -0,0 +1,135 @@
+import {
+  EmbeddingModelV3,
+  TooManyEmbeddingValuesForCallError,
+} from '@ai-sdk/provider';
+import {
+  combineHeaders,
+  createJsonResponseHandler,
+  postJsonToApi,
+  resolve,
+  parseProviderOptions,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import { googleVertexFailedResponseHandler } from './google-vertex-error';
+import {
+  GoogleVertexEmbeddingModelId,
+  googleVertexEmbeddingProviderOptions,
+} from './google-vertex-embedding-options';
+import { GoogleVertexConfig } from './google-vertex-config';
+
+export class GoogleVertexEmbeddingModel implements EmbeddingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly modelId: GoogleVertexEmbeddingModelId;
+  readonly maxEmbeddingsPerCall = 2048;
+  readonly supportsParallelCalls = true;
+
+  private readonly config: GoogleVertexConfig;
+
+  get provider(): string {
+    return this.config.provider;
+  }
+
+  constructor(
+    modelId: GoogleVertexEmbeddingModelId,
+    config: GoogleVertexConfig,
+  ) {
+    this.modelId = modelId;
+    this.config = config;
+  }
+
+  async doEmbed({
+    values,
+    headers,
+    abortSignal,
+    providerOptions,
+  }: Parameters<EmbeddingModelV3['doEmbed']>[0]): Promise<
+    Awaited<ReturnType<EmbeddingModelV3['doEmbed']>>
+  > {
+    let googleOptions = await parseProviderOptions({
+      provider: 'vertex',
+      providerOptions,
+      schema: googleVertexEmbeddingProviderOptions,
+    });
+
+    if (googleOptions == null) {
+      googleOptions = await parseProviderOptions({
+        provider: 'google',
+        providerOptions,
+        schema: googleVertexEmbeddingProviderOptions,
+      });
+    }
+
+    googleOptions = googleOptions ?? {};
+
+    if (values.length > this.maxEmbeddingsPerCall) {
+      throw new TooManyEmbeddingValuesForCallError({
+        provider: this.provider,
+        modelId: this.modelId,
+        maxEmbeddingsPerCall: this.maxEmbeddingsPerCall,
+        values,
+      });
+    }
+
+    const mergedHeaders = combineHeaders(
+      await resolve(this.config.headers),
+      headers,
+    );
+
+    const url = `${this.config.baseURL}/models/${this.modelId}:predict`;
+    const {
+      responseHeaders,
+      value: response,
+      rawValue,
+    } = await postJsonToApi({
+      url,
+      headers: mergedHeaders,
+      body: {
+        instances: values.map(value => ({
+          content: value,
+          task_type: googleOptions.taskType,
+          title: googleOptions.title,
+        })),
+        parameters: {
+          outputDimensionality: googleOptions.outputDimensionality,
+          autoTruncate: googleOptions.autoTruncate,
+        },
+      },
+      failedResponseHandler: googleVertexFailedResponseHandler,
+      successfulResponseHandler: createJsonResponseHandler(
+        googleVertexTextEmbeddingResponseSchema,
+      ),
+      abortSignal,
+      fetch: this.config.fetch,
+    });
+
+    return {
+      warnings: [],
+      embeddings: response.predictions.map(
+        prediction => prediction.embeddings.values,
+      ),
+      usage: {
+        tokens: response.predictions.reduce(
+          (tokenCount, prediction) =>
+            tokenCount + prediction.embeddings.statistics.token_count,
+          0,
+        ),
+      },
+      response: { headers: responseHeaders, body: rawValue },
+    };
+  }
+}
+
+// minimal version of the schema, focussed on what is needed for the implementation
+// this approach limits breakages when the API changes and increases efficiency
+const googleVertexTextEmbeddingResponseSchema = z.object({
+  predictions: z.array(
+    z.object({
+      embeddings: z.object({
+        values: z.array(z.number()),
+        statistics: z.object({
+          token_count: z.number(),
+        }),
+      }),
+    }),
+  ),
+});