Merge pull request #19 from profiq/azure-support

Azure support

build.gradle.kts

@@ -4,7 +4,7 @@ plugins {
 }
 
 group = "com.profiq"
-version = "0.1.6"
+version = "0.1.7"
 
 repositories {
     mavenCentral()
@@ -18,7 +18,8 @@ dependencies {
 // Configure Gradle IntelliJ Plugin
 // Read more: https://plugins.jetbrains.com/docs/intellij/tools-gradle-intellij-plugin.html
 intellij {
-    version.set("2024.1")
+    version.set("PY-2024.1")
+    type.set("PY")
     plugins.set(listOf())
 }
 

src/main/java/com/profiq/codexor/DocstringFormatter.java

@@ -0,0 +1,53 @@
+package com.profiq.codexor;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * Takes in a response from OpenAI API and formats it into a proper Python docstring.
+ */
+public class DocstringFormatter {
+
+    private final String openaiResponse;
+    private String[] docstringLines = null;
+
+    public DocstringFormatter(String openaiResponse) {
+        this.openaiResponse = openaiResponse;
+    }
+
+    public String docstringWithIndentation(String indentation) {
+        String[] docstringLines = getDocstringLines();
+
+        StringBuilder docstring = new StringBuilder();
+        for (String line : docstringLines) {
+            docstring.append(indentation).append(line).append("\n");
+        }
+
+        return docstring.toString();
+    }
+
+    private String[] getDocstringLines() {
+        if (docstringLines == null) {
+            Pattern docstringRegex = Pattern.compile("(([ \t]*)\"\"\".*\"\"\")", Pattern.DOTALL);
+            Matcher docstringMatcher = docstringRegex.matcher(openaiResponse);
+
+            if (!docstringMatcher.find()) {
+                throw new IllegalArgumentException("No docstring found in the OpenAI response.");
+            }
+
+            String docstring = docstringMatcher.group(1);
+            String indentation = docstringMatcher.group(2);
+            String[] docstringLinesRaw = docstring.split("\n");
+
+            docstringLines = new String[docstringLinesRaw.length];
+            int substringIdx = 0;
+
+            for (int i = 0; i < docstringLinesRaw.length; i++) {
+                substringIdx = Math.min(docstringLinesRaw[i].length(), indentation.length());
+                docstringLines[i] = docstringLinesRaw[i].substring(substringIdx);
+            }
+        }
+
+        return docstringLines;
+    }
+}

src/main/java/com/profiq/codexor/GenerateDocumentationAction.java

@@ -75,9 +75,8 @@ public void run(@NotNull ProgressIndicator indicator) {
                         if (response.statusCode() == 200) {
                             var responseParsed = (new Gson()).fromJson(response.body(), Response.class);
                             var responseText = responseParsed.getChoices()[0].getMessage().getContent();
-                            String docstring = parseDocstring(responseText);
                             var app = ApplicationManager.getApplication();
-                            app.invokeLater(() -> showEditor(e, docstring));
+                            app.invokeLater(() -> showEditor(e, responseText));
                         } else {
                             indicator.cancel();
                             showError(response.body());
@@ -111,35 +110,45 @@ public void run(@NotNull ProgressIndicator indicator) {
         }
     }
 
-    private static HttpRequest buildRequest(String apiKey, String model, String prompt, String code) {
-        var messages = new Message[]{new Message("user", prompt + "```" + code + "```")};
+    private HttpRequest buildRequest(String apiKey, String model, String prompt, String code) {
+        var endpointURL = getEndpointURL();
+        var isAzure = endpointURL.contains("azure");
+
+        if (isAzure) {
+            endpointURL = endpointURL + "openai/deployments/" + model + "/chat/completions?api-version=2024-02-01";
+        }
+
+        var messages = new Message[]{new Message(prompt + "```" + code + "```")};
         var requestBody = new Request(model, messages, 0);
         var requestJson = new Gson().toJson(requestBody);
 
-        return HttpRequest.newBuilder(URI.create("https://api.openai.com/v1/chat/completions"))
-            .header("Authorization", "Bearer " + apiKey)
-            .header("Content-Type", "application/json")
+        var requestBuilder = HttpRequest.newBuilder(URI.create(endpointURL));
+
+        if (isAzure) {
+            requestBuilder.header("api-key", apiKey);
+        } else {
+            requestBuilder.header("Authorization", "Bearer " + apiKey);
+        }
+
+        return requestBuilder.header("Content-Type", "application/json")
             .POST(HttpRequest.BodyPublishers.ofString(requestJson))
             .build();
     }
 
-    private String parseDocstring(String responseText) {
-        Pattern pattern = Pattern.compile(":\n(\\s*\"\"\".*\"\"\")", Pattern.DOTALL);
-        Matcher matcher = pattern.matcher(responseText);
-        if (matcher.find()) {
-            String docstring = matcher.group(1);
-            String[] docstringSplit = docstring.split("\"\"\"");
-            docstring = docstringSplit[0] + "\"\"\"" + docstringSplit[1] + "\"\"\"";
-            return docstring;
+    private String getEndpointURL() {
+        String endpointInput = PropertiesComponent.getInstance().getValue(SettingsConfigurable.ENDPOINT_SETTINGS_KEY);
+
+        if (endpointInput == null || endpointInput.isEmpty()) {
+            return "https://api.openai.com/v1/chat/completions";
         } else {
-            return "";
+            return endpointInput;
         }
     }
 
     private String getApiKey() throws IOException {
         String apiKeyInput = PropertiesComponent.getInstance().getValue(SettingsConfigurable.API_KEY_SETTING_KEY);
 
-        if (apiKeyInput == null || apiKeyInput.length() == 0) {
+        if (apiKeyInput == null || apiKeyInput.isEmpty()) {
             apiKeyInput = JOptionPane.showInputDialog("Enter your Open AI API key:");
             PropertiesComponent.getInstance().setValue(SettingsConfigurable.API_KEY_SETTING_KEY, apiKeyInput);
         }
@@ -175,7 +184,7 @@ private void showError(String text) {
         JOptionPane.showMessageDialog(null, text, "Error", JOptionPane.ERROR_MESSAGE);
     }
 
-    private void showEditor(AnActionEvent e, String docstring) {
+    private void showEditor(AnActionEvent e, String responseText) {
         var resultWindow = new JFrame();
         var virtualFile = e.getData(PlatformDataKeys.VIRTUAL_FILE);
         var fileType = FileTypeManager.getInstance().getFileTypeByFile(virtualFile);
@@ -185,7 +194,8 @@ private void showEditor(AnActionEvent e, String docstring) {
             return;
         }
 
-        var newDocument = EditorFactory.getInstance().createDocument(applyIndentation(docstring, ""));
+        DocstringFormatter docstringFormatter = new DocstringFormatter(responseText);
+        var newDocument = EditorFactory.getInstance().createDocument(docstringFormatter.docstringWithIndentation(""));
         var editor = EditorFactory.getInstance().createEditor(newDocument, null, fileType, false);
         editor.getContentComponent().setPreferredSize(new Dimension(1000, 800));
         var confirmBtn = new JButton("Insert");
@@ -196,7 +206,7 @@ private void showEditor(AnActionEvent e, String docstring) {
         int lineEndOffset = mainDocument.getLineEndOffset(caretPosition.line);
         int nextLineStartOffset = mainDocument.getLineStartOffset(caretPosition.line + 1);
         int offset = Math.min(lineEndOffset + 1, nextLineStartOffset);
-        String formattedDocstring = applyIndentation(docstring, correctIndentation);
+        String formattedDocstring = docstringFormatter.docstringWithIndentation(correctIndentation);
 
         confirmBtn.addActionListener(actionEvent -> {
             var app = ApplicationManager.getApplication();
@@ -215,31 +225,6 @@ private void showEditor(AnActionEvent e, String docstring) {
         resultWindow.setLocationRelativeTo(null);
     }
 
-    private String applyIndentation(String docstring, String indentation) {
-        Pattern pattern = Pattern.compile("^(\\s+)");
-        Matcher matcher = pattern.matcher(docstring);
-
-        if (matcher.find()) {
-            String leadingWhitespace = matcher.group(1);
-            int minIndent = leadingWhitespace.length();
-            String[] lines = docstring.split("\n");
-            StringBuilder unindented = new StringBuilder();
-
-            for (String line : lines) {
-                unindented.append(indentation);
-
-                if (line.length() >= minIndent) {
-                    unindented.append(line.substring(minIndent));
-                }
-                unindented.append("\n");
-            }
-
-            return unindented.toString();
-        }
-
-        return docstring;
-    }
-
     private String determineCorrectIndentation(String code) {
         Pattern pattern = Pattern.compile("\n(\\s+)");
         Matcher matcher = pattern.matcher(code);
@@ -250,4 +235,4 @@ private String determineCorrectIndentation(String code) {
 
         return "";
     }
-}
+}

src/main/java/com/profiq/codexor/Message.java

@@ -2,11 +2,10 @@
 
 public class Message {
 
-    private String role;
+    public String role = "user";
     private String content;
 
-    public Message(String role, String content) {
-        this.role = role;
+    public Message(String content) {
         this.content = content;
     }
 

src/main/java/com/profiq/codexor/Response.java

@@ -7,5 +7,4 @@ public class Response {
     Choice[] getChoices() {
         return choices;
     }
-
 }