Automatically sort segments for embedding with gemini into batches of 100

Author: jmartisk

model-providers/google/gemini/gemini-common/runtime/src/main/java/io/quarkiverse/langchain4j/gemini/common/GeminiEmbeddingModel.java

@@ -1,5 +1,6 @@
 package io.quarkiverse.langchain4j.gemini.common;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import dev.langchain4j.data.embedding.Embedding;
@@ -9,6 +10,7 @@
 
 public abstract class GeminiEmbeddingModel implements EmbeddingModel {
 
+    private static final int MAX_NUMBER_OF_SEGMENTS_PER_BATCH = 100;
     private final String modelId;
     private final Integer dimension;
     private final String taskType;
@@ -34,15 +36,24 @@ public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
         List<EmbedContentRequest> embedContentRequests = textSegments.stream()
                 .map(textSegment -> getEmbedContentRequest(modelId, textSegment.text()))
                 .toList();
-
-        EmbedContentResponses embedContentResponses = batchEmbedContents(
-                new EmbedContentRequests(embedContentRequests));
-
-        List<Embedding> embeddings = embedContentResponses.embeddings()
-                .stream()
-                .map(embedding -> Embedding.from(embedding.values()))
-                .toList();
-        return Response.from(embeddings);
+        List<Embedding> allEmbeddings = new ArrayList<>();
+        int numberOfEmbeddings = embedContentRequests.size();
+        int numberOfBatches = 1 + numberOfEmbeddings / MAX_NUMBER_OF_SEGMENTS_PER_BATCH;
+
+        for (int i = 0; i < numberOfBatches; i++) {
+            int startIndex = MAX_NUMBER_OF_SEGMENTS_PER_BATCH * i;
+            int lastIndex = Math.min(startIndex + MAX_NUMBER_OF_SEGMENTS_PER_BATCH, numberOfEmbeddings);
+
+            if (startIndex >= numberOfEmbeddings)
+                break;
+
+            EmbedContentResponses embedContentResponses = batchEmbedContents(
+                    new EmbedContentRequests(embedContentRequests.subList(startIndex, lastIndex)));
+            embedContentResponses.embeddings().stream()
+                    .map(embedding -> Embedding.from(embedding.values()))
+                    .forEach(allEmbeddings::add);
+        }
+        return Response.from(allEmbeddings);
     }
 
     private EmbedContentRequest getEmbedContentRequest(String model, String text) {