feat: Add support for reduced output dimensionality in GoogleGenerati…

…veAIEmbeddings (#544)
davidmigloz · Aug 31, 2024 · d588070 · d588070
1 parent e916ad3
commit d588070
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/packages/langchain_google/lib/src/embeddings/google_ai/google_ai_embeddings.dart b/packages/langchain_google/lib/src/embeddings/google_ai/google_ai_embeddings.dart
@@ -138,7 +138,6 @@ class GoogleGenerativeAIEmbeddings implements Embeddings {
 
   /// The number of dimensions the resulting output embeddings should have.
   /// Only supported in `text-embedding-004` and later models.
-  /// TODO https://github.com/google-gemini/generative-ai-dart/pull/149
   int? dimensions;
 
   /// The maximum number of documents to embed in a single request.
@@ -168,7 +167,7 @@ class GoogleGenerativeAIEmbeddings implements Embeddings {
               Content.text(doc.pageContent),
               taskType: TaskType.retrievalDocument,
               title: doc.metadata[docTitleKey],
-              // outputDimensionality: dimensions, TODO
+              outputDimensionality: dimensions,
             );
           }).toList(growable: false),
         );
@@ -187,7 +186,7 @@ class GoogleGenerativeAIEmbeddings implements Embeddings {
     final data = await _googleAiClient.embedContent(
       Content.text(query),
       taskType: TaskType.retrievalQuery,
-      // outputDimensionality: dimensions, TODO
+      outputDimensionality: dimensions,
     );
     return data.embedding.values;
   }

diff --git a/packages/langchain_google/test/embeddings/google_ai/google_ai_embeddings_test.dart b/packages/langchain_google/test/embeddings/google_ai/google_ai_embeddings_test.dart
@@ -49,8 +49,7 @@ void main() {
       expect(res[1].length, 768);
     });
 
-    // TODO https://github.com/google-gemini/generative-ai-dart/pull/149
-    test('Test shortening embeddings', skip: true, () async {
+    test('Test shortening embeddings', () async {
       embeddings.dimensions = 256;
       final res = await embeddings.embedQuery('Hello world');
       expect(res.length, 256);