Update Google AI Studio Addon to use langchain4j's google-ai-gemini module

Author: BlueTree242

google-aistudio/build.gradle.kts

@@ -1,6 +1,7 @@
 plugins {
     id("java")
     id("net.minecrell.plugin-yml.bukkit")
+    id("com.github.johnrengelman.shadow")
 }
 
 var pckg = rootProject.ext.get("pckg")
@@ -13,5 +14,12 @@ bukkit {
 }
 
 dependencies {
-    compileOnly(libs.retrofit2)
+    implementation(libs.langchain4j.googleai) {
+        isTransitive = false
+    }
+    bukkitLibrary(libs.retrofit2.gson)
+}
+
+tasks.shadowJar {
+    relocate("dev.langchain4j.model.googleai", "$pckg.dependencies.langchain4j.googleai")
 }

google-aistudio/src/main/java/dev/bluetree242/saaiaddons/aistudio/GoogleAiStudioChatModelProvider.java

@@ -1,36 +1,57 @@
 package dev.bluetree242.saaiaddons.aistudio;
 
-import dev.bluetree242.saaiaddons.aistudio.api.AiStudioChatLanguageModel;
 import dev.bluetree242.serverassistantai.api.ServerAssistantAIAPI;
 import dev.bluetree242.serverassistantai.api.config.option.OptionMap;
 import dev.bluetree242.serverassistantai.api.registry.chatmodel.ChatModelContext;
 import dev.bluetree242.serverassistantai.api.registry.chatmodel.ChatModelProvider;
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.data.message.ChatMessage;
+import dev.langchain4j.data.message.SystemMessage;
+import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.googleai.GeminiHarmBlockThreshold;
+import dev.langchain4j.model.googleai.GeminiHarmCategory;
+import dev.langchain4j.model.googleai.GoogleAiGeminiChatModel;
+import dev.langchain4j.model.output.Response;
 import lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.time.Duration;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 @RequiredArgsConstructor
-public class GoogleAiStudioChatModelProvider implements ChatModelProvider<AiStudioChatLanguageModel> {
+public class GoogleAiStudioChatModelProvider implements ChatModelProvider<GoogleAiStudioChatModelProvider.GoogleAiWrapper> {
+    private static final Map<GeminiHarmCategory, GeminiHarmBlockThreshold> safetySettings;
+
+    static {
+        Map<GeminiHarmCategory, GeminiHarmBlockThreshold> result = new HashMap<>();
+        for (GeminiHarmCategory value : GeminiHarmCategory.values()) {
+            result.put(value, GeminiHarmBlockThreshold.BLOCK_NONE);
+        }
+        safetySettings = Collections.unmodifiableMap(result);
+    }
+
     private final ServerAssistantAIAPI api;
 
     @NotNull
     @Override
-    public AiStudioChatLanguageModel provide(@NotNull ChatModelContext context) {
+    public GoogleAiWrapper provide(@NotNull ChatModelContext context) {
         OptionMap options = context.options();
         String model = options.getString("model");
-        //noinspection unchecked
         if (model.isBlank()) throw new IllegalStateException("Please set the model for Google AI Studio chat model.");
-        return AiStudioChatLanguageModel.builder()
-                .model(model)
+        return new GoogleAiWrapper(GoogleAiGeminiChatModel.builder()
+                .safetySettings(safetySettings)
+                .modelName(model)
                 .timeout(Duration.ofSeconds(options.getLong("timeout")))
                 .maxOutputTokens(options.getIntegerOrDefault("max_output_tokens", i -> i == 0, null))
                 .stopSequences(options.getList("stop").getStringList())
                 .apiKey(api.getCredentialsRegistry().getConfigured(GoogleAiStudioAddon.NAME, GoogleAiStudioCredentialsLoader.class))
-                .build();
+                .build());
     }
 
     @NotNull
@@ -58,4 +79,11 @@ public Map<String, Object> export(@NotNull ChatModelContext context) {
     public String getDisplayName(@Nullable ChatModelContext context) {
         return "Google AI Studio";
     }
+
+    public record GoogleAiWrapper(GoogleAiGeminiChatModel model) implements ChatLanguageModel {
+        @Override
+        public Response<AiMessage> generate(List<ChatMessage> messages) {
+            return model.generate(UserMessage.userMessage(messages.stream().filter(m -> m instanceof SystemMessage).map(m -> (SystemMessage) m).map(SystemMessage::text).collect(Collectors.joining("\n\n\n\n"))));
+        }
+    }
 }

google-aistudio/src/main/java/dev/bluetree242/saaiaddons/aistudio/GoogleAiStudioEmbeddingProvider.java

@@ -1,10 +1,10 @@
 package dev.bluetree242.saaiaddons.aistudio;
 
-import dev.bluetree242.saaiaddons.aistudio.api.AiStudioEmbeddingModel;
 import dev.bluetree242.serverassistantai.api.ServerAssistantAIAPI;
 import dev.bluetree242.serverassistantai.api.config.option.OptionMap;
 import dev.bluetree242.serverassistantai.api.registry.embedding.EmbeddingContext;
 import dev.bluetree242.serverassistantai.api.registry.embedding.EmbeddingModelProvider;
+import dev.langchain4j.model.googleai.GoogleAiEmbeddingModel;
 import lombok.RequiredArgsConstructor;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -13,17 +13,21 @@
 import java.util.Map;
 
 @RequiredArgsConstructor
-public class GoogleAiStudioEmbeddingProvider implements EmbeddingModelProvider<AiStudioEmbeddingModel> {
+public class GoogleAiStudioEmbeddingProvider implements EmbeddingModelProvider<GoogleAiEmbeddingModel> {
     private final ServerAssistantAIAPI api;
 
     @Override
-    public @NotNull AiStudioEmbeddingModel provide(@NotNull EmbeddingContext context) {
+    public @NotNull GoogleAiEmbeddingModel provide(@NotNull EmbeddingContext context) {
         OptionMap options = context.options();
         String apiKey = api.getCredentialsRegistry().getConfigured(GoogleAiStudioAddon.NAME, GoogleAiStudioCredentialsLoader.class);
         String model = options.getString("model");
         if (model.isBlank())
             throw new IllegalStateException("Please set the model for Google AI Studio embedding model.");
-        return new AiStudioEmbeddingModel(apiKey, model, Duration.ofSeconds(options.getLong("timeout")));
+        return GoogleAiEmbeddingModel.builder()
+                .apiKey(apiKey)
+                .modelName(model)
+                .timeout(Duration.ofSeconds(options.getLong("timeout")))
+                .build();
     }
 
     @NotNull