Feat #70: Implemented Streaming response + refactorings

devoxx · May 26, 2024 · 42c04de · 42c04de
1 parent b0d1ac2
commit 42c04de
Show file tree

Hide file tree

Showing 60 changed files with 578 additions and 453 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -5,19 +5,19 @@ plugins {
 }
 
 group = "com.devoxx.genie"
-version = "0.1.11"
+version = "0.1.12"
 
 repositories {
     mavenCentral()
 }
 
 dependencies {
-    implementation("dev.langchain4j:langchain4j:0.30.0")
-    implementation("dev.langchain4j:langchain4j-ollama:0.30.0")
-    implementation("dev.langchain4j:langchain4j-local-ai:0.30.0")
-    implementation("dev.langchain4j:langchain4j-open-ai:0.30.0")
-    implementation("dev.langchain4j:langchain4j-anthropic:0.30.0")
-    implementation("dev.langchain4j:langchain4j-mistral-ai:0.30.0")
+    implementation("dev.langchain4j:langchain4j:0.31.0")
+    implementation("dev.langchain4j:langchain4j-ollama:0.31.0")
+    implementation("dev.langchain4j:langchain4j-local-ai:0.31.0")
+    implementation("dev.langchain4j:langchain4j-open-ai:0.31.0")
+    implementation("dev.langchain4j:langchain4j-anthropic:0.31.0")
+    implementation("dev.langchain4j:langchain4j-mistral-ai:0.31.0")
     implementation("org.commonmark:commonmark:0.22.0")
 
     compileOnly("org.projectlombok:lombok:1.18.32")

diff --git a/src/main/java/com/devoxx/genie/action/AddFileAction.java b/src/main/java/com/devoxx/genie/action/AddFileAction.java
@@ -7,8 +7,6 @@
 import com.intellij.openapi.actionSystem.CommonDataKeys;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
-import com.intellij.openapi.wm.ToolWindow;
-import com.intellij.openapi.wm.ToolWindowManager;
 import org.jetbrains.annotations.NotNull;
 
 import static com.devoxx.genie.ui.util.WindowPluginUtil.ensureToolWindowVisible;
@@ -17,6 +15,7 @@ public class AddFileAction extends AnAction {
 
     /**
      * Add file to the window context.
+     *
      * @param e the action event
      */
     @Override

diff --git a/src/main/java/com/devoxx/genie/action/AddSnippetAction.java b/src/main/java/com/devoxx/genie/action/AddSnippetAction.java
@@ -29,6 +29,7 @@ public class AddSnippetAction extends AnAction {
 
     /**
      * Add a snippet to the tool window.
+     *
      * @param e the action event
      */
     @Override
@@ -53,6 +54,7 @@ public void actionPerformed(@NotNull AnActionEvent e) {
 
     /**
      * Add the selected file to the file list manager.
+     *
      * @param selectedFile the selected file
      */
     private static void addSelectedFile(VirtualFile selectedFile) {
@@ -65,9 +67,10 @@ private static void addSelectedFile(VirtualFile selectedFile) {
 
     /**
      * Create a virtual file and add it to the file list manager.
-     * @param originalFile the original file
+     *
+     * @param originalFile   the original file
      * @param selectionModel the selection model
-     * @param selectedText the selected text
+     * @param selectedText   the selected text
      */
     private void createAndAddVirtualFile(VirtualFile originalFile,
                                          SelectionModel selectionModel,

diff --git a/src/main/java/com/devoxx/genie/chatmodel/ChatModelFactory.java b/src/main/java/com/devoxx/genie/chatmodel/ChatModelFactory.java
@@ -10,13 +10,15 @@ public interface ChatModelFactory {
 
     /**
      * Create a chat model with the given parameters.
+     *
      * @param chatModel the chat model
      * @return the chat model
      */
     ChatLanguageModel createChatModel(ChatModel chatModel);
 
     /**
      * Create a streaming chat model with the given parameters.
+     *
      * @param chatModel the chat model
      * @return the streaming chat model
      */
@@ -26,6 +28,7 @@ default StreamingChatLanguageModel createStreamingChatModel(ChatModel chatModel)
 
     /**
      * List the available model names.
+     *
      * @return the list of model names
      */
     default List<String> getModelNames() {
@@ -34,6 +37,7 @@ default List<String> getModelNames() {
 
     /**
      * Get the model provider API key.
+     *
      * @return the API key
      */
     default String getApiKey() {

diff --git a/src/main/java/com/devoxx/genie/chatmodel/ChatModelFactoryProvider.java b/src/main/java/com/devoxx/genie/chatmodel/ChatModelFactoryProvider.java
@@ -1,14 +1,14 @@
 package com.devoxx.genie.chatmodel;
 
-import com.devoxx.genie.model.enumarations.ModelProvider;
 import com.devoxx.genie.chatmodel.anthropic.AnthropicChatModelFactory;
 import com.devoxx.genie.chatmodel.deepinfra.DeepInfraChatModelFactory;
+import com.devoxx.genie.chatmodel.gemini.GeminiChatModelFactory;
 import com.devoxx.genie.chatmodel.groq.GroqChatModelFactory;
+import com.devoxx.genie.chatmodel.jan.JanChatModelFactory;
 import com.devoxx.genie.chatmodel.mistral.MistralChatModelFactory;
 import com.devoxx.genie.chatmodel.ollama.OllamaChatModelFactory;
 import com.devoxx.genie.chatmodel.openai.OpenAIChatModelFactory;
-import com.devoxx.genie.chatmodel.gemini.GeminiChatModelFactory;
-import com.devoxx.genie.chatmodel.jan.JanChatModelFactory;
+import com.devoxx.genie.model.enumarations.ModelProvider;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Map;
@@ -29,10 +29,11 @@ public class ChatModelFactoryProvider {
         ModelProvider.DeepInfra, DeepInfraChatModelFactory::new,
         ModelProvider.Gemini, GeminiChatModelFactory::new,
         ModelProvider.Jan, JanChatModelFactory::new
-        );
+    );
 
     /**
      * Get the factory by provider.
+     *
      * @param provider the provider
      * @return the factory
      */

diff --git a/src/main/java/com/devoxx/genie/chatmodel/ChatModelProvider.java b/src/main/java/com/devoxx/genie/chatmodel/ChatModelProvider.java
@@ -41,6 +41,7 @@ public ChatModelProvider() {
 
     /**
      * Get the chat language model for selected model provider.
+     *
      * @param chatMessageContext the chat message context
      * @return the chat language model
      */
@@ -51,6 +52,7 @@ public ChatLanguageModel getChatLanguageModel(@NotNull ChatMessageContext chatMe
 
     /**
      * Get the streaming chat language model for selected model provider.
+     *
      * @param chatMessageContext the chat message context
      * @return the streaming chat language model
      */
@@ -61,6 +63,7 @@ public StreamingChatLanguageModel getStreamingChatLanguageModel(@NotNull ChatMes
 
     /**
      * Get the chat model factory for the selected model provider.
+     *
      * @param chatMessageContext the chat message context
      * @return the chat model factory
      */
@@ -75,6 +78,7 @@ public StreamingChatLanguageModel getStreamingChatLanguageModel(@NotNull ChatMes
 
     /**
      * Initialize chat model settings by default or by user settings.
+     *
      * @return the chat model
      */
     public @NotNull ChatModel initChatModel(@NotNull ChatMessageContext chatMessageContext) {
@@ -93,8 +97,9 @@ public StreamingChatLanguageModel getStreamingChatLanguageModel(@NotNull ChatMes
     /**
      * Set max output tokens.
      * Some extra work because of the settings state that didn't like the integer input field.
+     *
      * @param settingsState the settings state
-     * @param chatModel the chat model
+     * @param chatModel     the chat model
      */
     private static void setMaxOutputTokens(@NotNull SettingsState settingsState, ChatModel chatModel) {
         String maxOutputTokens = settingsState.getMaxOutputTokens();

diff --git a/src/main/java/com/devoxx/genie/chatmodel/deepinfra/DeepInfraChatModelFactory.java b/src/main/java/com/devoxx/genie/chatmodel/deepinfra/DeepInfraChatModelFactory.java
@@ -47,18 +47,18 @@ public String getApiKey() {
 
     @Override
     public List<String> getModelNames() {
-       return List.of(
-           "meta-llama/Meta-Llama-3-70B-Instruct",
-           "meta-llama/Meta-Llama-3-8B-Instruct",
-           "mistralai/Mixtral-8x7B-Instruct-v0.1",
-           "mistralai/Mixtral-8x22B-Instruct-v0.1",
-           "microsoft/WizardLM-2-8x22B",
-           "microsoft/WizardLM-2-7B",
-           "databricks/dbrx-instruct",
-           "openchat/openchat_3.5",
-           "google/gemma-7b-it",
-           "Phind/Phind-CodeLlama-34B-v2",
-           "bigcode/starcoder2-15b"
-       );
+        return List.of(
+            "meta-llama/Meta-Llama-3-70B-Instruct",
+            "meta-llama/Meta-Llama-3-8B-Instruct",
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            "mistralai/Mixtral-8x22B-Instruct-v0.1",
+            "microsoft/WizardLM-2-8x22B",
+            "microsoft/WizardLM-2-7B",
+            "databricks/dbrx-instruct",
+            "openchat/openchat_3.5",
+            "google/gemma-7b-it",
+            "Phind/Phind-CodeLlama-34B-v2",
+            "bigcode/starcoder2-15b"
+        );
     }
 }
diff --git a/src/main/java/com/devoxx/genie/chatmodel/groq/GroqChatModelFactory.java b/src/main/java/com/devoxx/genie/chatmodel/groq/GroqChatModelFactory.java
@@ -5,7 +5,6 @@
 import com.devoxx.genie.ui.SettingsState;
 import dev.langchain4j.model.chat.ChatLanguageModel;
 import dev.langchain4j.model.chat.StreamingChatLanguageModel;
-import dev.langchain4j.model.mistralai.MistralAiStreamingChatModel;
 import dev.langchain4j.model.openai.OpenAiChatModel;
 import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
 import org.jetbrains.annotations.NotNull;

diff --git a/src/main/java/com/devoxx/genie/chatmodel/jan/JanChatModelFactory.java b/src/main/java/com/devoxx/genie/chatmodel/jan/JanChatModelFactory.java
@@ -11,7 +11,6 @@
 import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import dev.langchain4j.model.localai.LocalAiChatModel;
 import dev.langchain4j.model.localai.LocalAiStreamingChatModel;
-import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
 import okhttp3.OkHttpClient;
 import org.jetbrains.annotations.NotNull;
 
@@ -52,6 +51,7 @@ public StreamingChatLanguageModel createStreamingChatModel(@NotNull ChatModel ch
 
     /**
      * Get the model names from the Jan service.
+     *
      * @return List of model names
      */
     @Override

diff --git a/src/main/java/com/devoxx/genie/chatmodel/ollama/OllamaChatModelFactory.java b/src/main/java/com/devoxx/genie/chatmodel/ollama/OllamaChatModelFactory.java
@@ -11,7 +11,6 @@
 import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import dev.langchain4j.model.ollama.OllamaChatModel;
 import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
-import okhttp3.OkHttpClient;
 import org.jetbrains.annotations.NotNull;
 
 import java.io.IOException;
@@ -21,9 +20,6 @@
 
 public class OllamaChatModelFactory implements ChatModelFactory {
 
-    // Moved client instance here for the sake of better performance
-    private final OkHttpClient client = new OkHttpClient();
-
     @Override
     public ChatLanguageModel createChatModel(@NotNull ChatModel chatModel) {
         return OllamaChatModel.builder()
@@ -49,13 +45,14 @@ public StreamingChatLanguageModel createStreamingChatModel(@NotNull ChatModel ch
 
     /**
      * Get the model names from the Ollama service.
+     *
      * @return List of model names
      */
     @Override
     public List<String> getModelNames() {
         List<String> modelNames = new ArrayList<>();
         try {
-            OllamaModelEntryDTO[] ollamaModels = new OllamaService(client).getModels();
+            OllamaModelEntryDTO[] ollamaModels = OllamaService.getInstance().getModels();
             for (OllamaModelEntryDTO model : ollamaModels) {
                 modelNames.add(model.getName());
             }

diff --git a/src/main/java/com/devoxx/genie/model/gemini/GeminiApi.java b/src/main/java/com/devoxx/genie/model/gemini/GeminiApi.java
@@ -1,7 +1,10 @@
 package com.devoxx.genie.model.gemini;
 
 import retrofit2.Call;
-import retrofit2.http.*;
+import retrofit2.http.Body;
+import retrofit2.http.Headers;
+import retrofit2.http.POST;
+import retrofit2.http.Url;
 
 public interface GeminiApi {
 

diff --git a/src/main/java/com/devoxx/genie/model/gemini/GeminiChatModel.java b/src/main/java/com/devoxx/genie/model/gemini/GeminiChatModel.java
@@ -4,10 +4,10 @@
 import com.devoxx.genie.model.gemini.model.Part;
 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.output.Response;
-import dev.langchain4j.data.message.SystemMessage;
 import lombok.Builder;
 import org.jetbrains.annotations.NotNull;
 
@@ -80,7 +80,8 @@ public Response<AiMessage> generate(@NotNull List<ChatMessage> messages) {
 
     /**
      * Create a message for the Gemini API
-     * @param role the user role
+     *
+     * @param role    the user role
      * @param message a chat message
      */
     private static Content createMessage(String role, String message) {

diff --git a/src/main/java/com/devoxx/genie/model/gemini/model/CitationMetadata.java b/src/main/java/com/devoxx/genie/model/gemini/model/CitationMetadata.java
@@ -1,7 +1,11 @@
 package com.devoxx.genie.model.gemini.model;
 
 // CitationMetadata.java
-import lombok.*;
+
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
 
 import java.util.List;
 

diff --git a/src/main/java/com/devoxx/genie/model/gemini/model/CitationSources.java b/src/main/java/com/devoxx/genie/model/gemini/model/CitationSources.java
@@ -1,7 +1,11 @@
 package com.devoxx.genie.model.gemini.model;
 
 // CitationSources.java
-import lombok.*;
+
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
 
 @Data
 @NoArgsConstructor

diff --git a/src/main/java/com/devoxx/genie/model/gemini/model/Content.java b/src/main/java/com/devoxx/genie/model/gemini/model/Content.java
@@ -1,7 +1,10 @@
 package com.devoxx.genie.model.gemini.model;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
-import lombok.*;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
 
 import java.util.List;
 

diff --git a/src/main/java/com/devoxx/genie/model/gemini/model/GenerationConfig.java b/src/main/java/com/devoxx/genie/model/gemini/model/GenerationConfig.java
@@ -1,7 +1,10 @@
 package com.devoxx.genie.model.gemini.model;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
-import lombok.*;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
 
 @Data
 @NoArgsConstructor

diff --git a/src/main/java/com/devoxx/genie/model/gemini/model/Part.java b/src/main/java/com/devoxx/genie/model/gemini/model/Part.java
@@ -1,6 +1,9 @@
 package com.devoxx.genie.model.gemini.model;
 
-import lombok.*;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
 
 @Data
 @NoArgsConstructor

diff --git a/src/main/java/com/devoxx/genie/model/gemini/model/SystemInstruction.java b/src/main/java/com/devoxx/genie/model/gemini/model/SystemInstruction.java
@@ -1,7 +1,10 @@
 package com.devoxx.genie.model.gemini.model;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
-import lombok.*;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
 
 import java.util.List;
 

diff --git a/src/main/java/com/devoxx/genie/service/ChatMemoryService.java b/src/main/java/com/devoxx/genie/service/ChatMemoryService.java
@@ -3,13 +3,19 @@
 import com.intellij.openapi.application.ApplicationManager;
 import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.memory.chat.MessageWindowChatMemory;
+import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
 
 public class ChatMemoryService {
 
-    private final MessageWindowChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
+    private final MessageWindowChatMemory chatMemory =
+        MessageWindowChatMemory.builder()
+            .id("devoxxgenie")
+            .chatMemoryStore(new InMemoryChatMemoryStore())
+            .maxMessages(10)    // TODO Make this configurable in the Settings page
+            .build();
 
     @NotNull
     public static ChatMemoryService getInstance() {