diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts index 00adfb22f..f42b58c1c 100644 --- a/js/plugins/googleai/src/embedder.ts +++ b/js/plugins/googleai/src/embedder.ts @@ -36,6 +36,14 @@ export const GeminiEmbeddingConfigSchema = z.object({ taskType: TaskTypeSchema.optional(), title: z.string().optional(), version: z.string().optional(), + /** + * The `dimensions` parameter allows you to specify the dimensionality of the embedding output. + * By default, the model generates embeddings with 768 dimensions. Models such as + * `text-embedding-004`, `text-embedding-005`, and `text-multilingual-embedding-002` + * allow the output dimensionality to be adjusted between 1 and 768. + * By selecting a smaller output dimensionality, users can save memory and storage space, leading to more efficient computations. + **/ + dimensions: z.number().min(1).max(768).optional(), }); export type GeminiEmbeddingConfig = z.infer; @@ -69,6 +77,8 @@ export const SUPPORTED_MODELS = { 'text-embedding-004': textEmbedding004, }; +const customDimensionsSupportedModelNames = [textEmbedding004.name]; + export function defineGoogleAIEmbedder( ai: Genkit, name: string, @@ -106,6 +116,11 @@ export function defineGoogleAIEmbedder( info: embedder.info!, }, async (input, options) => { + const dimensions = + customDimensionsSupportedModelNames.includes(embedder.name) && + options?.dimensions + ? options.dimensions + : undefined; const client = new GoogleGenerativeAI(apiKey!).getGenerativeModel({ model: options?.version || @@ -122,6 +137,7 @@ export function defineGoogleAIEmbedder( role: '', parts: [{ text: doc.text }], }, + outputDimensionality: dimensions, } as EmbedContentRequest); const values = response.embedding.values; return { embedding: values }; diff --git a/js/plugins/vertexai/src/embedder.ts b/js/plugins/vertexai/src/embedder.ts index e445a74ca..4b0fa6b2c 100644 --- a/js/plugins/vertexai/src/embedder.ts +++ b/js/plugins/vertexai/src/embedder.ts @@ -165,7 +165,7 @@ export function defineVertexAIEmbedder( title: options?.title, }; }), - dimensions ? { outputDimensionality: dimensions } : {} + { outputDimensionality: dimensions } ); return { embeddings: response.predictions.map((p) => ({ diff --git a/js/plugins/vertexai/src/predict.ts b/js/plugins/vertexai/src/predict.ts index 0d4f39302..991f80933 100644 --- a/js/plugins/vertexai/src/predict.ts +++ b/js/plugins/vertexai/src/predict.ts @@ -33,7 +33,7 @@ interface PredictionResponse { export type PredictClient = ( instances: I[], - parameters?: P + parameters: P ) => Promise>; export function predictModel( @@ -43,7 +43,7 @@ export function predictModel( ): PredictClient { return async ( instances: I[], - parameters?: P + parameters: P ): Promise> => { const fetch = (await import('node-fetch')).default;