From ca2ef423accb9139c5298faa2fd69b8460649313 Mon Sep 17 00:00:00 2001
From: Rodrigo Lazo Paz <rlazo@google.com>
Date: Tue, 17 Sep 2024 12:47:14 -0400
Subject: [PATCH 1/2] Add support for blockReasonMessage

Part of PromptFeedback
---
 .../kotlin/com/google/firebase/vertexai/common/server/Types.kt  | 1 +
 .../com/google/firebase/vertexai/internal/util/conversions.kt   | 1 +
 .../kotlin/com/google/firebase/vertexai/type/PromptFeedback.kt  | 2 ++
 3 files changed, 4 insertions(+)

diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/server/Types.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/server/Types.kt
index a5c0b00e461..16f96796bbf 100644
--- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/server/Types.kt
+++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/server/Types.kt
@@ -38,6 +38,7 @@ internal object FinishReasonSerializer :
 internal data class PromptFeedback(
   val blockReason: BlockReason? = null,
   val safetyRatings: List<SafetyRating>? = null,
+  val blockReasonMessage: String? = null,
 )
 
 @Serializable(BlockReasonSerializer::class)
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt
index b16d5ffad36..0a67271b3a1 100644
--- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt
+++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt
@@ -245,6 +245,7 @@ internal fun com.google.firebase.vertexai.common.server.PromptFeedback.toPublic(
   return com.google.firebase.vertexai.type.PromptFeedback(
     blockReason?.toPublic(),
     safetyRatings,
+    blockReasonMessage
   )
 }
 
diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/PromptFeedback.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/PromptFeedback.kt
index 7d9a0f5d438..f9d2a134628 100644
--- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/PromptFeedback.kt
+++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/PromptFeedback.kt
@@ -21,10 +21,12 @@ package com.google.firebase.vertexai.type
  *
  * @param blockReason The reason that content was blocked, if at all.
  * @param safetyRatings A list of relevant [SafetyRating].
+ * @param blockReasonMessage A message describing the reason that content was blocked, if any.
  */
 class PromptFeedback(
   val blockReason: BlockReason?,
   val safetyRatings: List<SafetyRating>,
+  val blockReasonMessage: String?
 )
 
 /** Describes why content was blocked. */

From 7505ed4ffea02be67dd4bda284c0e0116b8560b4 Mon Sep 17 00:00:00 2001
From: Rodrigo Lazo Paz <rlazo@google.com>
Date: Tue, 17 Sep 2024 17:25:48 -0400
Subject: [PATCH 2/2] Add blockReasonMessage to PromptFeedback

---
 .../firebase/vertexai/StreamingSnapshotTests.kt      | 12 ++++++++++++
 .../google/firebase/vertexai/UnarySnapshotTests.kt   | 12 ++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt
index 4dcfb5380cc..941f7647edd 100644
--- a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt
+++ b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt
@@ -123,6 +123,18 @@ internal class StreamingSnapshotTests {
       }
     }
 
+  @Test
+  fun `prompt blocked for safety with message`() =
+    goldenStreamingFile("streaming-failure-prompt-blocked-safety-with-message.txt") {
+      val responses = model.generateContentStream("prompt")
+
+      withTimeout(testTimeout) {
+        val exception = shouldThrow<PromptBlockedException> { responses.collect() }
+        exception.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+        exception.response.promptFeedback?.blockReasonMessage shouldBe "Reasons"
+      }
+    }
+
   @Test
   fun `empty content`() =
     goldenStreamingFile("streaming-failure-empty-content.txt") {
diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt
index eeb45bba43e..04f47a0949b 100644
--- a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt
+++ b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt
@@ -167,6 +167,18 @@ internal class UnarySnapshotTests {
       }
     }
 
+  @Test
+  fun `prompt blocked for safety with message`() =
+    goldenUnaryFile("unary-failure-prompt-blocked-safety-with-message.json") {
+      withTimeout(testTimeout) {
+        shouldThrow<PromptBlockedException> { model.generateContent("prompt") } should
+          {
+            it.response.promptFeedback?.blockReason shouldBe BlockReason.SAFETY
+            it.response.promptFeedback?.blockReasonMessage shouldContain "Reasons"
+          }
+      }
+    }
+
   @Test
   fun `empty content`() =
     goldenUnaryFile("unary-failure-empty-content.json") {