From 252db27577c7696db570e03316fba47c11bb585e Mon Sep 17 00:00:00 2001
From: Theo Kanning <TheoKanning@users.noreply.github.com>
Date: Fri, 19 Aug 2022 11:05:12 -0500
Subject: [PATCH] Update EditResult and fix the Edit test (#26)

OpenAI's example curl used an invalid engine, and the api returned a confusing error message.
Now everything works with text-davinci-edit-001
---
 .../theokanning/openai/edit/EditRequest.java  | 13 ++++++---
 .../theokanning/openai/edit/EditResult.java   |  5 ++++
 .../theokanning/openai/edit/EditUsage.java    | 27 +++++++++++++++++++
 .../java/com/theokanning/openai/EditTest.java |  8 +++---
 4 files changed, 44 insertions(+), 9 deletions(-)
 create mode 100644 api/src/main/java/com/theokanning/openai/edit/EditUsage.java

diff --git a/api/src/main/java/com/theokanning/openai/edit/EditRequest.java b/api/src/main/java/com/theokanning/openai/edit/EditRequest.java
index 9829cfd..5063cb4 100644
--- a/api/src/main/java/com/theokanning/openai/edit/EditRequest.java
+++ b/api/src/main/java/com/theokanning/openai/edit/EditRequest.java
@@ -1,9 +1,7 @@
 package com.theokanning.openai.edit;
 
-import lombok.AllArgsConstructor;
-import lombok.Builder;
-import lombok.Data;
-import lombok.NoArgsConstructor;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.*;
 
 /**
  * Given a prompt and an instruction, OpenAi will return an edited version of the prompt
@@ -25,8 +23,14 @@ public class EditRequest {
      * The instruction that tells the model how to edit the prompt.
      * For example, "Fix the spelling mistakes"
      */
+    @NonNull
     String instruction;
 
+    /**
+     * How many edits to generate for the input and instruction.
+     */
+    Integer n;
+
     /**
      * What sampling temperature to use. Higher values means the model will take more risks. Try 0.9 for more creative applications, and 0 (argmax sampling) for ones with a well-defined answer.
      *
@@ -39,5 +43,6 @@ public class EditRequest {
      *
      * We generally recommend altering this or {@link EditRequest#temperature} but not both.
      */
+    @JsonProperty("top_p")
     Double topP;
 }
diff --git a/api/src/main/java/com/theokanning/openai/edit/EditResult.java b/api/src/main/java/com/theokanning/openai/edit/EditResult.java
index b2b3b54..1e4b08c 100644
--- a/api/src/main/java/com/theokanning/openai/edit/EditResult.java
+++ b/api/src/main/java/com/theokanning/openai/edit/EditResult.java
@@ -26,4 +26,9 @@ public class EditResult {
      * A list of generated edits.
      */
     List<EditChoice> choices;
+
+    /**
+     * The API usage for this request
+     */
+    EditUsage usage;
 }
diff --git a/api/src/main/java/com/theokanning/openai/edit/EditUsage.java b/api/src/main/java/com/theokanning/openai/edit/EditUsage.java
new file mode 100644
index 0000000..1d08f5f
--- /dev/null
+++ b/api/src/main/java/com/theokanning/openai/edit/EditUsage.java
@@ -0,0 +1,27 @@
+package com.theokanning.openai.edit;
+
+import lombok.Data;
+
+/**
+ * An object containing the API usage for an edit request
+ *
+ * https://beta.openai.com/docs/api-reference/edits/create
+ */
+@Data
+public class EditUsage {
+
+    /**
+     * The number of prompt tokens consumed.
+     */
+    String promptTokens;
+
+    /**
+     * The number of completion tokens consumed.
+     */
+    String completionTokens;
+
+    /**
+     * The number of total tokens consumed.
+     */
+    String totalTokens;
+}
diff --git a/client/src/test/java/com/theokanning/openai/EditTest.java b/client/src/test/java/com/theokanning/openai/EditTest.java
index 39836de..556e52d 100644
--- a/client/src/test/java/com/theokanning/openai/EditTest.java
+++ b/client/src/test/java/com/theokanning/openai/EditTest.java
@@ -2,12 +2,10 @@ package com.theokanning.openai;
 
 import com.theokanning.openai.edit.EditRequest;
 import com.theokanning.openai.edit.EditResult;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 
-@Disabled // disabled until edit example CURL works
 public class EditTest {
 
     String token = System.getenv("OPENAI_TOKEN");
@@ -20,8 +18,8 @@ public class EditTest {
                 .instruction("Fix the spelling mistakes")
                 .build();
 
-        EditResult result = service.createEdit("text-ada-001", request);
+        EditResult result = service.createEdit("text-davinci-edit-001", request);
 
-        assertEquals("What day of the week is it?", result.getChoices().get(0).getText());
+        assertNotNull(result.getChoices().get(0).getText());
     }
 }