diff --git a/.changeset/vertex-embedding-default-options.md b/.changeset/vertex-embedding-default-options.md new file mode 100644 index 000000000000..1979d0003d63 --- /dev/null +++ b/.changeset/vertex-embedding-default-options.md @@ -0,0 +1,5 @@ +--- +'@ai-sdk/google-vertex': patch +--- + +fix(provider/google-vertex): support default embedding model options diff --git a/packages/google-vertex/src/google-vertex-embedding-model.test.ts b/packages/google-vertex/src/google-vertex-embedding-model.test.ts index fb12e7767eb0..023e536f3896 100644 --- a/packages/google-vertex/src/google-vertex-embedding-model.test.ts +++ b/packages/google-vertex/src/google-vertex-embedding-model.test.ts @@ -136,6 +136,67 @@ describe('GoogleVertexEmbeddingModel', () => { `); }); + it('should pass default provider model settings', async () => { + const provider = createGoogleVertex({ + project: 'test-project', + location: 'us-central1', + }); + + await provider + .textEmbeddingModel(mockModelId, { outputDimensionality: 1536 }) + .doEmbed({ + values: testValues, + }); + + expect(await server.calls[0].requestBodyJson).toMatchInlineSnapshot(` + { + "instances": [ + { + "content": "test text one", + }, + { + "content": "test text two", + }, + ], + "parameters": { + "outputDimensionality": 1536, + }, + } + `); + }); + + it('should prefer per-call provider options over default model settings', async () => { + const modelWithDefaultSettings = new GoogleVertexEmbeddingModel( + mockModelId, + mockConfig, + { autoTruncate: true, outputDimensionality: 1536 }, + ); + + await modelWithDefaultSettings.doEmbed({ + values: testValues, + providerOptions: { + googleVertex: { outputDimensionality: 768 }, + }, + }); + + expect(await server.calls[0].requestBodyJson).toMatchInlineSnapshot(` + { + "instances": [ + { + "content": "test text one", + }, + { + "content": "test text two", + }, + ], + "parameters": { + "autoTruncate": true, + "outputDimensionality": 768, + }, + } + `); + }); + it('should accept googleVertex as provider options key', async () => { await model.doEmbed({ values: testValues, diff --git a/packages/google-vertex/src/google-vertex-embedding-model.ts b/packages/google-vertex/src/google-vertex-embedding-model.ts index 61309ad5dc1e..fe6d0f69e497 100644 --- a/packages/google-vertex/src/google-vertex-embedding-model.ts +++ b/packages/google-vertex/src/google-vertex-embedding-model.ts @@ -17,16 +17,22 @@ import { googleVertexFailedResponseHandler } from './google-vertex-error'; import { googleVertexEmbeddingModelOptions, type GoogleVertexEmbeddingModelId, + type GoogleVertexEmbeddingModelOptions, } from './google-vertex-embedding-model-options'; import type { GoogleVertexConfig } from './google-vertex-config'; +interface GoogleVertexEmbeddingConfig extends GoogleVertexConfig { + settings?: GoogleVertexEmbeddingModelOptions; +} + export class GoogleVertexEmbeddingModel implements EmbeddingModelV4 { readonly specificationVersion = 'v4'; readonly modelId: GoogleVertexEmbeddingModelId; readonly maxEmbeddingsPerCall = 2048; readonly supportsParallelCalls = true; - private readonly config: GoogleVertexConfig; + private readonly config: GoogleVertexEmbeddingConfig; + private readonly settings: GoogleVertexEmbeddingModelOptions; static [WORKFLOW_SERIALIZE](model: GoogleVertexEmbeddingModel) { return serializeModelOptions({ @@ -37,7 +43,7 @@ export class GoogleVertexEmbeddingModel implements EmbeddingModelV4 { static [WORKFLOW_DESERIALIZE](options: { modelId: string; - config: GoogleVertexConfig; + config: GoogleVertexEmbeddingConfig; }) { return new GoogleVertexEmbeddingModel(options.modelId, options.config); } @@ -48,10 +54,18 @@ export class GoogleVertexEmbeddingModel implements EmbeddingModelV4 { constructor( modelId: GoogleVertexEmbeddingModelId, - config: GoogleVertexConfig, + config: GoogleVertexEmbeddingConfig, + settings?: GoogleVertexEmbeddingModelOptions, ) { this.modelId = modelId; - this.config = config; + this.config = + settings == null + ? config + : { + ...config, + settings, + }; + this.settings = this.config.settings ?? {}; } async doEmbed({ @@ -84,7 +98,10 @@ export class GoogleVertexEmbeddingModel implements EmbeddingModelV4 { }); } - googleOptions = googleOptions ?? {}; + googleOptions = { + ...this.settings, + ...(googleOptions ?? {}), + }; if (values.length > this.maxEmbeddingsPerCall) { throw new TooManyEmbeddingValuesForCallError({ diff --git a/packages/google-vertex/src/google-vertex-provider-base.test.ts b/packages/google-vertex/src/google-vertex-provider-base.test.ts index 571a9db759e8..2547159bf17d 100644 --- a/packages/google-vertex/src/google-vertex-provider-base.test.ts +++ b/packages/google-vertex/src/google-vertex-provider-base.test.ts @@ -111,6 +111,34 @@ describe('google-vertex-provider-base', () => { ); }); + it('should create embedding models with default model settings', () => { + const provider = createGoogleVertex({ + project: 'test-project', + location: 'test-location', + }); + const settings = { outputDimensionality: 1536 }; + + provider.embeddingModel('test-embedding-model', settings); + provider.textEmbeddingModel('test-text-embedding-model', settings); + + expect(GoogleVertexEmbeddingModel).toHaveBeenNthCalledWith( + 1, + 'test-embedding-model', + expect.objectContaining({ + provider: 'google.vertex.embedding', + }), + settings, + ); + expect(GoogleVertexEmbeddingModel).toHaveBeenNthCalledWith( + 2, + 'test-text-embedding-model', + expect.objectContaining({ + provider: 'google.vertex.embedding', + }), + settings, + ); + }); + it('should pass custom headers to the model constructor', () => { const customHeaders = { 'Custom-Header': 'custom-value' }; const provider = createGoogleVertex({ diff --git a/packages/google-vertex/src/google-vertex-provider-base.ts b/packages/google-vertex/src/google-vertex-provider-base.ts index 306c1468e33d..e3b4b652c2e7 100644 --- a/packages/google-vertex/src/google-vertex-provider-base.ts +++ b/packages/google-vertex/src/google-vertex-provider-base.ts @@ -19,7 +19,10 @@ import { import { VERSION } from './version'; import type { GoogleVertexConfig } from './google-vertex-config'; import { GoogleVertexEmbeddingModel } from './google-vertex-embedding-model'; -import type { GoogleVertexEmbeddingModelId } from './google-vertex-embedding-model-options'; +import type { + GoogleVertexEmbeddingModelId, + GoogleVertexEmbeddingModelOptions, +} from './google-vertex-embedding-model-options'; import { GoogleVertexImageModel } from './google-vertex-image-model'; import type { GoogleVertexImageModelId } from './google-vertex-image-settings'; import type { GoogleVertexModelId } from './google-vertex-options'; @@ -67,11 +70,20 @@ export interface GoogleVertexProvider extends ProviderV4 { tools: typeof googleVertexTools; + /** + * Creates a model for text embeddings. + */ + embeddingModel( + modelId: GoogleVertexEmbeddingModelId, + settings?: GoogleVertexEmbeddingModelOptions, + ): GoogleVertexEmbeddingModel; + /** * @deprecated Use `embeddingModel` instead. */ textEmbeddingModel( modelId: GoogleVertexEmbeddingModelId, + settings?: GoogleVertexEmbeddingModelOptions, ): GoogleVertexEmbeddingModel; /** @@ -206,8 +218,17 @@ export function createGoogleVertex( }); }; - const createEmbeddingModel = (modelId: GoogleVertexEmbeddingModelId) => - new GoogleVertexEmbeddingModel(modelId, createConfig('embedding')); + const createEmbeddingModel = ( + modelId: GoogleVertexEmbeddingModelId, + settings?: GoogleVertexEmbeddingModelOptions, + ) => + settings == null + ? new GoogleVertexEmbeddingModel(modelId, createConfig('embedding')) + : new GoogleVertexEmbeddingModel( + modelId, + createConfig('embedding'), + settings, + ); const createImageModel = (modelId: GoogleVertexImageModelId) => new GoogleVertexImageModel(modelId, {