From 5e14d4f62bb17315260e03b49a20b6371e15e4ad Mon Sep 17 00:00:00 2001
From: Theo Kanning <TheoKanning@users.noreply.github.com>
Date: Sun, 4 Dec 2022 12:37:28 -0600
Subject: [PATCH] Add usage data to Completion and Embedding APIs (#39)

Also changed EditResult to use the new shared object
---
 .../java/com/theokanning/openai/Usage.java    | 24 +++++++++++++++++++
 .../openai/completion/CompletionResult.java   |  6 +++++
 .../theokanning/openai/edit/EditResult.java   |  3 ++-
 .../theokanning/openai/edit/EditUsage.java    |  3 +++
 .../openai/embedding/EmbeddingResult.java     |  6 +++++
 5 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 api/src/main/java/com/theokanning/openai/Usage.java

diff --git a/api/src/main/java/com/theokanning/openai/Usage.java b/api/src/main/java/com/theokanning/openai/Usage.java
new file mode 100644
index 00000000..448dc18c
--- /dev/null
+++ b/api/src/main/java/com/theokanning/openai/Usage.java
@@ -0,0 +1,24 @@
+package com.theokanning.openai;
+
+import lombok.Data;
+
+/**
+ * The OpenAI resources used by a request
+ */
+@Data
+public class Usage {
+    /**
+     * The number of prompt tokens used.
+     */
+    long promptTokens;
+
+    /**
+     * The number of completion tokens used.
+     */
+    long completionTokens;
+
+    /**
+     * The number of total tokens used
+     */
+    long totalTokens;
+}
diff --git a/api/src/main/java/com/theokanning/openai/completion/CompletionResult.java b/api/src/main/java/com/theokanning/openai/completion/CompletionResult.java
index 7ba9fdd3..b43e4204 100644
--- a/api/src/main/java/com/theokanning/openai/completion/CompletionResult.java
+++ b/api/src/main/java/com/theokanning/openai/completion/CompletionResult.java
@@ -1,5 +1,6 @@
 package com.theokanning.openai.completion;
 
+import com.theokanning.openai.Usage;
 import lombok.Data;
 
 import java.util.List;
@@ -35,4 +36,9 @@ public class CompletionResult {
      * A list of generated completions.
      */
     List<CompletionChoice> choices;
+
+    /**
+     * The API usage for this request
+     */
+    Usage usage;
 }
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 9973e096..415ba749 100644
--- a/api/src/main/java/com/theokanning/openai/edit/EditResult.java
+++ b/api/src/main/java/com/theokanning/openai/edit/EditResult.java
@@ -1,5 +1,6 @@
 package com.theokanning.openai.edit;
 
+import com.theokanning.openai.Usage;
 import lombok.Data;
 
 import java.util.List;
@@ -30,5 +31,5 @@ public class EditResult {
     /**
      * The API usage for this request
      */
-    public EditUsage usage;
+    public Usage 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
index 1d08f5f3..7223890f 100644
--- a/api/src/main/java/com/theokanning/openai/edit/EditUsage.java
+++ b/api/src/main/java/com/theokanning/openai/edit/EditUsage.java
@@ -5,9 +5,12 @@
 /**
  * An object containing the API usage for an edit request
  *
+ * Deprecated, use {@link com.theokanning.openai.Usage} instead
+ *
  * https://beta.openai.com/docs/api-reference/edits/create
  */
 @Data
+@Deprecated
 public class EditUsage {
 
     /**
diff --git a/api/src/main/java/com/theokanning/openai/embedding/EmbeddingResult.java b/api/src/main/java/com/theokanning/openai/embedding/EmbeddingResult.java
index 6669717c..58d9b9c3 100644
--- a/api/src/main/java/com/theokanning/openai/embedding/EmbeddingResult.java
+++ b/api/src/main/java/com/theokanning/openai/embedding/EmbeddingResult.java
@@ -1,5 +1,6 @@
 package com.theokanning.openai.embedding;
 
+import com.theokanning.openai.Usage;
 import lombok.Data;
 
 import java.util.List;
@@ -26,4 +27,9 @@ public class EmbeddingResult {
      * A list of the calculated embeddings
      */
     List<Embedding> data;
+
+    /**
+     * The API usage for this request
+     */
+    Usage usage;
 }