Add CompletionModel class

StefanBratanov · Apr 2, 2024 · 3788acc · 3788acc
1 parent 092e5b0
commit 3788acc
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -35,7 +35,7 @@ OpenAI openAI = OpenAI.newBuilder(System.getenv("OPENAI_API_KEY")).build();
 
 ChatClient chatClient = openAI.chatClient();
 CreateChatCompletionRequest createChatCompletionRequest = CreateChatCompletionRequest.newBuilder()
-    .model("gpt-3.5-turbo")
+    .model(CompletionModel.GPT_3_5_TURBO)
     .message(ChatMessage.userMessage("Who won the world series in 2020?"))
     .build();
 ChatCompletion chatCompletion = chatClient.createChatCompletion(createChatCompletionRequest);
@@ -46,7 +46,7 @@ ChatCompletion chatCompletion = chatClient.createChatCompletion(createChatComple
 | API                                                                       | Status |
 |---------------------------------------------------------------------------|:------:|
 | [Audio](https://platform.openai.com/docs/api-reference/audio)             |   ✔️   |
-| [Chat](https://platform.openai.com/docs/api-reference/chat)               | ️  ✔️  |
+| [Chat](https://platform.openai.com/docs/api-reference/chat)               |   ✔️   |
 | [Embeddings](https://platform.openai.com/docs/api-reference/embeddings)   |   ✔️   |
 | [Fine-tuning](https://platform.openai.com/docs/api-reference/fine-tuning) |   ✔️   |
 | [Files](https://platform.openai.com/docs/api-reference/files)             |   ✔️   |
@@ -59,7 +59,7 @@ ChatCompletion chatCompletion = chatClient.createChatCompletion(createChatComple
 | API                                                                     | Status |
 |-------------------------------------------------------------------------|:------:|
 | [Assistants](https://platform.openai.com/docs/api-reference/assistants) |   ✔️   |
-| [Threads](https://platform.openai.com/docs/api-reference/threads)       | ️  ✔️  |
+| [Threads](https://platform.openai.com/docs/api-reference/threads)       |   ✔️   |
 | [Messages](https://platform.openai.com/docs/api-reference/messages)     |   ✔️   |
 | [Runs](https://platform.openai.com/docs/api-reference/runs)             |   ✔️   |
 
@@ -92,7 +92,7 @@ OpenAI openAI = OpenAI.newBuilder(System.getenv("OPENAI_API_KEY"))
 ```java
 ChatClient chatClient = openAI.chatClient();
 CreateChatCompletionRequest request = CreateChatCompletionRequest.newBuilder()
-    .model("gpt-3.5-turbo")
+    .model(CompletionModel.GPT_3_5_TURBO)
     .message(ChatMessage.userMessage("Who won the world series in 2020?"))
     .build();
 CompletableFuture<ChatCompletion> chatCompletion = chatClient.createChatCompletionAsync(request);

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/CompletionModel.java b/src/main/java/io/github/stefanbratanov/jvm/openai/CompletionModel.java
@@ -0,0 +1,32 @@
+package io.github.stefanbratanov.jvm.openai;
+
+import com.fasterxml.jackson.annotation.JsonValue;
+
+/**
+ * Represents the latest models that can be used to generate completions.
+ *
+ * <p>Note that this does not correspond to a static version and may change over time.
+ *
+ * @see <a href="https://platform.openai.com/docs/models/continuous-model-upgrades">Continuous model
+ * upgrades - OpenAI API</a>
+ */
+public enum CompletionModel {
+
+    GPT_4("gpt-4"),
+    GPT_4_TURBO_PREVIEW("gpt-4-turbo-preview"),
+    GPT_4_VISION_PREVIEW("gpt-4-vision-preview"),
+    GPT_4_32K("gpt-4-32k"),
+    GPT_3_5_TURBO("gpt-3.5-turbo"),
+    GPT_3_5_TURBO_16K("gpt-3.5-turbo-16k");
+
+    private final String id;
+
+    CompletionModel(String modelId) {
+        this.id = modelId;
+    }
+
+    @JsonValue
+    public String getId() {
+        return this.id;
+    }
+}
diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/CreateChatCompletionRequest.java b/src/main/java/io/github/stefanbratanov/jvm/openai/CreateChatCompletionRequest.java
@@ -41,7 +41,7 @@ public static Builder newBuilder() {
 
   public static class Builder {
 
-    private static final String DEFAULT_MODEL = "gpt-3.5-turbo";
+    private static final String DEFAULT_MODEL = CompletionModel.GPT_3_5_TURBO.getId();
 
     private final List<ChatMessage> messages = new LinkedList<>();
 
@@ -88,6 +88,14 @@ public Builder model(String model) {
       return this;
     }
 
+    /**
+     * @param model {@link CompletionModel} to use
+     */
+    public Builder model(CompletionModel model) {
+      this.model = model.getId();
+      return this;
+    }
+
     /**
      * @param frequencyPenalty Number between -2.0 and 2.0. Positive values penalize new tokens
      *     based on their existing frequency in the text so far, decreasing the model's likelihood