Merge pull request #410 from samkerr4coding/feat/issue-81

feat(issue-81): Adding streaming ability with gemini models
devoxx · Dec 19, 2024 · 315b075 · 315b075
2 parents d23c6d7 + 76fc9c6
commit 315b075
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/src/main/java/com/devoxx/genie/chatmodel/cloud/google/GoogleChatModelFactory.java b/src/main/java/com/devoxx/genie/chatmodel/cloud/google/GoogleChatModelFactory.java
@@ -5,9 +5,12 @@
 import com.devoxx.genie.model.LanguageModel;
 import com.devoxx.genie.model.enumarations.ModelProvider;
 import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import dev.langchain4j.model.googleai.GoogleAiGeminiChatModel;
+import dev.langchain4j.model.googleai.GoogleAiGeminiStreamingChatModel;
 import org.jetbrains.annotations.NotNull;
 
+import java.time.Duration;
 import java.util.List;
 
 public class GoogleChatModelFactory implements ChatModelFactory {
@@ -24,6 +27,16 @@ public ChatLanguageModel createChatModel(@NotNull ChatModel chatModel) {
             .build();
     }
 
+    @Override
+    public StreamingChatLanguageModel createStreamingChatModel(@NotNull ChatModel chatModel) {
+        return GoogleAiGeminiStreamingChatModel.builder()
+                .apiKey(getApiKey(MODEL_PROVIDER))
+                .modelName(chatModel.getModelName())
+                .temperature(chatModel.getTemperature())
+                .maxOutputTokens(chatModel.getMaxTokens())
+                .build();
+    }
+
     @Override
     public List<LanguageModel> getModels() {
         return getModels(MODEL_PROVIDER);