Merge branch 'main' into dconeybe/dataconnect/TestModernizationAndSom…

…eNewFeatures
firebase · Nov 5, 2024 · ded7347 · ded7347
2 parents 12b56b0 + b9013d5
commit ded7347
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 1 deletion.
diff --git a/firebase-vertexai/CHANGELOG.md b/firebase-vertexai/CHANGELOG.md
@@ -1,4 +1,5 @@
 # Unreleased
+* [fixed] Improved error message when using an invalid location. (#6428)
 * [fixed] Fixed issue where Firebase App Check error tokens were unintentionally missing from the requests. (#6409)
 * [fixed] Clarified in the documentation that `Schema.integer` and `Schema.float` only provide hints to the model. (#6420)
 

diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/APIController.kt
@@ -42,7 +42,9 @@ import io.ktor.client.statement.bodyAsText
 import io.ktor.http.ContentType
 import io.ktor.http.HttpStatusCode
 import io.ktor.http.contentType
+import io.ktor.http.withCharset
 import io.ktor.serialization.kotlinx.json.json
+import io.ktor.utils.io.charsets.Charset
 import kotlin.math.max
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
@@ -225,6 +227,15 @@ internal interface HeaderProvider {
 
 private suspend fun validateResponse(response: HttpResponse) {
   if (response.status == HttpStatusCode.OK) return
+
+  val htmlContentType = ContentType.Text.Html.withCharset(Charset.forName("utf-8"))
+  if (response.status == HttpStatusCode.NotFound && response.contentType() == htmlContentType)
+    throw ServerException(
+      """URL not found. Please verify the location used to create the `FirebaseVertexAI` object
+          | See https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#available-regions
+          | for the list of available locations. Raw response: ${response.bodyAsText()}"""
+        .trimMargin()
+    )
   val text = response.bodyAsText()
   val error =
     try {

diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/GenerativeModelTesting.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/GenerativeModelTesting.kt
@@ -24,10 +24,13 @@ import com.google.firebase.vertexai.common.shared.Content
 import com.google.firebase.vertexai.common.shared.TextPart
 import com.google.firebase.vertexai.common.util.doBlocking
 import com.google.firebase.vertexai.type.RequestOptions
+import com.google.firebase.vertexai.type.ServerException
 import com.google.firebase.vertexai.type.content
 import io.kotest.assertions.json.shouldContainJsonKey
 import io.kotest.assertions.json.shouldContainJsonKeyValue
+import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.collections.shouldNotBeEmpty
+import io.kotest.matchers.string.shouldContain
 import io.kotest.matchers.types.shouldBeInstanceOf
 import io.ktor.client.engine.mock.MockEngine
 import io.ktor.client.engine.mock.respond
@@ -44,7 +47,7 @@ internal class GenerativeModelTesting {
   private val TEST_CLIENT_ID = "test"
 
   @Test
-  fun addition() = doBlocking {
+  fun `system calling in request`() = doBlocking {
     val mockEngine = MockEngine {
       respond(
         generateContentResponseAsJsonString("text response"),
@@ -84,6 +87,46 @@ internal class GenerativeModelTesting {
     }
   }
 
+  @Test
+  fun `exception thrown when using invalid location`() = doBlocking {
+    val mockEngine = MockEngine {
+      respond(
+        """<!DOCTYPE html>
+           <html lang=en>
+            <title>Error 404 (Not Found)!!1</title>
+        """
+          .trimIndent(),
+        HttpStatusCode.NotFound,
+        headersOf(HttpHeaders.ContentType, "text/html; charset=utf-8")
+      )
+    }
+
+    val apiController =
+      APIController(
+        "super_cool_test_key",
+        "gemini-1.5-flash",
+        RequestOptions(),
+        mockEngine,
+        TEST_CLIENT_ID,
+        null,
+      )
+
+    // Creating the
+    val generativeModel =
+      GenerativeModel(
+        "projects/PROJECTID/locations/INVALID_LOCATION/publishers/google/models/gemini-1.5-flash",
+        controller = apiController
+      )
+
+    val exception =
+      shouldThrow<ServerException> {
+        withTimeout(5.seconds) { generativeModel.generateContent("my test prompt") }
+      }
+
+    // Let's not be too strict on the wording to avoid breaking the test unnecessarily.
+    exception.message shouldContain "location"
+  }
+
   private fun generateContentResponseAsJsonString(text: String): String {
     return JSON.encodeToString(
       GenerateContentResponse(listOf(Candidate(Content(parts = listOf(TextPart(text))))))