embedContents

Author: michael-ameri

gemini-api/src/main/java/swiss/ameri/gemini/api/GenAi.java

@@ -242,8 +242,78 @@ public CompletableFuture<GeneratedContent> generateContent(GenerativeModel model
         });
     }
 
+    /**
+     * Embedding is a technique used to represent information as a list of floating point numbers in an array.
+     * With Gemini, you can represent text (words, sentences, and blocks of text) in a vectorized form,
+     * making it easier to compare and contrast embeddings.
+     * For example, two texts that share a similar subject or sentiment should have similar embeddings,
+     * which can be identified through mathematical comparison techniques such as cosine similarity.
+     *
+     * @param model                to use. Currently, only {@link ModelVariant#TEXT_EMBEDDING_004} is allowed.
+     * @param taskType             Optional. Optional task type for which the embeddings will be used. For possible values, see {@link TaskType}
+     * @param title                Optional. An optional title for the text. Only applicable when TaskType is RETRIEVAL_DOCUMENT.
+     *                             Note: Specifying a title for RETRIEVAL_DOCUMENT provides better quality embeddings for retrieval.
+     * @param outputDimensionality Optional. Optional reduced dimension for the output embedding.
+     *                             If set, excessive values in the output embedding are truncated from the end.
+     *                             Supported by newer models since 2024, and the earlier model (models/embedding-001) cannot specify this value.
+     * @return List of values
+     * @apiNote Only {@link swiss.ameri.gemini.api.Content.TextContent} are allowed.
+     */
+    public CompletableFuture<List<ContentEmbedding>> embedContents(
+            GenerativeModel model,
+            String taskType,
+            String title,
+            Long outputDimensionality
+    ) {
+        return execute(() -> {
+
+            var requests = convertGenerationContents(model)
+                    .stream()
+                    .map(generationContent -> new EmbedContentRequest(
+                            model.modelName(),
+                            generationContent,
+                            taskType,
+                            title,
+                            outputDimensionality
+                    ))
+                    .toList();
+
+            var request = new BatchEmbedContentRequest(requests);
+
+            CompletableFuture<HttpResponse<String>> response = client.sendAsync(
+                    HttpRequest.newBuilder()
+                            .POST(HttpRequest.BodyPublishers.ofString(
+                                    jsonParser.toJson(request)
+                            ))
+                            .uri(URI.create("%s/%s:batchEmbedContents?key=%s".formatted(urlPrefix, model.modelName(), apiKey)))
+                            .build(),
+                    HttpResponse.BodyHandlers.ofString()
+            );
+            return response
+                    .thenApply(HttpResponse::body)
+                    .thenApply(body -> {
+                        try {
+                            BatchEmbedContentResponse becr = jsonParser.fromJson(body, BatchEmbedContentResponse.class);
+                            if (becr.embeddings() == null) {
+                                throw new RuntimeException();
+                            }
+                            return becr
+                                    .embeddings();
+                        } catch (Exception e) {
+                            throw new RuntimeException("Unexpected body:\n" + body, e);
+                        }
+                    });
+
+        });
+    }
+
     private static GenerateContentRequest convert(GenerativeModel model) {
-        List<GenerationContent> generationContents = model.contents().stream()
+        List<GenerationContent> generationContents = convertGenerationContents(model);
+        return new GenerateContentRequest(model.modelName(), generationContents, model.safetySettings(), model.generationConfig());
+    }
+
+    private static List<GenerationContent> convertGenerationContents(GenerativeModel model) {
+        return model.contents().stream()
                 .map(content -> {
                     // change to "switch" over sealed type with jdk 21
                     if (content instanceof Content.TextContent textContent) {
@@ -294,7 +364,6 @@ private static GenerateContentRequest convert(GenerativeModel model) {
                     }
                 })
                 .toList();
-        return new GenerateContentRequest(model.modelName(), generationContents, model.safetySettings(), model.generationConfig());
     }
 
     private <T> T execute(ThrowingSupplier<T> supplier) {
@@ -379,6 +448,35 @@ public record TypedSafetyRating(
 
     }
 
+    /**
+     * A list of floats representing an embedding.
+     *
+     * @param values A list of floats representing an embedding.
+     */
+    public record ContentEmbedding(
+            List<Double> values
+    ) {
+    }
+
+    private record BatchEmbedContentRequest(
+            List<EmbedContentRequest> requests
+    ) {
+    }
+
+    private record EmbedContentRequest(
+            String model,
+            GenerationContent content,
+            String taskType,
+            String title,
+            Long outputDimensionality
+    ) {
+    }
+
+    private record BatchEmbedContentResponse(
+            List<ContentEmbedding> embeddings
+    ) {
+    }
+
     private record CountTokenRequest(
             GenerateContentRequest generateContentRequest
     ) {

gemini-api/src/main/java/swiss/ameri/gemini/api/TaskType.java

@@ -0,0 +1,46 @@
+package swiss.ameri.gemini.api;
+
+/**
+ * Type of task for which the embedding will be used.
+ */
+public enum TaskType {
+    /**
+     * Unset value, which will default to one of the other enum values.
+     */
+    TASK_TYPE_UNSPECIFIED,
+
+    /**
+     * Specifies the given text is a query in a search/retrieval setting.
+     */
+    RETRIEVAL_QUERY,
+
+    /**
+     * Specifies the given text is a document from the corpus being searched.
+     */
+    RETRIEVAL_DOCUMENT,
+
+    /**
+     * Specifies the given text will be used for Semantic Textual Similarity (STS).
+     */
+    SEMANTIC_SIMILARITY,
+
+    /**
+     * Specifies that the given text will be classified.
+     */
+    CLASSIFICATION,
+
+    /**
+     * Specifies that the embeddings will be used for clustering.
+     */
+    CLUSTERING,
+
+    /**
+     * Specifies that the given text will be used for question answering.
+     */
+    QUESTION_ANSWERING,
+
+    /**
+     * Specifies that the given text will be used for fact verification.
+     */
+    FACT_VERIFICATION
+}

gemini-tester/src/main/java/swiss/ameri/gemini/tester/GeminiTester.java

@@ -7,6 +7,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Base64;
+import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -39,11 +40,45 @@ public static void main(String[] args) throws Exception {
             generateContentStream(genAi);
             multiChatTurn(genAi);
             textAndImage(genAi);
+            embedContents(genAi);
         }
 
 
     }
 
+    private static void embedContents(GenAi genAi) {
+        System.out.println("----- embed contents");
+        var model = GenerativeModel.builder()
+                .modelName(ModelVariant.TEXT_EMBEDDING_004)
+                .addContent(Content.textContent(
+                        Content.Role.USER,
+                        "Write a 50 word story about a magic backpack."
+                ))
+                .addContent(Content.textContent(
+                        Content.Role.MODEL,
+                        "bla bla bla bla"
+                ))
+                .addSafetySetting(SafetySetting.of(
+                        SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                        SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH
+                ))
+                .generationConfig(new GenerationConfig(
+                        null,
+                        null,
+                        null,
+                        null,
+                        null,
+                        null,
+                        null
+                ))
+                .build();
+
+        List<GenAi.ContentEmbedding> embeddings = genAi.embedContents(model, null, null, null).join();
+        System.out.println("Embedding count: " + embeddings.size());
+        System.out.println("Values per embedding: " + embeddings.stream().map(GenAi.ContentEmbedding::values).map(List::size).toList());
+
+    }
+
     private static void countTokens(GenAi genAi) {
         System.out.println("----- count tokens");
         var model = createStoryModel();