Add support for usage metadata (#91)

google-gemini · Mar 22, 2024 · b2afc55 · b2afc55
1 parent 4638406
commit b2afc55
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 0 deletions.
diff --git a/common/src/main/kotlin/com/google/ai/client/generativeai/common/Response.kt b/common/src/main/kotlin/com/google/ai/client/generativeai/common/Response.kt
@@ -27,8 +27,16 @@ sealed interface Response
 data class GenerateContentResponse(
   val candidates: List<Candidate>? = null,
   val promptFeedback: PromptFeedback? = null,
+  val usageMetadata: UsageMetadata? = null
 ) : Response
 
 @Serializable data class CountTokensResponse(val totalTokens: Int) : Response
 
 @Serializable data class GRpcErrorResponse(val error: GRpcError) : Response
+
+@Serializable
+data class UsageMetadata(
+  val promptTokenCount: Int,
+  val candidatesTokenCount: Int,
+  val totalTokenCount: Int
+)
diff --git a/common/src/test/java/com/google/ai/client/generativeai/common/UnarySnapshotTests.kt b/common/src/test/java/com/google/ai/client/generativeai/common/UnarySnapshotTests.kt
@@ -23,6 +23,7 @@ import com.google.ai.client.generativeai.common.util.goldenUnaryFile
 import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.should
 import io.kotest.matchers.shouldBe
+import io.kotest.matchers.shouldNotBe
 import io.ktor.http.HttpStatusCode
 import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.withTimeout
@@ -132,6 +133,19 @@ internal class UnarySnapshotTests {
       }
     }
 
+  @Test
+  fun `response includes usage metadata`() =
+    goldenUnaryFile("success-usage-metadata.json") {
+      withTimeout(testTimeout) {
+        val response = apiController.generateContent(textGenerateContentRequest("prompt"))
+
+        response.candidates?.isEmpty() shouldBe false
+        response.candidates?.first()?.finishReason shouldBe FinishReason.STOP
+        response.usageMetadata shouldNotBe null
+        response.usageMetadata?.totalTokenCount shouldBe 363
+      }
+    }
+
   @Test
   fun `citation returns correctly when using alternative name`() =
     goldenUnaryFile("success-citations-altname.json") {

diff --git a/common/src/test/resources/golden-files/unary/success-usage-metadata.json b/common/src/test/resources/golden-files/unary/success-usage-metadata.json
@@ -0,0 +1,59 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "text": "Mountain View, California, United States"
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0,
+      "safetyRatings": [
+        {
+          "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_HATE_SPEECH",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_HARASSMENT",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+          "probability": "NEGLIGIBLE"
+        }
+      ]
+    }
+  ],
+  "usageMetadata": {
+    "promptTokenCount": 6,
+    "candidatesTokenCount": 357,
+    "totalTokenCount": 363
+  },
+  "promptFeedback": {
+    "safetyRatings": [
+      {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "probability": "NEGLIGIBLE"
+      },
+      {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "probability": "NEGLIGIBLE"
+      },
+      {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "probability": "NEGLIGIBLE"
+      },
+      {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "probability": "NEGLIGIBLE"
+      }
+    ]
+  }
+}