Add support for SafetyRating.HarmSeverity in common

common/src/main/kotlin/com/google/ai/client/generativeai/common/server/Types.kt

@@ -73,7 +73,9 @@ data class CitationSources(
 data class SafetyRating(
   val category: HarmCategory,
   val probability: HarmProbability,
-  val blocked: Boolean? = null // TODO(): any reason not to default to false?
+  val blocked: Boolean? = null, // TODO(): any reason not to default to false?
+  val severity: HarmSeverity? = null,
+  val severityScore: Float? = null,
 )
 
 @Serializable(HarmProbabilitySerializer::class)
@@ -86,6 +88,16 @@ enum class HarmProbability {
   HIGH
 }
 
+@Serializable
+enum class HarmSeverity {
+  UNKNOWN,
+  @SerialName("HARM_SEVERITY_UNSPECIFIED") UNSPECIFIED,
+  @SerialName("HARM_SEVERITY_NEGLIGIBLE") NEGLIGIBLE,
+  @SerialName("HARM_SEVERITY_LOW") LOW,
+  @SerialName("HARM_SEVERITY_MEDIUM") MEDIUM,
+  @SerialName("HARM_SEVERITY_HIGH") HIGH
+}
+
 @Serializable(FinishReasonSerializer::class)
 enum class FinishReason {
   UNKNOWN,

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.HarmSeverity
 import com.google.ai.client.generativeai.common.shared.HarmCategory
 import com.google.ai.client.generativeai.common.util.goldenUnaryFile
 import io.kotest.assertions.throwables.shouldThrow
@@ -70,6 +71,21 @@ internal class UnarySnapshotTests {
       }
     }
 
+  @Test
+  fun `safetyRatings including severity`() =
+    goldenUnaryFile("success-including-severity.json") {
+      withTimeout(testTimeout) {
+        val response = apiController.generateContent(textGenerateContentRequest("prompt"))
+
+        response.candidates?.isEmpty() shouldBe false
+        response.candidates?.first()?.safetyRatings?.isEmpty() shouldBe false
+        response.candidates?.first()?.safetyRatings?.all {
+          it.severity == HarmSeverity.NEGLIGIBLE
+        } shouldBe true
+        response.candidates?.first()?.safetyRatings?.all { it.severityScore != null } shouldBe true
+      }
+    }
+
   @Test
   fun `prompt blocked for safety`() =
     goldenUnaryFile("failure-prompt-blocked-safety.json") {

common/src/test/resources/golden-files/unary/success-including-severity.json

@@ -0,0 +1,50 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "role": "model",
+        "parts": [
+          {
+            "text": "## One Thousand and One Nights: A Summary\n\nOne Thousand and One Nights, also known as Arabian Nights, is a collection of Middle Eastern and South Asian stories."
+          }
+        ]
+      },
+      "finishReason": "STOP",
+      "safetyRatings": [
+        {
+          "category": "HARM_CATEGORY_HATE_SPEECH",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.062331032,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.052134257
+        },
+        {
+          "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.04240383,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.06325052
+        },
+        {
+          "category": "HARM_CATEGORY_HARASSMENT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.06359858,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.021990221
+        },
+        {
+          "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.39030153,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.10650458
+        }
+      ]
+    }
+  ],
+  "usageMetadata": {
+    "promptTokenCount": 11,
+    "candidatesTokenCount": 592,
+    "totalTokenCount": 603
+  }
+}