Enable compatibility with Azure OpenAI (#415)

Author: drdavella

core-codemods/src/main/java/io/codemodder/codemods/LogFailedLoginCodemod.java

@@ -67,8 +67,8 @@ public LogFailedLoginCodemod(
                   When no logger is in scope, the new code emits a log message to the console.
                   """
                     .replace('\n', ' '))),
-        StandardModel.GPT_4O,
-        StandardModel.GPT_4);
+        StandardModel.GPT_4O_2024_05_13,
+        StandardModel.GPT_4_0613);
   }
 
   @Override

core-codemods/src/main/java/io/codemodder/codemods/SensitiveDataLoggingCodemod.java

@@ -1,5 +1,6 @@
 package io.codemodder.codemods;
 
+import com.azure.ai.openai.models.ChatRequestUserMessage;
 import com.contrastsecurity.sarif.PhysicalLocation;
 import com.contrastsecurity.sarif.Result;
 import com.contrastsecurity.sarif.Run;
@@ -8,10 +9,10 @@
 import com.fasterxml.jackson.databind.ObjectReader;
 import com.github.javaparser.ast.CompilationUnit;
 import com.github.javaparser.ast.stmt.Statement;
-import com.theokanning.openai.completion.chat.*;
 import io.codemodder.*;
 import io.codemodder.javaparser.JavaParserChanger;
 import io.codemodder.plugins.llm.OpenAIService;
+import io.codemodder.plugins.llm.StandardModel;
 import io.codemodder.providers.sarif.semgrep.SemgrepScan;
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -98,21 +99,10 @@ private SensitivityAndFixAnalysis performSensitivityAnalysis(
               """
             .formatted(startLine, codeSnippet);
 
-    ChatCompletionRequest request =
-        ChatCompletionRequest.builder()
-            .temperature(0D)
-            .model("gpt-4o-2024-05-13")
-            .n(1)
-            .messages(List.of(new ChatMessage(ChatMessageRole.USER.value(), prompt)))
-            .build();
-    ChatCompletionResult completion = service.createChatCompletion(request);
-    ChatCompletionChoice response = completion.getChoices().get(0);
-    String responseText = response.getMessage().getContent();
-    if (responseText.startsWith("```json") && responseText.endsWith("```")) {
-      responseText =
-          responseText.substring("```json".length(), responseText.length() - "```".length());
-    }
-    return reader.readValue(responseText);
+    return service.getResponseForPrompt(
+        List.of(new ChatRequestUserMessage(prompt)),
+        StandardModel.GPT_4O_2024_05_13,
+        SensitivityAndFixAnalysisDTO.class);
   }
 
   /**

framework/codemodder-base/build.gradle.kts

@@ -29,7 +29,7 @@ dependencies {
     api(libs.javaparser.symbolsolver.model)
     api(libs.javadiff)
     api(libs.jtokkit)
-    api(libs.openai.service)
+    api("com.azure:azure-ai-openai:1.0.0-beta.10")
     api("io.github.classgraph:classgraph:4.8.160")
 
     implementation(libs.tuples)

plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/EnvironmentBasedModelMapper.java

@@ -0,0 +1,22 @@
+package io.codemodder.plugins.llm;
+
+import java.util.HashMap;
+
+/** Mapper that maps models to their deployment names based on environment variables. */
+final class EnvironmentBasedModelMapper implements ModelMapper {
+  private static final String DEPLOYMENT_TEMPLATE = "CODEMODDER_AZURE_OPENAI_%s_DEPLOYMENT";
+
+  private final HashMap<Model, String> map = new HashMap<>();
+
+  EnvironmentBasedModelMapper() {
+    for (Model m : StandardModel.values()) {
+      final var deployment = System.getenv(String.format(DEPLOYMENT_TEMPLATE, m));
+      map.put(m, deployment == null ? m.id() : deployment);
+    }
+  }
+
+  @Override
+  public String getModelName(Model model) {
+    return map.get(model);
+  }
+}

plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/LLMServiceModule.java

@@ -5,18 +5,36 @@
 /** Provides configured LLM services. */
 public final class LLMServiceModule extends AbstractModule {
 
-  private static final String TOKEN_NAME = "CODEMODDER_OPENAI_API_KEY";
+  private static final String OPENAI_KEY_NAME = "CODEMODDER_OPENAI_API_KEY";
+  private static final String AZURE_OPENAI_KEY_NAME = "CODEMODDER_AZURE_OPENAI_API_KEY";
+  private static final String AZURE_OPENAI_ENDPOINT = "CODEMODDER_AZURE_OPENAI_ENDPOINT";
 
   @Override
   protected void configure() {
-    bind(OpenAIService.class).toProvider(() -> new OpenAIService(getToken()));
+    final var azureOpenAIKey = System.getenv(AZURE_OPENAI_KEY_NAME);
+    final var azureOpenAIEndpoint = System.getenv(AZURE_OPENAI_ENDPOINT);
+    if ((azureOpenAIEndpoint == null) != (azureOpenAIKey == null)) {
+      throw new IllegalArgumentException(
+          "Both or neither of "
+              + AZURE_OPENAI_KEY_NAME
+              + " and "
+              + AZURE_OPENAI_ENDPOINT
+              + " must be set");
+    }
+    if (azureOpenAIKey != null) {
+      bind(OpenAIService.class)
+          .toProvider(() -> OpenAIService.fromAzureOpenAI(azureOpenAIKey, azureOpenAIEndpoint));
+      return;
+    }
+
+    bind(OpenAIService.class).toProvider(() -> OpenAIService.fromOpenAI(getOpenAIToken()));
   }
 
-  private String getToken() {
-    String token = System.getenv(TOKEN_NAME);
-    if (token == null) {
-      throw new IllegalArgumentException(TOKEN_NAME + " environment variable must be set");
+  private String getOpenAIToken() {
+    final var openAIKey = System.getenv(OPENAI_KEY_NAME);
+    if (openAIKey == null) {
+      throw new IllegalArgumentException(OPENAI_KEY_NAME + " environment variable must be set");
     }
-    return token;
+    return openAIKey;
   }
 }

plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/Model.java

@@ -1,6 +1,5 @@
 package io.codemodder.plugins.llm;
 
-import com.theokanning.openai.completion.chat.ChatMessage;
 import java.util.List;
 
 /**
@@ -26,5 +25,5 @@ public interface Model {
    * @param messages the list of messages for which to estimate token usage
    * @return estimated tokens that would be consumed by the model
    */
-  int tokens(List<ChatMessage> messages);
+  int tokens(List<String> messages);
 }

plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/ModelMapper.java

@@ -0,0 +1,12 @@
+package io.codemodder.plugins.llm;
+
+/** Maps models to their deployment names. */
+interface ModelMapper {
+  /**
+   * Maps the given model to its deployment name.
+   *
+   * @param model the model to map
+   * @return the deployment name of the model
+   */
+  String getModelName(Model model);
+}

plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/OpenAIService.java

@@ -1,69 +1,118 @@
 package io.codemodder.plugins.llm;
 
-import com.theokanning.openai.client.OpenAiApi;
-import com.theokanning.openai.completion.chat.ChatCompletionRequest;
-import com.theokanning.openai.completion.chat.ChatCompletionResult;
-import com.theokanning.openai.service.OpenAiService;
-import io.reactivex.Flowable;
-import io.reactivex.functions.Function;
-import java.net.SocketTimeoutException;
+import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIClientBuilder;
+import com.azure.ai.openai.models.ChatCompletions;
+import com.azure.ai.openai.models.ChatCompletionsJsonResponseFormat;
+import com.azure.ai.openai.models.ChatCompletionsOptions;
+import com.azure.ai.openai.models.ChatRequestMessage;
+import com.azure.core.credential.AzureKeyCredential;
+import com.azure.core.credential.KeyCredential;
+import com.azure.core.http.policy.RetryPolicy;
+import com.azure.core.util.HttpClientOptions;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
 import java.time.Duration;
-import java.util.concurrent.TimeUnit;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import retrofit2.HttpException;
+import java.util.List;
+import java.util.Objects;
 
-/**
- * A custom service class to call the {@link OpenAiApi}, since the out-of-the box {@link
- * OpenAiService} <a href="https://github.com/TheoKanning/openai-java/issues/189">does not support
- * automatic retries</a>.
- */
+/** A custom service class to wrap the {@link OpenAIClient} */
 public class OpenAIService {
-  private final OpenAiApi api;
+  private final OpenAIClient api;
+  private static final int TIMEOUT_SECONDS = 90;
+  private final ModelMapper modelMapper;
 
-  public OpenAIService(final String token) {
-    this.api = OpenAiService.buildApi(token, Duration.ofSeconds(90));
+  private static OpenAIClientBuilder builder(final KeyCredential key) {
+    HttpClientOptions clientOptions = new HttpClientOptions();
+    clientOptions.setReadTimeout(Duration.ofSeconds(TIMEOUT_SECONDS));
+    return new OpenAIClientBuilder()
+        .retryPolicy(new RetryPolicy())
+        .clientOptions(clientOptions)
+        .credential(key);
   }
 
-  public ChatCompletionResult createChatCompletion(final ChatCompletionRequest request) {
-    return this.api
-        .createChatCompletion(request)
-        .retryWhen(new OpenAIRetryStrategy())
-        .blockingGet();
+  OpenAIService(final ModelMapper mapper, final KeyCredential key) {
+    this.modelMapper = mapper;
+    this.api = builder(key).buildClient();
+  }
+
+  OpenAIService(final ModelMapper mapper, final KeyCredential key, final String endpoint) {
+    this.modelMapper = mapper;
+    this.api = builder(key).endpoint(endpoint).buildClient();
   }
-}
 
-/**
- * When there is a retryable error from OpenAI -- either a timeout or a retryable <a
- * href="https://platform.openai.com/docs/guides/error-codes/api-errors">error code</a> -- this will
- * retry the request up to 3 times, with a delay of {@code 1 * retryCount} seconds between retries.
- */
-class OpenAIRetryStrategy implements Function<Flowable<? extends Throwable>, Flowable<Object>> {
-  private static final int MAX_RETRY_COUNT = 3;
-  private static final Logger logger = LoggerFactory.getLogger(OpenAIRetryStrategy.class);
+  /**
+   * Creates a new {@link OpenAIService} instance with the given OpenAI token.
+   *
+   * @param token the token to use
+   * @return the new instance
+   */
+  public static OpenAIService fromOpenAI(final String token) {
+    return new OpenAIService(
+        new EnvironmentBasedModelMapper(), new KeyCredential(Objects.requireNonNull(token)));
+  }
 
-  private int retryCount = 0;
+  /**
+   * Creates a new {@link OpenAIService} instance with the given Azure OpenAI token and endpoint.
+   *
+   * @param token the token to use
+   * @param endpoint the endpoint to use
+   * @return the new instance
+   */
+  public static OpenAIService fromAzureOpenAI(final String token, final String endpoint) {
+    return new OpenAIService(
+        new EnvironmentBasedModelMapper(),
+        new AzureKeyCredential(Objects.requireNonNull(token)),
+        Objects.requireNonNull(endpoint));
+  }
 
-  @Override
-  public Flowable<Object> apply(final Flowable<? extends Throwable> flowable) {
-    return flowable.flatMap(
-        e -> {
-          if (++retryCount <= MAX_RETRY_COUNT && isRetryable(e)) {
-            logger.warn("retrying after {}s: {}", retryCount, e);
-            return Flowable.timer(retryCount, TimeUnit.SECONDS);
-          } else {
-            return Flowable.error(e);
-          }
-        });
+  /**
+   * Gets the completion for the given messages.
+   *
+   * @param messages the messages
+   * @param modelOrDeploymentName the model or deployment name
+   * @return the completion
+   */
+  public String getJSONCompletion(
+      final List<ChatRequestMessage> messages, final Model modelOrDeploymentName) {
+    ChatCompletionsOptions options =
+        new ChatCompletionsOptions(messages)
+            .setTemperature(0D)
+            .setN(1)
+            .setResponseFormat(new ChatCompletionsJsonResponseFormat());
+    final var modelName = modelMapper.getModelName(modelOrDeploymentName);
+    ChatCompletions completions = this.api.getChatCompletions(modelName, options);
+    return completions.getChoices().get(0).getMessage().getContent().trim();
   }
 
-  private boolean isRetryable(final Throwable e) {
-    if (e instanceof SocketTimeoutException) {
-      return true;
-    } else if (e instanceof HttpException) {
-      int code = ((HttpException) e).code();
-      return code == 429 || code == 500 || code == 503;
+  /**
+   * Returns an object of the given type based on the completion for the given messages.
+   *
+   * @param messages the messages
+   * @param modelName the model name
+   * @return the completion
+   */
+  public <T> T getResponseForPrompt(
+      final List<ChatRequestMessage> messages, final Model modelName, final Class<T> responseType)
+      throws IOException {
+    String json = getJSONCompletion(messages, modelName);
+
+    // we've seen this with turbo/lesser models
+    if (json.startsWith("```json") && json.endsWith("```")) {
+      json = json.substring(7, json.length() - 3);
+    }
+
+    // try to deserialize the content as is
+    ObjectMapper mapper = new ObjectMapper();
+    try {
+      return mapper.readValue(json, responseType);
+    } catch (IOException e) {
+      int firstBorder = json.indexOf("```json");
+      int lastBorder = json.lastIndexOf("```");
+      if (firstBorder != -1 && lastBorder != -1 && lastBorder > firstBorder) {
+        json = json.substring(firstBorder + 7, lastBorder);
+      }
+      return mapper.readValue(json, responseType);
     }
-    return false;
   }
 }