From 8991795b7353c3354f125395cd9e29cdba427b80 Mon Sep 17 00:00:00 2001 From: Daniel D'Avella Date: Thu, 11 Jul 2024 13:49:05 -0400 Subject: [PATCH 1/2] Skip LLM codemods when no service is available --- .../codemodder/codemods/DefaultCodemods.java | 1 + .../plugins/llm/LLMServiceModule.java | 18 ++++++++++------- .../codemodder/plugins/llm/OpenAIService.java | 20 +++++++++++++++++++ .../plugins/llm/SarifPluginLLMCodemod.java | 19 ++++++++++++++++++ ...ForBinaryVerificationAndFixingCodemod.java | 6 ++---- .../llm/SarifToLLMForMultiOutcomeCodemod.java | 8 +++----- 6 files changed, 56 insertions(+), 16 deletions(-) create mode 100644 plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java diff --git a/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java b/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java index 2dde62f81..7c1da5585 100644 --- a/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java +++ b/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java @@ -35,6 +35,7 @@ public static List> asList() { JDBCResourceLeakCodemod.class, JEXLInjectionCodemod.class, JSPScriptletXSSCodemod.class, + // LogFailedLoginCodemod.class, LimitReadlineCodemod.class, MavenSecureURLCodemod.class, OutputResourceLeakCodemod.class, diff --git a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/LLMServiceModule.java b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/LLMServiceModule.java index 3fe2ae10b..7ca180d01 100644 --- a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/LLMServiceModule.java +++ b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/LLMServiceModule.java @@ -1,6 +1,8 @@ package io.codemodder.plugins.llm; import com.google.inject.AbstractModule; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Provides configured LLM services. */ public final class LLMServiceModule extends AbstractModule { @@ -8,6 +10,7 @@ public final class LLMServiceModule extends AbstractModule { 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"; + private static final Logger logger = LoggerFactory.getLogger(LLMServiceModule.class); @Override protected void configure() { @@ -22,19 +25,20 @@ protected void configure() { + " must be set"); } if (azureOpenAIKey != null) { + logger.info("Using Azure OpenAI service with endpoint {}", azureOpenAIEndpoint); bind(OpenAIService.class) .toProvider(() -> OpenAIService.fromAzureOpenAI(azureOpenAIKey, azureOpenAIEndpoint)); return; } - bind(OpenAIService.class).toProvider(() -> OpenAIService.fromOpenAI(getOpenAIToken())); - } - - 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"); + if (openAIKey != null) { + logger.info("Using OpenAI service"); + bind(OpenAIService.class).toProvider(() -> OpenAIService.fromOpenAI(openAIKey)); + return; } - return openAIKey; + + logger.info("No LLM service available"); + bind(OpenAIService.class).toProvider(OpenAIService::noServiceAvailable); } } diff --git a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/OpenAIService.java b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/OpenAIService.java index 5477579d7..0fe4d63d8 100644 --- a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/OpenAIService.java +++ b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/OpenAIService.java @@ -21,6 +21,7 @@ public class OpenAIService { private final OpenAIClient api; private static final int TIMEOUT_SECONDS = 90; private final ModelMapper modelMapper; + private boolean serviceAvailable = true; private static OpenAIClientBuilder builder(final KeyCredential key) { HttpClientOptions clientOptions = new HttpClientOptions(); @@ -31,6 +32,12 @@ private static OpenAIClientBuilder builder(final KeyCredential key) { .credential(key); } + OpenAIService(final boolean serviceAvailable) { + this.serviceAvailable = serviceAvailable; + this.modelMapper = null; + this.api = null; + } + OpenAIService(final ModelMapper mapper, final KeyCredential key) { this.modelMapper = mapper; this.api = builder(key).buildClient(); @@ -66,6 +73,19 @@ public static OpenAIService fromAzureOpenAI(final String token, final String end Objects.requireNonNull(endpoint)); } + public static OpenAIService noServiceAvailable() { + return new OpenAIService(false); + } + + /** + * Returns whether the service is available. + * + * @return whether the service is available + */ + public boolean isServiceAvailable() { + return serviceAvailable; + } + /** * Gets the completion for the given messages. * diff --git a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java new file mode 100644 index 000000000..4c7461b69 --- /dev/null +++ b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java @@ -0,0 +1,19 @@ +package io.codemodder.plugins.llm; + +import io.codemodder.RuleSarif; +import io.codemodder.SarifPluginRawFileChanger; +import java.util.Objects; + +public abstract class SarifPluginLLMCodemod extends SarifPluginRawFileChanger { + protected final OpenAIService openAI; + + public SarifPluginLLMCodemod(RuleSarif sarif, final OpenAIService openAI) { + super(sarif); + this.openAI = Objects.requireNonNull(openAI); + } + + @Override + public boolean shouldRun() { + return openAI.isServiceAvailable(); + } +} diff --git a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForBinaryVerificationAndFixingCodemod.java b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForBinaryVerificationAndFixingCodemod.java index 3bdbcffb0..cda978cc2 100644 --- a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForBinaryVerificationAndFixingCodemod.java +++ b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForBinaryVerificationAndFixingCodemod.java @@ -30,15 +30,13 @@ * */ public abstract class SarifToLLMForBinaryVerificationAndFixingCodemod - extends SarifPluginRawFileChanger { + extends SarifPluginLLMCodemod { - private final OpenAIService openAI; private final Model model; protected SarifToLLMForBinaryVerificationAndFixingCodemod( final RuleSarif sarif, final OpenAIService openAI, final Model model) { - super(sarif); - this.openAI = Objects.requireNonNull(openAI); + super(sarif, openAI); this.model = Objects.requireNonNull(model); } diff --git a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForMultiOutcomeCodemod.java b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForMultiOutcomeCodemod.java index c64ddbb61..26168b01a 100644 --- a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForMultiOutcomeCodemod.java +++ b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifToLLMForMultiOutcomeCodemod.java @@ -38,11 +38,10 @@ *

To accomplish that, we need the analysis to "bucket" the code into one of the above * categories. */ -public abstract class SarifToLLMForMultiOutcomeCodemod extends SarifPluginRawFileChanger { +public abstract class SarifToLLMForMultiOutcomeCodemod extends SarifPluginLLMCodemod { private static final Logger logger = LoggerFactory.getLogger(SarifToLLMForMultiOutcomeCodemod.class); - private final OpenAIService openAI; private final List remediationOutcomes; private final Model categorizationModel; private final Model codeChangingModel; @@ -65,8 +64,7 @@ protected SarifToLLMForMultiOutcomeCodemod( final List remediationOutcomes, final Model categorizationModel, final Model codeChangingModel) { - super(sarif); - this.openAI = Objects.requireNonNull(openAI); + super(sarif, openAI); this.remediationOutcomes = Objects.requireNonNull(remediationOutcomes); if (remediationOutcomes.size() < 2) { throw new IllegalArgumentException("must have 2+ remediation outcome"); @@ -78,7 +76,7 @@ protected SarifToLLMForMultiOutcomeCodemod( @Override public CodemodFileScanningResult onFileFound( final CodemodInvocationContext context, final List results) { - logger.info("processing: {}", context.path()); + logger.debug("processing: {}", context.path()); List changes = new ArrayList<>(); for (Result result : results) { From cfb7b01cde83f09f746acfff9dca7d429a3874ba Mon Sep 17 00:00:00 2001 From: Daniel D'Avella Date: Thu, 11 Jul 2024 15:16:10 -0400 Subject: [PATCH 2/2] Address code review comments --- .../main/java/io/codemodder/codemods/DefaultCodemods.java | 1 - .../io/codemodder/plugins/llm/SarifPluginLLMCodemod.java | 7 +++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java b/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java index 7c1da5585..2dde62f81 100644 --- a/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java +++ b/core-codemods/src/main/java/io/codemodder/codemods/DefaultCodemods.java @@ -35,7 +35,6 @@ public static List> asList() { JDBCResourceLeakCodemod.class, JEXLInjectionCodemod.class, JSPScriptletXSSCodemod.class, - // LogFailedLoginCodemod.class, LimitReadlineCodemod.class, MavenSecureURLCodemod.class, OutputResourceLeakCodemod.class, diff --git a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java index 4c7461b69..089ecc51f 100644 --- a/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java +++ b/plugins/codemodder-plugin-llm/src/main/java/io/codemodder/plugins/llm/SarifPluginLLMCodemod.java @@ -4,6 +4,7 @@ import io.codemodder.SarifPluginRawFileChanger; import java.util.Objects; +/** A base class for LLM codemods that process SARIF and use the OpenAI service. */ public abstract class SarifPluginLLMCodemod extends SarifPluginRawFileChanger { protected final OpenAIService openAI; @@ -12,6 +13,12 @@ public SarifPluginLLMCodemod(RuleSarif sarif, final OpenAIService openAI) { this.openAI = Objects.requireNonNull(openAI); } + /** + * Indicates whether the codemod should run. + * + *

Subclasses can override this method to add additional hecks but should call + * super.shouldRun() to ensure the OpenAI service is available. + */ @Override public boolean shouldRun() { return openAI.isServiceAvailable();