GenerativeModelBuilder

Author: michael-ameri

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

@@ -67,8 +67,9 @@ public List<Model> listModels() {
         });
     }
 
-    /***
+    /**
      * Get information of a model. Can be used to create a {@link GenerativeModel}.
+     *
      * @param model of which the information is wanted.
      * @see #listModels()
      */

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

@@ -1,5 +1,6 @@
 package swiss.ameri.gemini.api;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -17,14 +18,86 @@ public record GenerativeModel(
         GenerationConfig generationConfig
 ) {
 
-    public static GenerativeModel of(
-            ModelVariant modelVariant,
-            List<Content> contents,
-            List<SafetySetting> safetySettings,
-            GenerationConfig generationConfig
-    ) {
-        // todo add builder, which accepts modelVariant
-        return new GenerativeModel(modelVariant.variant(), contents, safetySettings, generationConfig);
+    /**
+     * Create a {@link GenerativeModelBuilder}.
+     *
+     * @return an empty {@link GenerativeModelBuilder}
+     */
+    public static GenerativeModelBuilder builder() {
+        return new GenerativeModelBuilder();
+    }
+
+    /**
+     * A builder for {@link GenerativeModel}. Currently, does not validate the fields when building the model. Not thread-safe.
+     */
+    public static class GenerativeModelBuilder {
+        private String modelName;
+        private GenerationConfig generationConfig;
+        private final List<Content> contents = new ArrayList<>();
+        private final List<SafetySetting> safetySettings = new ArrayList<>();
+
+        /**
+         * Set the model name.
+         *
+         * @param modelName to be set
+         * @return this
+         */
+        public GenerativeModelBuilder modelName(String modelName) {
+            this.modelName = modelName;
+            return this;
+        }
+
+        /**
+         * Set the model name.
+         *
+         * @param modelVariant to be set
+         * @return this
+         */
+        public GenerativeModelBuilder modelName(ModelVariant modelVariant) {
+            return modelName(modelVariant == null ? null : modelVariant.variant());
+        }
+
+        /**
+         * Add content
+         *
+         * @param content to be added
+         * @return this
+         */
+        public GenerativeModelBuilder addContent(Content content) {
+            this.contents.add(content);
+            return this;
+        }
+
+        /**
+         * Add safety setting
+         *
+         * @param safetySetting to be added
+         * @return this
+         */
+        public GenerativeModelBuilder addSafetySetting(SafetySetting safetySetting) {
+            this.safetySettings.add(safetySetting);
+            return this;
+        }
+
+        /**
+         * Set the generation config
+         *
+         * @param generationConfig to be set
+         * @return this
+         */
+        public GenerativeModelBuilder generationConfig(GenerationConfig generationConfig) {
+            this.generationConfig = generationConfig;
+            return this;
+        }
+
+        /**
+         * Build the model based on this builder.
+         *
+         * @return a completed (not necessarily validated) {@link GenerativeModel}
+         */
+        public GenerativeModel build() {
+            return new GenerativeModel(modelName, contents, safetySettings, generationConfig);
+        }
     }
 
 }

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

@@ -17,8 +17,8 @@ public static SafetySetting of(
             HarmBlockThreshold threshold
     ) {
         return new SafetySetting(
-                category.name(),
-                threshold.name()
+                category == null ? null : category.name(),
+                threshold == null ? null : threshold.name()
         );
     }
 

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

@@ -4,7 +4,6 @@
 import swiss.ameri.gemini.gson.GsonJsonParser;
 import swiss.ameri.gemini.spi.JsonParser;
 
-import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -38,27 +37,26 @@ public static void main(String[] args) throws Exception {
         );
 
         System.out.println("-----");
-        var model = GenerativeModel.of(
-                ModelVariant.GEMINI_1_0_PRO,
-                List.of(
-                        new Content.TextContent(
-                                Content.Role.USER.roleName(),
-                                "Write a 300 word story about a magic backpack."
-                        )
-                ),
-                List.of(
-                        SafetySetting.of(SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH)
-                ),
-                new GenerationConfig(
+        var model = GenerativeModel.builder()
+                .modelName(ModelVariant.GEMINI_1_0_PRO)
+                .addContent(new Content.TextContent(
+                        Content.Role.USER.roleName(),
+                        "Write a 300 word story about a magic backpack."
+                ))
+                .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();
         genAi.generateContent(model)
                 .thenAccept(System.out::println)
                 .get(20, TimeUnit.SECONDS);