Allow using existing Llama Server instead of running locally (#345)

* Add setting to use existing Llama server

* minor UI improvements

* support infill template configuration

---------

Co-authored-by: Carl-Robert Linnupuu <[email protected]>

Author: PhilKes

src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestProvider.java

@@ -1,5 +1,6 @@
 package ee.carlrobert.codegpt.codecompletions;
 
+import ee.carlrobert.codegpt.completions.llama.LlamaModel;
 import ee.carlrobert.codegpt.settings.state.LlamaSettingsState;
 import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest;
 import ee.carlrobert.llm.client.openai.completion.request.OpenAITextCompletionRequest;
@@ -26,13 +27,25 @@ public OpenAITextCompletionRequest buildOpenAIRequest() {
   }
 
   public LlamaCompletionRequest buildLlamaRequest() {
-    var promptTemplate = LlamaSettingsState.getInstance().getInfillPromptTemplate();
-    var prompt = promptTemplate.buildPrompt(details.getPrefix(), details.getSuffix());
+    InfillPromptTemplate promptTemplate = getLlamaInfillPromptTemplate();
+    String prompt = promptTemplate.buildPrompt(details.getPrefix(), details.getSuffix());
     return new LlamaCompletionRequest.Builder(prompt)
         .setN_predict(MAX_TOKENS)
         .setStream(true)
         .setTemperature(0.1)
         .setStop(promptTemplate.getStopTokens())
         .build();
   }
+
+  private InfillPromptTemplate getLlamaInfillPromptTemplate() {
+    var settings = LlamaSettingsState.getInstance();
+    if (!settings.isRunLocalServer()) {
+      return settings.getRemoteModelInfillPromptTemplate();
+    }
+    if (settings.isUseCustomModel()) {
+      return settings.getLocalModelInfillPromptTemplate();
+    }
+    return LlamaModel.findByHuggingFaceModel(settings.getHuggingFaceModel())
+        .getInfillPromptTemplate();
+  }
 }

src/main/java/ee/carlrobert/codegpt/completions/CompletionClientProvider.java

@@ -5,6 +5,7 @@
 import ee.carlrobert.codegpt.CodeGPTPlugin;
 import ee.carlrobert.codegpt.completions.you.YouUserManager;
 import ee.carlrobert.codegpt.credentials.AzureCredentialsManager;
+import ee.carlrobert.codegpt.credentials.LlamaCredentialsManager;
 import ee.carlrobert.codegpt.credentials.OpenAICredentialsManager;
 import ee.carlrobert.codegpt.settings.advanced.AdvancedSettingsState;
 import ee.carlrobert.codegpt.settings.state.AzureSettingsState;
@@ -13,6 +14,7 @@
 import ee.carlrobert.llm.client.azure.AzureClient;
 import ee.carlrobert.llm.client.azure.AzureCompletionRequestParams;
 import ee.carlrobert.llm.client.llama.LlamaClient;
+import ee.carlrobert.llm.client.llama.LlamaClient.Builder;
 import ee.carlrobert.llm.client.openai.OpenAIClient;
 import ee.carlrobert.llm.client.you.UTMParameters;
 import ee.carlrobert.llm.client.you.YouClient;
@@ -72,9 +74,17 @@ public static YouClient getYouClient() {
   }
 
   public static LlamaClient getLlamaClient() {
-    return new LlamaClient.Builder()
-        .setPort(LlamaSettingsState.getInstance().getServerPort())
-        .build(getDefaultClientBuilder());
+    LlamaSettingsState llamaSettingsState = LlamaSettingsState.getInstance();
+    Builder builder = new Builder()
+        .setPort(llamaSettingsState.getServerPort());
+    if (!llamaSettingsState.isRunLocalServer()) {
+      builder.setHost(llamaSettingsState.getBaseHost());
+      String apiKey = LlamaCredentialsManager.getInstance().getApiKey();
+      if (apiKey != null && !apiKey.isBlank()) {
+        builder.setApiKey(apiKey);
+      }
+    }
+    return builder.build(getDefaultClientBuilder());
   }
 
   private static OkHttpClient.Builder getDefaultClientBuilder() {

src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java

@@ -105,9 +105,14 @@ public LlamaCompletionRequest buildLlamaCompletionRequest(
       Message message,
       ConversationType conversationType) {
     var settings = LlamaSettingsState.getInstance();
-    var promptTemplate = settings.isUseCustomModel()
-        ? settings.getPromptTemplate()
-        : LlamaModel.findByHuggingFaceModel(settings.getHuggingFaceModel()).getPromptTemplate();
+    PromptTemplate promptTemplate;
+    if (settings.isRunLocalServer()) {
+      promptTemplate = settings.isUseCustomModel()
+          ? settings.getLocalModelPromptTemplate()
+          : LlamaModel.findByHuggingFaceModel(settings.getHuggingFaceModel()).getPromptTemplate();
+    } else {
+      promptTemplate = settings.getRemoteModelPromptTemplate();
+    }
 
     var systemPrompt = COMPLETION_SYSTEM_PROMPT;
     if (conversationType == ConversationType.FIX_COMPILE_ERRORS) {

src/main/java/ee/carlrobert/codegpt/completions/llama/LlamaModel.java

@@ -4,6 +4,7 @@
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
 
+import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate;
 import ee.carlrobert.codegpt.completions.HuggingFaceModel;
 import java.util.Collections;
 import java.util.List;
@@ -17,6 +18,7 @@ public enum LlamaModel {
           + "support for large input contexts, and zero-shot instruction following ability for "
           + "programming tasks.",
       PromptTemplate.LLAMA,
+      InfillPromptTemplate.LLAMA,
       List.of(
           HuggingFaceModel.CODE_LLAMA_7B_Q3,
           HuggingFaceModel.CODE_LLAMA_7B_Q4,
@@ -45,6 +47,7 @@ public enum LlamaModel {
           + "in both English and Chinese. It achieves state-of-the-art performance among "
           + "open-source code models on multiple programming languages and various benchmarks.",
       PromptTemplate.ALPACA,
+      InfillPromptTemplate.DEEPSEEK_CODER,
       List.of(
           HuggingFaceModel.DEEPSEEK_CODER_1_3B_Q3,
           HuggingFaceModel.DEEPSEEK_CODER_1_3B_Q4,
@@ -85,16 +88,27 @@ public enum LlamaModel {
   private final String label;
   private final String description;
   private final PromptTemplate promptTemplate;
+  private final InfillPromptTemplate infillPromptTemplate;
   private final List<HuggingFaceModel> huggingFaceModels;
 
   LlamaModel(
       String label,
       String description,
       PromptTemplate promptTemplate,
       List<HuggingFaceModel> huggingFaceModels) {
+    this(label, description, promptTemplate, null, huggingFaceModels);
+  }
+
+  LlamaModel(
+      String label,
+      String description,
+      PromptTemplate promptTemplate,
+      InfillPromptTemplate infillPromptTemplate,
+      List<HuggingFaceModel> huggingFaceModels) {
     this.label = label;
     this.description = description;
     this.promptTemplate = promptTemplate;
+    this.infillPromptTemplate = infillPromptTemplate;
     this.huggingFaceModels = huggingFaceModels;
   }
 
@@ -125,6 +139,10 @@ public PromptTemplate getPromptTemplate() {
     return promptTemplate;
   }
 
+  public InfillPromptTemplate getInfillPromptTemplate() {
+    return infillPromptTemplate;
+  }
+
   public List<HuggingFaceModel> getHuggingFaceModels() {
     return huggingFaceModels;
   }

src/main/java/ee/carlrobert/codegpt/completions/llama/LlamaServerAgent.java

@@ -39,19 +39,15 @@ public void startAgent(
       Runnable onSuccess,
       Runnable onServerTerminated) {
     ApplicationManager.getApplication().invokeLater(() -> {
-      if (!params.isUseCustomServer()) {
-        try {
-          serverProgressPanel.updateText(
-              CodeGPTBundle.get("llamaServerAgent.buildingProject.description"));
-          makeProcessHandler = new OSProcessHandler(getMakeCommandLinde());
-          makeProcessHandler.addProcessListener(
-              getMakeProcessListener(params, serverProgressPanel, onSuccess, onServerTerminated));
-          makeProcessHandler.startNotify();
-        } catch (ExecutionException e) {
-          throw new RuntimeException(e);
-        }
-      } else {
-        startServer(params, serverProgressPanel, onSuccess, onServerTerminated);
+      try {
+        serverProgressPanel.updateText(
+            CodeGPTBundle.get("llamaServerAgent.buildingProject.description"));
+        makeProcessHandler = new OSProcessHandler(getMakeCommandLinde());
+        makeProcessHandler.addProcessListener(
+            getMakeProcessListener(params, serverProgressPanel, onSuccess, onServerTerminated));
+        makeProcessHandler.startNotify();
+      } catch (ExecutionException e) {
+        throw new RuntimeException(e);
       }
     });
   }
@@ -83,31 +79,23 @@ public void onTextAvailable(@NotNull ProcessEvent event, @NotNull Key outputType
 
       @Override
       public void processTerminated(@NotNull ProcessEvent event) {
-        startServer(params, serverProgressPanel, onSuccess, onServerTerminated);
+        try {
+          LOG.info("Booting up llama server");
+
+          serverProgressPanel.updateText(
+              CodeGPTBundle.get("llamaServerAgent.serverBootup.description"));
+          startServerProcessHandler = new OSProcessHandler.Silent(getServerCommandLine(params));
+          startServerProcessHandler.addProcessListener(
+              getProcessListener(params.getPort(), onSuccess, onServerTerminated));
+          startServerProcessHandler.startNotify();
+        } catch (ExecutionException ex) {
+          LOG.error("Unable to start llama server", ex);
+          throw new RuntimeException(ex);
+        }
       }
     };
   }
 
-  private void startServer(
-      LlamaServerStartupParams params,
-      ServerProgressPanel serverProgressPanel,
-      Runnable onSuccess,
-      Runnable onServerTerminated) {
-    try {
-      LOG.info("Booting up llama server");
-
-      serverProgressPanel.updateText(
-          CodeGPTBundle.get("llamaServerAgent.serverBootup.description"));
-      startServerProcessHandler = new OSProcessHandler.Silent(getServerCommandLine(params));
-      startServerProcessHandler.addProcessListener(
-          getProcessListener(params.getPort(), onSuccess, onServerTerminated));
-      startServerProcessHandler.startNotify();
-    } catch (ExecutionException ex) {
-      LOG.error("Unable to start llama server", ex);
-      throw new RuntimeException(ex);
-    }
-  }
-
   private ProcessListener getProcessListener(
       int port,
       Runnable onSuccess,
@@ -164,8 +152,8 @@ private static GeneralCommandLine getMakeCommandLinde() {
 
   private GeneralCommandLine getServerCommandLine(LlamaServerStartupParams params) {
     GeneralCommandLine commandLine = new GeneralCommandLine().withCharset(StandardCharsets.UTF_8);
-    commandLine.setExePath("./" + params.getServerFileName());
-    commandLine.withWorkDirectory(params.getServerDirectory());
+    commandLine.setExePath("./server");
+    commandLine.withWorkDirectory(CodeGPTPlugin.getLlamaSourcePath());
     commandLine.addParameters(
         "-m", params.getModelPath(),
         "-c", String.valueOf(params.getContextLength()),

src/main/java/ee/carlrobert/codegpt/completions/llama/LlamaServerStartupParams.java

@@ -1,50 +1,28 @@
 package ee.carlrobert.codegpt.completions.llama;
 
-import java.io.File;
 import java.util.List;
 
 public class LlamaServerStartupParams {
 
-  private final String serverPath;
-  private final boolean useCustomServer;
   private final String modelPath;
   private final int contextLength;
   private final int threads;
   private final int port;
   private final List<String> additionalParameters;
 
   public LlamaServerStartupParams(
-      String serverPath,
-      boolean useCustomServer, String modelPath,
+      String modelPath,
       int contextLength,
       int threads,
       int port,
       List<String> additionalParameters) {
-    this.serverPath = serverPath;
-    this.useCustomServer = useCustomServer;
     this.modelPath = modelPath;
     this.contextLength = contextLength;
     this.threads = threads;
     this.port = port;
     this.additionalParameters = additionalParameters;
   }
 
-  public String getServerPath() {
-    return serverPath;
-  }
-
-  public String getServerFileName() {
-    return serverPath.substring(serverPath.lastIndexOf(File.separator) + 1);
-  }
-
-  public String getServerDirectory() {
-    return serverPath.substring(0, serverPath.lastIndexOf(File.separator) + 1);
-  }
-
-  public boolean isUseCustomServer() {
-    return useCustomServer;
-  }
-
   public String getModelPath() {
     return modelPath;
   }

src/main/java/ee/carlrobert/codegpt/credentials/LlamaCredentialsManager.java

@@ -0,0 +1,36 @@
+package ee.carlrobert.codegpt.credentials;
+
+import com.intellij.credentialStore.CredentialAttributes;
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.components.Service;
+import org.jetbrains.annotations.Nullable;
+
+@Service
+public final class LlamaCredentialsManager {
+
+  private static final CredentialAttributes llamaApiKeyCredentialAttributes =
+      CredentialsUtil.createCredentialAttributes("LLAMA_API_KEY");
+
+  private String llamaApiKey;
+
+  private LlamaCredentialsManager() {
+    llamaApiKey = CredentialsUtil.getPassword(llamaApiKeyCredentialAttributes);
+  }
+
+  public static LlamaCredentialsManager getInstance() {
+    return ApplicationManager.getApplication().getService(LlamaCredentialsManager.class);
+  }
+
+  public boolean isApiKeySet() {
+    return llamaApiKey != null && !llamaApiKey.isEmpty();
+  }
+
+  public @Nullable String getApiKey() {
+    return llamaApiKey;
+  }
+
+  public void setApiKey(String llamaApiKey) {
+    this.llamaApiKey = llamaApiKey;
+    CredentialsUtil.setPassword(llamaApiKeyCredentialAttributes, llamaApiKey);
+  }
+}

src/main/java/ee/carlrobert/codegpt/settings/SettingsConfigurable.java

@@ -6,6 +6,7 @@
 import ee.carlrobert.codegpt.CodeGPTBundle;
 import ee.carlrobert.codegpt.conversations.ConversationsState;
 import ee.carlrobert.codegpt.credentials.AzureCredentialsManager;
+import ee.carlrobert.codegpt.credentials.LlamaCredentialsManager;
 import ee.carlrobert.codegpt.credentials.OpenAICredentialsManager;
 import ee.carlrobert.codegpt.settings.state.AzureSettingsState;
 import ee.carlrobert.codegpt.settings.state.LlamaSettingsState;
@@ -75,6 +76,8 @@ public void apply() {
     AzureCredentialsManager.getInstance().setApiKey(serviceSelectionForm.getAzureOpenAIApiKey());
     AzureCredentialsManager.getInstance()
         .setAzureActiveDirectoryToken(serviceSelectionForm.getAzureActiveDirectoryToken());
+    LlamaCredentialsManager.getInstance()
+        .setApiKey(serviceSelectionForm.getLlamaServerPreferencesForm().getApiKey());
 
     var settings = SettingsState.getInstance();
     settings.setDisplayName(settingsComponent.getDisplayName());