Add Azure OpenAI Service support (#96)

* Add Azure OpenAI Service support

* Bump version number

* Bump openai-client dependency to 1.0.12

* Improve settings panel logic, bump openai-client

---------

Co-authored-by: Shalak <github.author@shalak.eu.org>

build.gradle.kts

@@ -4,7 +4,7 @@ plugins {
 }
 
 group = "ee.carlrobert"
-version = "1.10.4"
+version = "1.10.5"
 
 repositories {
   mavenCentral()
@@ -22,7 +22,7 @@ dependencies {
   implementation("com.fifesoft:rsyntaxtextarea:3.3.2")
   implementation("com.vladsch.flexmark:flexmark-all:0.64.0")
   implementation("org.apache.commons:commons-text:1.10.0")
-  implementation("ee.carlrobert:openai-client:1.0.9")
+  implementation("ee.carlrobert:openai-client:1.0.14")
   implementation("com.knuddels:jtokkit:0.2.0")
 }
 

src/main/java/ee/carlrobert/codegpt/client/ClientProvider.java

@@ -4,6 +4,10 @@ import ee.carlrobert.codegpt.state.settings.SettingsState;
 import ee.carlrobert.codegpt.state.settings.advanced.AdvancedSettingsState;
 import ee.carlrobert.openai.client.OpenAIClient;
 import ee.carlrobert.openai.client.ProxyAuthenticator;
+import ee.carlrobert.openai.client.azure.AzureChatCompletionClient;
+import ee.carlrobert.openai.client.azure.AzureClientRequestParams;
+import ee.carlrobert.openai.client.azure.AzureTextCompletionClient;
+import ee.carlrobert.openai.client.completion.CompletionClient;
 import ee.carlrobert.openai.client.completion.chat.ChatCompletionClient;
 import ee.carlrobert.openai.client.completion.text.TextCompletionClient;
 import ee.carlrobert.openai.client.dashboard.DashboardClient;
@@ -16,22 +20,35 @@ public class ClientProvider {
   public static DashboardClient getDashboardClient() {
     return getClientBuilder().buildDashboardClient();
   }
-
-  public static ChatCompletionClient getChatCompletionClient() {
+  public static CompletionClient getChatCompletionClient(SettingsState settings) {
+    if (settings.useAzure) {
+      return getClientBuilder().buildAzureChatCompletionClient(getAzureRequestParams());
+    }
     return getClientBuilder().buildChatCompletionClient();
   }
 
-  public static TextCompletionClient getTextCompletionClient() {
+  public static CompletionClient getTextCompletionClient(SettingsState settings) {
+    if (settings.useAzure) {
+      return getClientBuilder().buildAzureTextCompletionClient(getAzureRequestParams());
+    }
     return getClientBuilder().buildTextCompletionClient();
   }
 
+  private static AzureClientRequestParams getAzureRequestParams() {
+    var settings = SettingsState.getInstance();
+    return new AzureClientRequestParams(settings.resourceName, settings.deploymentId, settings.apiVersion);
+  }
+
   private static OpenAIClient.Builder getClientBuilder() {
     var settings = SettingsState.getInstance();
     var builder = new OpenAIClient.Builder(settings.apiKey) // TODO: ENV var?
-        .setOrganization(settings.organization)
         .setConnectTimeout(60L, TimeUnit.SECONDS)
         .setReadTimeout(30L, TimeUnit.SECONDS);
 
+    if (!settings.useAzure) {
+      builder.setOrganization(settings.organization);
+    }
+
     var advancedSettings = AdvancedSettingsState.getInstance();
     var proxyHost = advancedSettings.proxyHost;
     var proxyPort = advancedSettings.proxyPort;

src/main/java/ee/carlrobert/codegpt/client/RequestHandler.java

@@ -60,11 +60,13 @@ public class RequestHandler implements ActionListener {
   private EventSource startCall(Message message, EventListener eventListener) {
     var settings = SettingsState.getInstance();
     var requestProvider = new CompletionRequestProvider(message.getPrompt(), conversation);
+
     if (settings.isChatCompletionOptionSelected) {
-      return ClientProvider.getChatCompletionClient().stream(
-          requestProvider.buildChatCompletionRequest(settings.chatCompletionBaseModel), eventListener);
+      return ClientProvider.getChatCompletionClient(settings).stream(
+              requestProvider.buildChatCompletionRequest(settings.chatCompletionBaseModel),
+              eventListener);
     }
-    return ClientProvider.getTextCompletionClient().stream(
-        requestProvider.buildTextCompletionRequest(settings.textCompletionBaseModel), eventListener);
-  }
+    return ClientProvider.getTextCompletionClient(settings).stream(
+            requestProvider.buildTextCompletionRequest(settings.textCompletionBaseModel),
+            eventListener);  }
 }

src/main/java/ee/carlrobert/codegpt/state/settings/SettingsComponent.java

@@ -5,7 +5,6 @@ import com.intellij.ui.TitledSeparator;
 import com.intellij.ui.components.JBCheckBox;
 import com.intellij.ui.components.JBRadioButton;
 import com.intellij.ui.components.JBTextField;
-import com.intellij.util.ui.CheckBox;
 import com.intellij.util.ui.FormBuilder;
 import com.intellij.util.ui.JBUI;
 import com.intellij.util.ui.UI;
@@ -14,20 +13,23 @@ import ee.carlrobert.openai.client.completion.CompletionModel;
 import ee.carlrobert.openai.client.completion.chat.ChatCompletionModel;
 import ee.carlrobert.openai.client.completion.text.TextCompletionModel;
 import java.awt.Desktop;
+import java.awt.event.ItemEvent;
 import java.io.IOException;
 import java.net.URISyntaxException;
-import javax.swing.ButtonGroup;
-import javax.swing.JComponent;
-import javax.swing.JPanel;
+import javax.swing.*;
 import javax.swing.event.HyperlinkEvent;
 
 public class SettingsComponent {
 
   private final JPanel mainPanel;
+  private final JBTextField resourceNameField;
+  private final JBTextField deploymentIdField;
+  private final JBTextField apiVersionField;
   private final JBTextField apiKeyField;
   private final JBTextField organizationField;
   private final JBTextField displayNameField;
   private final JBCheckBox useOpenAIAccountNameCheckBox;
+  private final JBCheckBox useAzureCheckbox;
   private final ComboBox<CompletionModel> chatCompletionBaseModelComboBox;
   private final ComboBox<CompletionModel> textCompletionBaseModelComboBox;
   private final JBRadioButton useChatCompletionRadioButton;
@@ -35,6 +37,10 @@ public class SettingsComponent {
 
   public SettingsComponent(SettingsState settings) {
     apiKeyField = new JBTextField(settings.apiKey, 40);
+    useAzureCheckbox = new JBCheckBox("Use Azure OpenAI service API", false);
+    resourceNameField = new JBTextField(settings.resourceName, 40);
+    deploymentIdField = new JBTextField(settings.resourceName, 40);
+    apiVersionField = new JBTextField(settings.resourceName, 40);
     organizationField = new JBTextField(settings.organization, 40);
     displayNameField = new JBTextField(settings.displayName, 20);
     useOpenAIAccountNameCheckBox = new JBCheckBox("Use OpenAI account name", true);
@@ -86,6 +92,37 @@ public class SettingsComponent {
     apiKeyField.setText(apiKey);
   }
 
+  public void setUseAzureCheckbox(boolean selected) {
+    useAzureCheckbox.setSelected(selected);
+  }
+
+  public boolean isUseAzure() {
+    return useAzureCheckbox.isSelected();
+  }
+
+  public String getResourceName() {
+    return resourceNameField.getText();
+  }
+
+  public void setResourceName(String resourceName) {
+    resourceNameField.setText(resourceName);
+  }
+
+  public String getDeploymentId() {
+    return deploymentIdField.getText();
+  }
+
+  public void setDeploymentId(String deploymentId) {
+    deploymentIdField.setText(deploymentId);
+  }
+  public String getApiVersion() {
+    return apiVersionField.getText();
+  }
+
+  public void setApiVersionField(String apiVersion) {
+    apiVersionField.setText(apiVersion);
+  }
+
   public String getOrganization() {
     return organizationField.getText();
   }
@@ -168,9 +205,13 @@ public class SettingsComponent {
         textCompletionBaseModelComboBox, "Model:", false);
     textCompletionModelsPanel.setBorder(JBUI.Borders.emptyLeft(16));
 
+    var azureRelatedFieldsPanel = createAzureServicePanel();
+
     var panel = FormBuilder.createFormBuilder()
         .addComponent(FormBuilder.createFormBuilder()
             .addComponent(apiKeyFieldPanel)
+            .addComponent(useAzureCheckbox)
+            .addComponent(azureRelatedFieldsPanel)
             .addComponent(organizationFieldPanel)
             .addVerticalGap(8)
             .addComponent(displayNameFieldPanel)
@@ -188,9 +229,49 @@ public class SettingsComponent {
             .getPanel())
         .getPanel();
     panel.setBorder(JBUI.Borders.emptyLeft(16));
+
+    useAzureCheckbox.addItemListener(e -> {
+      azureRelatedFieldsPanel.setVisible(e.getStateChange() == ItemEvent.SELECTED);
+      organizationFieldPanel.setVisible(e.getStateChange() != ItemEvent.SELECTED);
+    });
+
     return panel;
   }
 
+  private JPanel createAzureServicePanel() {
+    JPanel azureRelatedFieldsPanel = new JPanel();
+    var resourceNameFieldPanel = UI.PanelFactory.panel(resourceNameField)
+            .withLabel("Resource name:")
+            .resizeX(false)
+            .withComment(
+                    "Azure OpenAI Service resource name")
+            .createPanel();
+    var deploymentIdFieldPanel = UI.PanelFactory.panel(deploymentIdField)
+            .withLabel("Deployment ID:")
+            .resizeX(false)
+            .withComment(
+                    "Azure OpenAI Service deployment ID")
+            .createPanel();
+    var apiVersionFieldPanel = UI.PanelFactory.panel(apiVersionField)
+            .withLabel("API version:")
+            .resizeX(false)
+            .withComment(
+                    "API version to be used for Azure OpenAI Service")
+            .createPanel();
+    azureRelatedFieldsPanel.setLayout(new BoxLayout(azureRelatedFieldsPanel, BoxLayout.Y_AXIS));
+    azureRelatedFieldsPanel.setBorder(JBUI.Borders.emptyLeft(16));
+
+    azureRelatedFieldsPanel.add(resourceNameFieldPanel);
+    azureRelatedFieldsPanel.add(deploymentIdFieldPanel);
+    azureRelatedFieldsPanel.add(apiVersionFieldPanel);
+    SwingUtils.setEqualLabelWidths(deploymentIdFieldPanel, resourceNameFieldPanel);
+    SwingUtils.setEqualLabelWidths(apiVersionFieldPanel, resourceNameFieldPanel);
+
+    azureRelatedFieldsPanel.setVisible(false);
+
+    return azureRelatedFieldsPanel;
+  }
+
   private void registerButtons() {
     ButtonGroup completionButtonGroup = new ButtonGroup();
     completionButtonGroup.add(useChatCompletionRadioButton);

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

@@ -37,6 +37,10 @@ public class SettingsConfigurable implements Configurable {
   public boolean isModified() {
     var settings = SettingsState.getInstance();
     return !settingsComponent.getApiKey().equals(settings.apiKey) ||
+        settingsComponent.isUseAzure() != settings.useAzure ||
+        !settingsComponent.getResourceName().equals(settings.resourceName) ||
+        !settingsComponent.getDeploymentId().equals(settings.deploymentId) ||
+        !settingsComponent.getApiVersion().equals(settings.apiVersion) ||
         !settingsComponent.getOrganization().equals(settings.organization) ||
         !settingsComponent.getDisplayName().equals(settings.displayName) ||
         settingsComponent.isUseOpenAIAccountName() != settings.useOpenAIAccountName ||
@@ -70,6 +74,10 @@ public class SettingsConfigurable implements Configurable {
     }
 
     settings.apiKey = settingsComponent.getApiKey();
+    settings.useAzure = settingsComponent.isUseAzure();
+    settings.resourceName = settingsComponent.getResourceName();
+    settings.deploymentId = settingsComponent.getDeploymentId();
+    settings.apiVersion = settingsComponent.getApiVersion();
     settings.organization = settingsComponent.getOrganization();
     settings.displayName = settingsComponent.getDisplayName();
     settings.useOpenAIAccountName = settingsComponent.isUseOpenAIAccountName();
@@ -85,6 +93,10 @@ public class SettingsConfigurable implements Configurable {
     settingsComponent.setUseChatCompletionSelected(settings.isChatCompletionOptionSelected);
     settingsComponent.setUseTextCompletionSelected(settings.isTextCompletionOptionSelected);
     settingsComponent.setApiKey(settings.apiKey);
+    settingsComponent.setUseAzureCheckbox(settings.useAzure);
+    settingsComponent.setResourceName(settings.resourceName);
+    settingsComponent.setDeploymentId(settings.deploymentId);
+    settingsComponent.setApiVersionField(settings.apiVersion);
     settingsComponent.setOrganization(settings.organization);
     settingsComponent.setDisplayName(settings.displayName);
     settingsComponent.setUseOpenAIAccountNameCheckBox(settings.useOpenAIAccountName);

src/main/java/ee/carlrobert/codegpt/state/settings/SettingsState.java

@@ -19,6 +19,10 @@ import org.jetbrains.annotations.Nullable;
 public class SettingsState implements PersistentStateComponent<SettingsState> {
 
   public String apiKey = "";
+  public boolean useAzure;
+  public String resourceName = "";
+  public String deploymentId = "";
+  public String apiVersion = "";
   public String organization = "";
   public String displayName = getDisplayName();
   public boolean useOpenAIAccountName = true;