From 5d74cfde5a9b3005c6d43bbd7c8d9521dfed8dad Mon Sep 17 00:00:00 2001
From: Theo Kanning <theokanning@gmail.com>
Date: Mon, 14 Sep 2020 21:56:31 -0500
Subject: [PATCH] Add Completion api

Might add a builder or something to make setting specific parameters easier
---
 .../main/java/example/OpenAiApiExample.java   | 12 +++++++++--
 openai/src/main/java/openai/OpenAiApi.java    |  7 +++++++
 .../src/main/java/openai/OpenAiService.java   |  6 ++++++
 .../openai/completion/CompletionChoice.java   | 11 ++++++++++
 .../openai/completion/CompletionRequest.java  | 21 +++++++++++++++++++
 .../openai/completion/CompletionResult.java   | 14 +++++++++++++
 6 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 openai/src/main/java/openai/completion/CompletionChoice.java
 create mode 100644 openai/src/main/java/openai/completion/CompletionRequest.java
 create mode 100644 openai/src/main/java/openai/completion/CompletionResult.java

diff --git a/example/src/main/java/example/OpenAiApiExample.java b/example/src/main/java/example/OpenAiApiExample.java
index b0b1d3b..2f14546 100644
--- a/example/src/main/java/example/OpenAiApiExample.java
+++ b/example/src/main/java/example/OpenAiApiExample.java
@@ -1,6 +1,7 @@
 package example;
 
 import openai.OpenAiService;
+import openai.completion.CompletionRequest;
 import openai.engine.Engine;
 
 class OpenAiApiExample {
@@ -8,11 +9,18 @@ class OpenAiApiExample {
         String token = System.getenv("OPENAI_TOKEN");
         OpenAiService service = new OpenAiService(token);
 
-        System.out.println("Getting available engines");
+        System.out.println("Getting available engines...");
         service.getEngines().forEach(System.out::println);
 
-        System.out.println("Getting ada engine");
+        System.out.println("Getting ada engine...");
         Engine ada = service.getEngine("ada");
         System.out.println(ada);
+
+        System.out.println("Creating completion...");
+        CompletionRequest completionRequest = new CompletionRequest();
+        completionRequest.setPrompt("Somebody once told me the world is gonna roll me");
+        completionRequest.setEcho(true);
+        service.createCompletion("ada", completionRequest).getChoices().forEach(System.out::println);
+
     }
 }
\ No newline at end of file
diff --git a/openai/src/main/java/openai/OpenAiApi.java b/openai/src/main/java/openai/OpenAiApi.java
index 90c6f28..9a38e19 100644
--- a/openai/src/main/java/openai/OpenAiApi.java
+++ b/openai/src/main/java/openai/OpenAiApi.java
@@ -1,8 +1,12 @@
 package openai;
 
 import io.reactivex.Single;
+import openai.completion.CompletionRequest;
+import openai.completion.CompletionResult;
 import openai.engine.Engine;
+import retrofit2.http.Body;
 import retrofit2.http.GET;
+import retrofit2.http.POST;
 import retrofit2.http.Path;
 
 public interface OpenAiApi {
@@ -12,4 +16,7 @@ public interface OpenAiApi {
 
     @GET("/v1/engines/{engine_id}")
     Single<Engine> getEngine(@Path("engine_id") String engineId);
+
+    @POST("/v1/engines/{engine_id}/completions")
+    Single<CompletionResult> createCompletion(@Path("engine_id") String engineId, @Body CompletionRequest request);
 }
diff --git a/openai/src/main/java/openai/OpenAiService.java b/openai/src/main/java/openai/OpenAiService.java
index 612172d..3cc4bcd 100644
--- a/openai/src/main/java/openai/OpenAiService.java
+++ b/openai/src/main/java/openai/OpenAiService.java
@@ -6,6 +6,8 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.PropertyNamingStrategy;
 import okhttp3.ConnectionPool;
 import okhttp3.OkHttpClient;
+import openai.completion.CompletionRequest;
+import openai.completion.CompletionResult;
 import openai.engine.Engine;
 import retrofit2.Retrofit;
 import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
@@ -46,4 +48,8 @@ public class OpenAiService {
     public Engine getEngine(String engineId) {
         return api.getEngine(engineId).blockingGet();
     }
+
+    public CompletionResult createCompletion(String engineId, CompletionRequest request) {
+        return api.createCompletion(engineId, request).blockingGet();
+    }
 }
diff --git a/openai/src/main/java/openai/completion/CompletionChoice.java b/openai/src/main/java/openai/completion/CompletionChoice.java
new file mode 100644
index 0000000..64c8f50
--- /dev/null
+++ b/openai/src/main/java/openai/completion/CompletionChoice.java
@@ -0,0 +1,11 @@
+package openai.completion;
+
+import lombok.Data;
+
+@Data
+public class CompletionChoice {
+    String text;
+    Integer index;
+    // todo add logprobs
+    String finish_reason;
+}
diff --git a/openai/src/main/java/openai/completion/CompletionRequest.java b/openai/src/main/java/openai/completion/CompletionRequest.java
new file mode 100644
index 0000000..1591321
--- /dev/null
+++ b/openai/src/main/java/openai/completion/CompletionRequest.java
@@ -0,0 +1,21 @@
+package openai.completion;
+
+import lombok.Data;
+
+import java.util.List;
+
+@Data
+public class CompletionRequest {
+    String prompt;
+    Integer maxTokens;
+    Double temperature;
+    Double topP;
+    Integer n;
+    Boolean stream;
+    Integer logprobs;
+    Boolean echo;
+    List<String> stop; //todo test this
+    Double presencePenalty;
+    Double frequencyPenalty;
+    Integer bestOf;
+}
diff --git a/openai/src/main/java/openai/completion/CompletionResult.java b/openai/src/main/java/openai/completion/CompletionResult.java
new file mode 100644
index 0000000..7c1df97
--- /dev/null
+++ b/openai/src/main/java/openai/completion/CompletionResult.java
@@ -0,0 +1,14 @@
+package openai.completion;
+
+import lombok.Data;
+
+import java.util.List;
+
+@Data
+public class CompletionResult {
+    String id;
+    String object;
+    long created;
+    String model;
+    List<CompletionChoice> choices;
+}