Merge pull request #68 from scijava/chatgpt2

make OpenAI API key optional and prompt prefix, model name configurable
scijava · Nov 28, 2023 · d57a9a5 · d57a9a5
2 parents c5e501b + 7dd7aba
commit d57a9a5
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 7 deletions.
diff --git a/src/main/java/org/scijava/ui/swing/script/OpenAIOptions.java b/src/main/java/org/scijava/ui/swing/script/OpenAIOptions.java
@@ -41,14 +41,42 @@
 @Plugin(type = OptionsPlugin.class, menuPath = "Edit>Options>OpenAI...")
 public class OpenAIOptions extends OptionsPlugin {
 
-	@Parameter(label = "OpenAI key")
+	@Parameter(label = "OpenAI key", required = false)
 	private String openAIKey;
 
+	@Parameter(label = "OpenAI Model name")
+	private String modelName = "gpt-3.5-turbo-0613";
+
+	@Parameter(label = "Prompt prefix (added before custom prompt)", style = "text area")
+	private String promptPrefix = "Write code in {programming_language}.\n" +
+			"Write concise and high quality code for ImageJ/Fiji.\n" +
+			"Put minimal comments explaining what the code does.\n" +
+			"The code should do the following:\n" +
+			"{custom_prompt}";
+
 	public String getOpenAIKey() {
 		return openAIKey;
 	}
 
 	public void setOpenAIKey(final String openAIKey) {
 		this.openAIKey = openAIKey;
 	}
+
+	public String getPromptPrefix() {
+		return promptPrefix;
+	}
+
+	public void setPromptPrefix(final String promptPrefix) {
+		this.promptPrefix = promptPrefix;
+	}
+
+	public String getModelName() {
+		return modelName;
+	}
+
+	public void setModelName(final String modelName) {
+		this.modelName = modelName;
+	}
+
+
 }
diff --git a/src/main/java/org/scijava/ui/swing/script/TextEditor.java b/src/main/java/org/scijava/ui/swing/script/TextEditor.java
@@ -3250,11 +3250,9 @@ public void askChatGPTtoGenerateCode() {
 
 			// setup default prompt
 			String prompt =
-					"Write code in " + getCurrentLanguage().getLanguageName() + ".\n" +
-					"Write concise and high quality code for ImageJ/Fiji.\n" +
-					"Put minimal comments explaining what the code does.\n" +
-					"The code should do the following:\n" +
-					getTextArea().getSelectedText();
+					promptPrefix()
+							.replace("{programming_language}",  getCurrentLanguage().getLanguageName() )
+							.replace("{custom_prompt}", getTextArea().getSelectedText());
 
 			String answer = askChatGPT(prompt);
 
@@ -3289,7 +3287,7 @@ private String askChatGPT(String text) {
 
 		ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest
 				.builder()
-				.model("gpt-3.5-turbo-0613")
+				.model(modelName())
 				.messages(messages).build();
 
 		ChatMessage responseMessage = service.createChatCompletion(chatCompletionRequest).getChoices().get(0).getMessage();
@@ -3312,6 +3310,31 @@ private String apiKey() {
 		return System.getenv("OPENAI_API_KEY");
 	}
 
+	private String modelName() {
+		if (optionsService != null) {
+			final OpenAIOptions openAIOptions =
+					optionsService.getOptions(OpenAIOptions.class);
+			if (openAIOptions != null) {
+				final String key = openAIOptions.getModelName();
+				if (key != null && !key.isEmpty()) return key;
+			}
+		}
+		return null;
+	}
+
+	private String promptPrefix() {
+		if (optionsService != null) {
+			final OpenAIOptions openAIOptions =
+					optionsService.getOptions(OpenAIOptions.class);
+			if (openAIOptions != null) {
+				final String promptPrefix = openAIOptions.getPromptPrefix();
+				if (promptPrefix != null && !promptPrefix.isEmpty()) return promptPrefix;
+			}
+		}
+		return "";
+	}
+
+
 	public void extractSourceJar(final File file) {
 		try {
 			final FileFunctions functions = new FileFunctions(this);