VertexAI: convert exceptions from countTokens() to VertexAIException (#6105)

Author: tanzimfh

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/GenerativeModel.kt

@@ -195,7 +195,11 @@ internal constructor(
    * @return A [CountTokensResponse] containing the amount of tokens in the prompt.
    */
   suspend fun countTokens(vararg prompt: Content): CountTokensResponse {
-    return controller.countTokens(constructCountTokensRequest(*prompt)).toPublic()
+    try {
+      return controller.countTokens(constructCountTokensRequest(*prompt)).toPublic()
+    } catch (e: Throwable) {
+      throw FirebaseVertexAIException.from(e)
+    }
   }
 
   /**

firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt

@@ -427,4 +427,10 @@ internal class UnarySnapshotTests {
         response.totalBillableCharacters shouldBe 0
       }
     }
+
+  @Test
+  fun `countTokens fails with model not found`() =
+    goldenUnaryFile("failure-model-not-found.json", HttpStatusCode.NotFound) {
+      withTimeout(testTimeout) { shouldThrow<ServerException> { model.countTokens("prompt") } }
+    }
 }