Include probabilyScore too

google-gemini · Apr 18, 2024 · eb414d9 · eb414d9
1 parent b85e522
commit eb414d9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/common/src/main/kotlin/com/google/ai/client/generativeai/common/server/Types.kt b/common/src/main/kotlin/com/google/ai/client/generativeai/common/server/Types.kt
@@ -74,6 +74,7 @@ data class SafetyRating(
   val category: HarmCategory,
   val probability: HarmProbability,
   val blocked: Boolean? = null, // TODO(): any reason not to default to false?
+  val probabilityScore: Float? = null,
   val severity: HarmSeverity? = null,
   val severityScore: Float? = null,
 )

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
@@ -18,6 +18,7 @@ package com.google.ai.client.generativeai.common
 
 import com.google.ai.client.generativeai.common.server.BlockReason
 import com.google.ai.client.generativeai.common.server.FinishReason
+import com.google.ai.client.generativeai.common.server.HarmProbability
 import com.google.ai.client.generativeai.common.server.HarmSeverity
 import com.google.ai.client.generativeai.common.shared.HarmCategory
 import com.google.ai.client.generativeai.common.util.goldenUnaryFile
@@ -79,6 +80,11 @@ internal class UnarySnapshotTests {
 
         response.candidates?.isEmpty() shouldBe false
         response.candidates?.first()?.safetyRatings?.isEmpty() shouldBe false
+        response.candidates?.first()?.safetyRatings?.all {
+          it.probability == HarmProbability.NEGLIGIBLE
+        } shouldBe true
+        response.candidates?.first()?.safetyRatings?.all { it.probabilityScore != null } shouldBe
+          true
         response.candidates?.first()?.safetyRatings?.all {
           it.severity == HarmSeverity.NEGLIGIBLE
         } shouldBe true