[Vertex AI] Remove CountTokensError enum (#13736)

firebase · Sep 27, 2024 · 78c341f · 78c341f
1 parent 4c9cc64
commit 78c341f
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 17 deletions.
diff --git a/FirebaseVertexAI/CHANGELOG.md b/FirebaseVertexAI/CHANGELOG.md
@@ -25,6 +25,9 @@
 - [changed] **Breaking Change**: The `ImageConversionError` enum is no longer
   public; image conversion errors are still reported as
   `GenerateContentError.promptImageContentError`. (#13735)
+- [changed] **Breaking Change**: The `CountTokensError` enum has been removed;
+  errors occurring in `GenerativeModel.countTokens(...)` are now thrown directly
+  instead of being wrapped in a `CountTokensError.internalError`. (#13736)
 - [changed] The default request timeout is now 180 seconds instead of the
   platform-default value of 60 seconds for a `URLRequest`; this timeout may
   still be customized in `RequestOptions`. (#13722)

diff --git a/FirebaseVertexAI/Sources/GenerativeModel.swift b/FirebaseVertexAI/Sources/GenerativeModel.swift
@@ -269,16 +269,12 @@ public final class GenerativeModel {
   /// invalid.
   public func countTokens(_ content: @autoclosure () throws -> [ModelContent]) async throws
     -> CountTokensResponse {
-    do {
-      let countTokensRequest = try CountTokensRequest(
-        model: modelResourceName,
-        contents: content(),
-        options: requestOptions
-      )
-      return try await generativeAIService.loadRequest(request: countTokensRequest)
-    } catch {
-      throw CountTokensError.internalError(underlying: error)
-    }
+    let countTokensRequest = try CountTokensRequest(
+      model: modelResourceName,
+      contents: content(),
+      options: requestOptions
+    )
+    return try await generativeAIService.loadRequest(request: countTokensRequest)
   }
 
   /// Returns a `GenerateContentError` (for public consumption) from an internal error.
@@ -291,9 +287,3 @@ public final class GenerativeModel {
     return GenerateContentError.internalError(underlying: error)
   }
 }
-
-/// An error thrown in `GenerativeModel.countTokens(_:)`.
-@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public enum CountTokensError: Error {
-  case internalError(underlying: Error)
-}
diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift
@@ -1258,7 +1258,7 @@ final class GenerativeModelTests: XCTestCase {
     do {
       _ = try await model.countTokens("Why is the sky blue?")
       XCTFail("Request should not have succeeded.")
-    } catch let CountTokensError.internalError(rpcError as RPCError) {
+    } catch let rpcError as RPCError {
       XCTAssertEqual(rpcError.httpResponseCode, 404)
       XCTAssertEqual(rpcError.status, .notFound)
       XCTAssert(rpcError.message.hasPrefix("models/test-model-name is not found"))