OPIK-611 changed gemini model creation to map struct

comet-ml · Jan 13, 2025 · 46fc67f · 46fc67f
1 parent 138ca1b
commit 46fc67f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 22 deletions.
diff --git a/...-backend/src/main/java/com/comet/opik/domain/llmproviders/LlmProviderClientGenerator.java b/...-backend/src/main/java/com/comet/opik/domain/llmproviders/LlmProviderClientGenerator.java
@@ -14,6 +14,8 @@
 
 @RequiredArgsConstructor
 public class LlmProviderClientGenerator {
+    private static final int MAX_RETRIES = 1;
+
     private final @NonNull LlmProviderClientConfig llmProviderClientConfig;
 
     public AnthropicClient newAnthropicClient(@NonNull String apiKey) {
@@ -57,32 +59,14 @@ public OpenAiClient newOpenAiClient(@NonNull String apiKey) {
                 .build();
     }
 
-    // TODO: mapper
     public GoogleAiGeminiChatModel newGeminiClient(@NonNull String apiKey, @NonNull ChatCompletionRequest request) {
-        return GoogleAiGeminiChatModel.builder()
-                .apiKey(apiKey)
-                .modelName(request.model())
-                .maxOutputTokens(request.maxCompletionTokens())
-                .maxRetries(1)
-                .stopSequences(request.stop())
-                .temperature(request.temperature())
-                .topP(request.topP())
-                .timeout(llmProviderClientConfig.getCallTimeout().toJavaDuration())
-                .build();
+        return LlmProviderGeminiMapper.INSTANCE.toGeminiChatModel(apiKey, request,
+                llmProviderClientConfig.getCallTimeout().toJavaDuration(), MAX_RETRIES);
     }
 
-    // TODO: mapper
     public GoogleAiGeminiStreamingChatModel newGeminiStreamingClient(
             @NonNull String apiKey, @NonNull ChatCompletionRequest request) {
-        return GoogleAiGeminiStreamingChatModel.builder()
-                .apiKey(apiKey)
-                .modelName(request.model())
-                .maxOutputTokens(request.maxCompletionTokens())
-                .maxRetries(1)
-                .stopSequences(request.stop())
-                .temperature(request.temperature())
-                .topP(request.topP())
-                .timeout(llmProviderClientConfig.getCallTimeout().toJavaDuration())
-                .build();
+        return LlmProviderGeminiMapper.INSTANCE.toGeminiStreamingChatModel(apiKey, request,
+                llmProviderClientConfig.getCallTimeout().toJavaDuration(), MAX_RETRIES);
     }
 }
diff --git a/...pik-backend/src/main/java/com/comet/opik/domain/llmproviders/LlmProviderGeminiMapper.java b/...pik-backend/src/main/java/com/comet/opik/domain/llmproviders/LlmProviderGeminiMapper.java
@@ -11,6 +11,8 @@
 import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.data.message.SystemMessage;
 import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.model.googleai.GoogleAiGeminiChatModel;
+import dev.langchain4j.model.googleai.GoogleAiGeminiStreamingChatModel;
 import dev.langchain4j.model.output.Response;
 import jakarta.ws.rs.BadRequestException;
 import lombok.NonNull;
@@ -19,6 +21,7 @@
 import org.mapstruct.Named;
 import org.mapstruct.factory.Mappers;
 
+import java.time.Duration;
 import java.util.List;
 
 @Mapper
@@ -76,4 +79,20 @@ default Usage mapToUsage(@NonNull Response<AiMessage> response) {
                 .totalTokens(response.tokenUsage().totalTokenCount())
                 .build();
     }
+
+    @Mapping(expression = "java(request.model())", target = "modelName")
+    @Mapping(expression = "java(request.maxCompletionTokens())", target = "maxOutputTokens")
+    @Mapping(expression = "java(request.stop())", target = "stopSequences")
+    @Mapping(expression = "java(request.temperature())", target = "temperature")
+    @Mapping(expression = "java(request.topP())", target = "topP")
+    GoogleAiGeminiChatModel toGeminiChatModel(
+            @NonNull String apiKey, @NonNull ChatCompletionRequest request, @NonNull Duration timeout, int maxRetries);
+
+    @Mapping(expression = "java(request.model())", target = "modelName")
+    @Mapping(expression = "java(request.maxCompletionTokens())", target = "maxOutputTokens")
+    @Mapping(expression = "java(request.stop())", target = "stopSequences")
+    @Mapping(expression = "java(request.temperature())", target = "temperature")
+    @Mapping(expression = "java(request.topP())", target = "topP")
+    GoogleAiGeminiStreamingChatModel toGeminiStreamingChatModel(
+            @NonNull String apiKey, @NonNull ChatCompletionRequest request, @NonNull Duration timeout, int maxRetries);
 }