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 index 339f6c7..7abfb09 100644 --- 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 getModels() { return getModels(MODEL_PROVIDER);