Added lookup for models

codemakerai · Apr 3, 2024 · 1e4c989 · 1e4c989
1 parent 2d53acc
commit 1e4c989
Show file tree

Hide file tree

Showing 10 changed files with 140 additions and 4 deletions.
diff --git a/src/main/java/ai/codemaker/sdkv2/client/Client.java b/src/main/java/ai/codemaker/sdkv2/client/Client.java
@@ -14,6 +14,8 @@
 import ai.codemaker.sdkv2.client.model.CreateContextResponse;
 import ai.codemaker.sdkv2.client.model.DiscoverContextRequest;
 import ai.codemaker.sdkv2.client.model.DiscoverContextResponse;
+import ai.codemaker.sdkv2.client.model.ListModelsRequest;
+import ai.codemaker.sdkv2.client.model.ListModelsResponse;
 import ai.codemaker.sdkv2.client.model.PredictRequest;
 import ai.codemaker.sdkv2.client.model.PredictResponse;
 import ai.codemaker.sdkv2.client.model.ProcessRequest;
@@ -38,4 +40,6 @@ public interface Client {
     CreateContextResponse createContext(CreateContextRequest request);
 
     RegisterContextResponse registerContext(RegisterContextRequest request);
+
+    ListModelsResponse listModels(ListModelsRequest request);
 }
diff --git a/src/main/java/ai/codemaker/sdkv2/client/DefaultClient.java b/src/main/java/ai/codemaker/sdkv2/client/DefaultClient.java
@@ -16,7 +16,10 @@
 import ai.codemaker.sdkv2.client.model.DiscoverContextResponse;
 import ai.codemaker.sdkv2.client.model.Input;
 import ai.codemaker.sdkv2.client.model.Language;
+import ai.codemaker.sdkv2.client.model.ListModelsRequest;
+import ai.codemaker.sdkv2.client.model.ListModelsResponse;
 import ai.codemaker.sdkv2.client.model.Mode;
+import ai.codemaker.sdkv2.client.model.Model;
 import ai.codemaker.sdkv2.client.model.Modify;
 import ai.codemaker.sdkv2.client.model.Options;
 import ai.codemaker.sdkv2.client.model.Output;
@@ -156,6 +159,15 @@ public RegisterContextResponse registerContext(RegisterContextRequest request) {
         return createRegisterContextResponse(registerContextResponse);
     }
 
+    @Override
+    public ListModelsResponse listModels(ListModelsRequest request) {
+        final Codemakerai.ListModelsRequest listModelsRequest = createListModelsRequest(request);
+
+        final Codemakerai.ListModelsResponse listModelsResponse = doListModels(listModelsRequest);
+
+        return createListModelsResponse(listModelsResponse);
+    }
+
     private Codemakerai.DiscoverSourceContextResponse doDiscoverContext(Codemakerai.DiscoverSourceContextRequest request) {
         try {
             return client().discoverContext(request);
@@ -252,6 +264,18 @@ private Codemakerai.PredictResponse doPredict(Codemakerai.PredictRequest request
         }
     }
 
+    private Codemakerai.ListModelsResponse doListModels(Codemakerai.ListModelsRequest request) {
+        try {
+            return client().listModels(request);
+        } catch (StatusRuntimeException e) {
+            logger.error("Error calling service {} {}", e.getStatus().getCode(), e.getStatus().getDescription(), e);
+            if (e.getStatus().getCode() == Status.Code.PERMISSION_DENIED) {
+                throw new UnauthorizedException("Unauthorized request.");
+            }
+            throw e;
+        }
+    }
+
     private CodemakerServiceGrpc.CodemakerServiceBlockingStub client() {
         return client.withDeadlineAfter(DEFAULT_TIMEOUT_IN_MILLIS, TimeUnit.MILLISECONDS);
     }
@@ -382,6 +406,19 @@ private Codemakerai.PredictRequest createPredictRequest(PredictRequest request)
                 .build();
     }
 
+    private Codemakerai.ListModelsRequest createListModelsRequest(ListModelsRequest request) {
+        return Codemakerai.ListModelsRequest.newBuilder()
+                .build();
+    }
+
+    private ListModelsResponse createListModelsResponse(Codemakerai.ListModelsResponse listModelsResponse) {
+        final List<Model> models = listModelsResponse.getModelsList().stream()
+                .map(value -> new Model(value.getId(), value.getName()))
+                .toList();
+
+        return new ListModelsResponse(models);
+    }
+
     private Codemakerai.Input createInput(Input request) {
         return Codemakerai.Input.newBuilder()
                 .setSource(encodeInput(request))

diff --git a/src/main/java/ai/codemaker/sdkv2/client/model/ListModelsRequest.java b/src/main/java/ai/codemaker/sdkv2/client/model/ListModelsRequest.java
@@ -0,0 +1,8 @@
+/*
+ * Copyright 2023 CodeMaker AI Inc. All rights reserved.
+ */
+
+package ai.codemaker.sdkv2.client.model;
+
+public class ListModelsRequest {
+}
diff --git a/src/main/java/ai/codemaker/sdkv2/client/model/ListModelsResponse.java b/src/main/java/ai/codemaker/sdkv2/client/model/ListModelsResponse.java
@@ -0,0 +1,20 @@
+/*
+ * Copyright 2023 CodeMaker AI Inc. All rights reserved.
+ */
+
+package ai.codemaker.sdkv2.client.model;
+
+import java.util.List;
+
+public class ListModelsResponse {
+
+    private final List<Model> models;
+
+    public ListModelsResponse(List<Model> models) {
+        this.models = models;
+    }
+
+    public List<Model> getModels() {
+        return models;
+    }
+}
diff --git a/src/main/java/ai/codemaker/sdkv2/client/model/Model.java b/src/main/java/ai/codemaker/sdkv2/client/model/Model.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2023 CodeMaker AI Inc. All rights reserved.
+ */
+
+package ai.codemaker.sdkv2.client.model;
+
+public class Model {
+
+    private final String id;
+
+    private final String name;
+
+    public Model(String id, String name) {
+        this.id = id;
+        this.name = name;
+    }
+
+    public String getId() {
+        return id;
+    }
+
+    public String getName() {
+        return name;
+    }
+}
diff --git a/src/main/kotlin/ai/codemaker/jetbrains/service/CodeMakerService.kt b/src/main/kotlin/ai/codemaker/jetbrains/service/CodeMakerService.kt
@@ -150,6 +150,11 @@ class CodeMakerService(private val project: Project) {
         }
     }
 
+    fun listModels(): List<Model> {
+        val resp = client.listModels(createListModelsRequest())
+        return resp.models
+    }
+
     private fun process(mode: Mode, title: String, path: VirtualFile?, modify: Modify, codePath: String?, prompt: String? = null) {
         runInBackground(title) {
             try {
@@ -451,6 +456,10 @@ class CodeMakerService(private val project: Project) {
         )
     }
 
+    private fun createListModelsRequest(): ListModelsRequest? {
+        return ListModelsRequest()
+    }
+
     private fun isExtendedContextSupported(mode: Mode): Boolean {
         return mode == Mode.CODE
                 || mode == Mode.EDIT_CODE

diff --git a/src/main/kotlin/ai/codemaker/jetbrains/settings/AppSettingsComponent.kt b/src/main/kotlin/ai/codemaker/jetbrains/settings/AppSettingsComponent.kt
@@ -3,18 +3,24 @@
  */
 package ai.codemaker.jetbrains.settings
 
+import ai.codemaker.jetbrains.service.CodeMakerService
+import ai.codemaker.sdkv2.client.model.Model
 import com.intellij.ide.BrowserUtil
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.ui.components.*
 import com.intellij.util.ui.FormBuilder
+import javax.swing.DefaultComboBoxModel
+import javax.swing.JButton
 import javax.swing.JPanel
 
-class AppSettingsComponent {
+class AppSettingsComponent(project: Project) {
     private val defaultModel = "default"
 
     val panel: JPanel
     private val apiKeyText = JBTextField()
     private val modelCombo = ComboBox(arrayOf(defaultModel))
+    private val reloadModelsButton = JButton("Refresh")
     private val codeActionsEnabledCheck = JBCheckBox()
     private val autocompletionEnabledCheck = JBCheckBox()
     private val multilineAutocompletionEnabledCheck = JBCheckBox()
@@ -30,11 +36,23 @@ class AppSettingsComponent {
         }
         createAccountLabel.setExternalLinkIcon()
 
+        reloadModelsButton.addActionListener({
+            val service: CodeMakerService = project.getService(CodeMakerService::class.java)
+            val models = service.listModels()
+
+            val allModels = mutableListOf<String>()
+            allModels.add(defaultModel)
+            allModels.addAll(models.map { it.name })
+
+            modelCombo.model = DefaultComboBoxModel(allModels.toTypedArray())
+        })
+
         panel = FormBuilder.createFormBuilder()
                 .addComponent(createAccountLabel)
                 .addSeparator()
                 .addLabeledComponent(JBLabel("API Key: "), apiKeyText, 1, false)
                 .addLabeledComponent(JBLabel("Model: "), modelCombo, 1, false)
+                .addComponentToRightColumn(reloadModelsButton)
                 .addSeparator()
                 .addLabeledComponent(JBLabel("Enable autocompletion: "), autocompletionEnabledCheck, 1, false)
                 .addLabeledComponent(JBLabel("Enable multiline autocompletion: "), multilineAutocompletionEnabledCheck, 1, false)

diff --git a/src/main/kotlin/ai/codemaker/jetbrains/settings/AppSettingsConfigurable.kt b/src/main/kotlin/ai/codemaker/jetbrains/settings/AppSettingsConfigurable.kt
@@ -4,9 +4,10 @@
 package ai.codemaker.jetbrains.settings
 
 import com.intellij.openapi.options.Configurable
+import com.intellij.openapi.project.Project
 import javax.swing.JComponent
 
-class AppSettingsConfigurable : Configurable {
+class AppSettingsConfigurable(private val project: Project) : Configurable {
 
     private var settingsComponent: AppSettingsComponent? = null
 
@@ -15,7 +16,7 @@ class AppSettingsConfigurable : Configurable {
     }
 
     override fun createComponent(): JComponent? {
-        settingsComponent = AppSettingsComponent()
+        settingsComponent = AppSettingsComponent(project)
         return settingsComponent!!.panel
     }
 

diff --git a/src/main/proto/codemakerai.proto b/src/main/proto/codemakerai.proto
@@ -10,6 +10,7 @@ service CodemakerService {
   rpc DiscoverContext (DiscoverSourceContextRequest) returns (DiscoverSourceContextResponse);
   rpc CreateContext (CreateSourceContextRequest) returns (CreateSourceContextResponse);
   rpc RegisterContext (RegisterSourceContextRequest) returns (RegisterSourceContextResponse);
+  rpc ListModels (ListModelsRequest) returns (ListModelsResponse);
 }
 
 enum Mode {
@@ -79,6 +80,11 @@ message CodeSnippetContext {
   double score = 4;
 }
 
+message Model {
+  string id = 1;
+  string name = 2;
+}
+
 message ProcessOptions {
   Modify modify = 1;
   string codePath = 2;
@@ -181,4 +187,12 @@ message RegisterSourceContextRequest {
 
 message RegisterSourceContextResponse {
 
+}
+
+message ListModelsRequest {
+
+}
+
+message ListModelsResponse {
+  repeated Model models = 1;
 }
diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml
@@ -177,7 +177,7 @@
         <notificationGroup id="CodeMaker AI"
                            displayType="BALLOON"/>
 
-        <applicationConfigurable id="ai.codemaker.jetbrains.settings.CodeMakerSettings"
+        <projectConfigurable id="ai.codemaker.jetbrains.settings.CodeMakerSettings"
                                  parentId="tools"
                                  displayName="CodeMaker AI"
                                  instance="ai.codemaker.jetbrains.settings.AppSettingsConfigurable"/>