From 936a004495d7cb70d8861ec9d30ea953094d04ba Mon Sep 17 00:00:00 2001
From: Torsten Rudolf <trudolf@squareup.com>
Date: Fri, 4 Apr 2025 17:36:31 +1100
Subject: [PATCH] feat: add missing usage fields Addresses issue #434

---
 .../openai/client/TestChatCompletions.kt      | 67 ++++++++++++++-----
 .../com.aallam.openai.api/core/Usage.kt       | 42 ++++++++++++
 2 files changed, 91 insertions(+), 18 deletions(-)

diff --git a/openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestChatCompletions.kt b/openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestChatCompletions.kt
index f3a35d32..9dfa8424 100644
--- a/openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestChatCompletions.kt
+++ b/openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestChatCompletions.kt
@@ -68,7 +68,7 @@ class TestChatCompletions : TestOpenAI() {
                 function(
                     name = "currentWeather",
                     parameters =
-                    """
+                        """
                     {
                       "type": "object",
                       "properties": {
@@ -137,23 +137,33 @@ class TestChatCompletions : TestOpenAI() {
 
     @Test
     fun jsonSchema() = test {
-        val schemaJson = JsonObject(mapOf(
-            "type" to JsonPrimitive("object"),
-            "properties" to JsonObject(mapOf(
-                "question" to JsonObject(mapOf(
-                    "type" to JsonPrimitive("string"),
-                    "description" to JsonPrimitive("The question that was asked")
-                )),
-                "response" to JsonObject(mapOf(
-                    "type" to JsonPrimitive("string"),
-                    "description" to JsonPrimitive("The answer to the question")
-                ))
-            )),
-            "required" to JsonArray(listOf(
-                JsonPrimitive("question"),
-                JsonPrimitive("response")
-            ))
-        ))
+        val schemaJson = JsonObject(
+            mapOf(
+                "type" to JsonPrimitive("object"),
+                "properties" to JsonObject(
+                    mapOf(
+                        "question" to JsonObject(
+                            mapOf(
+                                "type" to JsonPrimitive("string"),
+                                "description" to JsonPrimitive("The question that was asked")
+                            )
+                        ),
+                        "response" to JsonObject(
+                            mapOf(
+                                "type" to JsonPrimitive("string"),
+                                "description" to JsonPrimitive("The answer to the question")
+                            )
+                        )
+                    )
+                ),
+                "required" to JsonArray(
+                    listOf(
+                        JsonPrimitive("question"),
+                        JsonPrimitive("response")
+                    )
+                )
+            )
+        )
 
         val jsonSchema = JsonSchema(
             name = "AnswerSchema",
@@ -301,4 +311,25 @@ class TestChatCompletions : TestOpenAI() {
         assertNotNull(results.last().usage?.completionTokens)
         assertNotNull(results.last().usage?.totalTokens)
     }
+
+    @Test
+    fun usage_fields() = test {
+        val request = chatCompletionRequest {
+            model = ModelId("gpt-3.5-turbo")
+            messages {
+                message {
+                    role = ChatRole.User
+                    content = "What's the weather like in Boston?"
+                }
+            }
+        }
+        val response = openAI.chatCompletion(request)
+        assertNotNull(response.usage)
+        assertNotNull(response.usage?.promptTokens)
+        assertNotNull(response.usage?.completionTokens)
+        assertNotNull(response.usage?.totalTokens)
+        assertNotNull(response.usage?.promptTokensDetails)
+        assertNotNull(response.usage?.promptTokensDetails?.cachedTokens)
+        assertNotNull(response.usage?.completionTokensDetails)
+    }
 }
diff --git a/openai-core/src/commonMain/kotlin/com.aallam.openai.api/core/Usage.kt b/openai-core/src/commonMain/kotlin/com.aallam.openai.api/core/Usage.kt
index bbc7d672..c0b8c1de 100644
--- a/openai-core/src/commonMain/kotlin/com.aallam.openai.api/core/Usage.kt
+++ b/openai-core/src/commonMain/kotlin/com.aallam.openai.api/core/Usage.kt
@@ -17,4 +17,46 @@ public data class Usage(
      * Count of total tokens.
      */
     @SerialName("total_tokens") public val totalTokens: Int? = null,
+    /**
+     * Breakdown of tokens used in the prompt.
+     */
+    @SerialName("prompt_tokens_details") public val promptTokensDetails: PromptTokensDetails? = null,
+    /**
+     * Breakdown of tokens used in a completion.
+     */
+    @SerialName("completion_tokens_details") public val completionTokensDetails: CompletionTokensDetails? = null,
+)
+
+@Serializable
+public data class PromptTokensDetails(
+    /**
+     * Cached tokens present in the prompt.
+     */
+    public val cachedTokens: Int? = null,
+    /**
+     * Audio input tokens present in the prompt.
+     */
+    public val audioTokens: Int? = null,
+)
+
+@Serializable
+public data class CompletionTokensDetails(
+    /**
+     * Tokens generated by the model for reasoning.
+     */
+    public val reasoningTokens: Int? = null,
+    /**
+     * Audio input tokens generated by the model.
+     */
+    public val audioTokens: Int? = null,
+    /**
+     * When using Predicted Outputs, the number of tokens in the prediction that appeared in the completion.
+     */
+    public val acceptedPredictionTokens: Int? = null,
+    /**
+     * When using Predicted Outputs, the number of tokens in the prediction that did not appear in the completion.
+     * However, like reasoning tokens, these tokens are still counted in the total completion tokens for
+     * purposes of billing, output, and context window limits.
+     */
+    public val rejectedPredictionTokens: Int? = null,
 )