Use forked google-generative-ai in Vertex AI samples (#12574)

Author: andrewheard

FirebaseVertexAI/Sample/ChatSample/Views/ErrorDetailsView.swift

@@ -142,22 +142,6 @@ struct ErrorDetailsView: View {
             SafetyRatingsSection(ratings: ratings)
           }
 
-        case GenerateContentError.invalidAPIKey:
-          Section("Error Type") {
-            Text("Invalid API Key")
-          }
-
-          Section("Details") {
-            SubtitleFormRow(title: "Error description", value: error.localizedDescription)
-            SubtitleMarkdownFormRow(
-              title: "Help",
-              value: """
-              The `API_KEY` provided in the `GoogleService-Info.plist` file is invalid. Download a
-              new copy of the file from the [Firebase Console](https://console.firebase.google.com).
-              """
-            )
-          }
-
         default:
           Section("Error Type") {
             Text("Some other error")
@@ -222,11 +206,3 @@ struct ErrorDetailsView: View {
 
   return ErrorDetailsView(error: error)
 }
-
-#Preview("Invalid API Key") {
-  ErrorDetailsView(error: GenerateContentError.invalidAPIKey)
-}
-
-#Preview("Unsupported User Location") {
-  ErrorDetailsView(error: GenerateContentError.unsupportedUserLocation)
-}

FirebaseVertexAI/Sources/Errors.swift

@@ -30,14 +30,6 @@ struct RPCError: Error {
     self.status = status
     self.details = details
   }
-
-  func isInvalidAPIKeyError() -> Bool {
-    return errorInfo?.reason == "API_KEY_INVALID"
-  }
-
-  func isUnsupportedUserLocationError() -> Bool {
-    return message == RPCErrorMessage.unsupportedUserLocation.rawValue
-  }
 }
 
 extension RPCError: Decodable {
@@ -179,10 +171,6 @@ enum RPCStatus: String, Decodable {
   case dataLoss = "DATA_LOSS"
 }
 
-enum RPCErrorMessage: String {
-  case unsupportedUserLocation = "User location is not supported for the API use."
-}
-
 enum InvalidCandidateError: Error {
   case emptyContent(underlyingError: Error)
   case malformedContent(underlyingError: Error)

FirebaseVertexAI/Sources/GenerateContentError.swift

@@ -28,17 +28,4 @@ public enum GenerateContentError: Error {
 
   /// A response didn't fully complete. See the `FinishReason` for more information.
   case responseStoppedEarly(reason: FinishReason, response: GenerateContentResponse)
-
-  /// The provided API key is invalid.
-  case invalidAPIKey(message: String)
-
-  /// The user's location (region) is not supported by the API.
-  ///
-  /// See the Google documentation for a
-  /// [list of regions](https://ai.google.dev/available_regions#available_regions)
-  /// (countries and territories) where the API is available.
-  ///
-  /// - Important: The API is only available in
-  /// [specific regions](https://ai.google.dev/available_regions#available_regions).
-  case unsupportedUserLocation
 }

FirebaseVertexAI/Sources/GenerativeAIRequest.swift

@@ -38,8 +38,8 @@ public struct RequestOptions {
   /// - Parameters:
   ///   - timeout The request’s timeout interval in seconds; if not specified uses the default value
   ///   for a `URLRequest`.
-  ///   - apiVersion The API version to use in requests to the backend; defaults to "v1".
-  public init(timeout: TimeInterval? = nil, apiVersion: String = "v1") {
+  ///   - apiVersion The API version to use in requests to the backend; defaults to "v2beta".
+  public init(timeout: TimeInterval? = nil, apiVersion: String = "v2beta") {
     self.timeout = timeout
     self.apiVersion = apiVersion
   }

FirebaseVertexAI/Sources/GenerativeAIService.swift

@@ -12,22 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import FirebaseAppCheckInterop
+import FirebaseCore
 import Foundation
 
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 struct GenerativeAIService {
   /// Gives permission to talk to the backend.
   private let apiKey: String
 
+  private let appCheck: AppCheckInterop?
+
   private let urlSession: URLSession
 
-  init(apiKey: String, urlSession: URLSession) {
+  init(apiKey: String, appCheck: AppCheckInterop?, urlSession: URLSession) {
     self.apiKey = apiKey
+    self.appCheck = appCheck
     self.urlSession = urlSession
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
-    let urlRequest = try urlRequest(request: request)
+    let urlRequest = try await urlRequest(request: request)
 
     #if DEBUG
       printCURLCommand(from: urlRequest)
@@ -59,7 +64,7 @@ struct GenerativeAIService {
       Task {
         let urlRequest: URLRequest
         do {
-          urlRequest = try self.urlRequest(request: request)
+          urlRequest = try await self.urlRequest(request: request)
         } catch {
           continuation.finish(throwing: error)
           return
@@ -146,13 +151,24 @@ struct GenerativeAIService {
 
   // MARK: - Private Helpers
 
-  private func urlRequest<T: GenerativeAIRequest>(request: T) throws -> URLRequest {
+  private func urlRequest<T: GenerativeAIRequest>(request: T) async throws -> URLRequest {
     var urlRequest = URLRequest(url: request.url)
     urlRequest.httpMethod = "POST"
     urlRequest.setValue(apiKey, forHTTPHeaderField: "x-goog-api-key")
-    urlRequest.setValue("genai-swift/\(GenerativeAISwift.version)",
-                        forHTTPHeaderField: "x-goog-api-client")
+    // TODO: Determine the right client header to use.
+    //    urlRequest.setValue("genai-swift/\(GenerativeAISwift.version))",
+    //                        forHTTPHeaderField: "x-goog-api-client")
     urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
+
+    if let appCheck {
+      let tokenResult = await appCheck.getToken(forcingRefresh: false)
+      urlRequest.setValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck")
+      if let error = tokenResult.error {
+        Logging.default
+          .debug("[GoogleGenerativeAI] Failed to fetch AppCheck token. Error: \(error)")
+      }
+    }
+
     let encoder = JSONEncoder()
     encoder.keyEncodingStrategy = .convertToSnakeCase
     urlRequest.httpBody = try encoder.encode(request)

FirebaseVertexAI/Sources/GenerativeAISwift.swift

@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+
 import Foundation
 
 #if !os(macOS) && !os(iOS)
@@ -20,7 +21,5 @@ import Foundation
 /// Constants associated with the GenerativeAISwift SDK
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
 public enum GenerativeAISwift {
-  /// String value of the SDK version
-  public static let version = "0.4.8"
-  static let baseURL = "https://generativelanguage.googleapis.com"
+  static let baseURL = "https://staging-firebaseml.sandbox.googleapis.com"
 }

FirebaseVertexAI/Sources/GenerativeModel.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import FirebaseAppCheckInterop
 import Foundation
 
 /// A type that represents a remote multimodal model (like Gemini), with the ability to generate
@@ -44,31 +45,21 @@ public final class GenerativeModel {
   ///   - apiKey: The API key for your project.
   ///   - generationConfig: The content generation parameters your model should use.
   ///   - safetySettings: A value describing what types of harmful content your model should allow.
-  ///   - requestOptions Configuration parameters for sending requests to the backend.
-  public convenience init(name: String,
-                          apiKey: String,
-                          generationConfig: GenerationConfig? = nil,
-                          safetySettings: [SafetySetting]? = nil,
-                          requestOptions: RequestOptions = RequestOptions()) {
-    self.init(
-      name: name,
-      apiKey: apiKey,
-      generationConfig: generationConfig,
-      safetySettings: safetySettings,
-      requestOptions: requestOptions,
-      urlSession: .shared
-    )
-  }
-
-  /// The designated initializer for this class.
+  ///   - requestOptions: Configuration parameters for sending requests to the backend.
+  ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
   init(name: String,
        apiKey: String,
        generationConfig: GenerationConfig? = nil,
        safetySettings: [SafetySetting]? = nil,
-       requestOptions: RequestOptions = RequestOptions(),
-       urlSession: URLSession) {
+       requestOptions: RequestOptions,
+       appCheck: AppCheckInterop?,
+       urlSession: URLSession = .shared) {
     modelResourceName = GenerativeModel.modelResourceName(name: name)
-    generativeAIService = GenerativeAIService(apiKey: apiKey, urlSession: urlSession)
+    generativeAIService = GenerativeAIService(
+      apiKey: apiKey,
+      appCheck: appCheck,
+      urlSession: urlSession
+    )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings
     self.requestOptions = requestOptions
@@ -282,10 +273,6 @@ public final class GenerativeModel {
   private static func generateContentError(from error: Error) -> GenerateContentError {
     if let error = error as? GenerateContentError {
       return error
-    } else if let error = error as? RPCError, error.isInvalidAPIKeyError() {
-      return GenerateContentError.invalidAPIKey(message: error.message)
-    } else if let error = error as? RPCError, error.isUnsupportedUserLocationError() {
-      return GenerateContentError.unsupportedUserLocation
     }
     return GenerateContentError.internalError(underlying: error)
   }

FirebaseVertexAI/Sources/VertexAI.swift

@@ -12,13 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import Foundation
-
 import FirebaseAppCheckInterop
 import FirebaseCore
-
-// Exports the GoogleGenerativeAI module to users of the SDK.
-@_exported import GoogleGenerativeAI
+import Foundation
 
 // Avoids exposing internal FirebaseCore APIs to Swift users.
 @_implementationOnly import FirebaseCoreExtension
@@ -31,17 +27,20 @@ open class VertexAI: NSObject {
   /// Returns an instance of `GoogleGenerativeAI.GenerativeModel` that uses the Vertex AI API.
   ///
   /// This instance is configured with the default `FirebaseApp`.
-  public static func generativeModel(modelName: String, location: String) -> GoogleGenerativeAI
-    .GenerativeModel {
+  ///
+  /// TODO: Add RequestOptions to public API.
+  public static func generativeModel(modelName: String, location: String) -> GenerativeModel {
     guard let app = FirebaseApp.app() else {
       fatalError("No instance of the default Firebase app was found.")
     }
     return generativeModel(app: app, modelName: modelName, location: location)
   }
 
   /// Returns an instance of `GoogleGenerativeAI.GenerativeModel` that uses the Vertex AI API.
+  ///
+  /// TODO: Add RequestOptions to public API.
   public static func generativeModel(app: FirebaseApp, modelName: String,
-                                     location: String) -> GoogleGenerativeAI.GenerativeModel {
+                                     location: String) -> GenerativeModel {
     guard let provider = ComponentType<VertexAIProvider>.instance(for: VertexAIProvider.self,
                                                                   in: app.container) else {
       fatalError("No \(VertexAIProvider.self) instance found for Firebase app: \(app.name)")
@@ -64,18 +63,15 @@ open class VertexAI: NSObject {
   private let modelResouceName: String
 
   lazy var model: GenerativeModel = {
-    let options = RequestOptions(
-      apiVersion: "v2beta",
-      endpoint: "staging-firebaseml.sandbox.googleapis.com",
-      hooks: [addAppCheckHeader]
-    )
     guard let apiKey = app.options.apiKey else {
       fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.")
     }
     return GenerativeModel(
       name: modelResouceName,
       apiKey: apiKey,
-      requestOptions: options
+      // TODO: Add RequestOptions to public API.
+      requestOptions: RequestOptions(),
+      appCheck: appCheck
     )
   }()
 
@@ -104,21 +100,4 @@ open class VertexAI: NSObject {
 
     return "projects/\(projectID)/locations/\(location)/publishers/google/models/\(modelName)"
   }
-
-  // MARK: Request Hooks
-
-  /// Adds an App Check token to the provided request if App Check is included in the app.
-  ///
-  /// This demonstrates how an App Check token can be added to requests; it is currently ignored by
-  /// the backend.
-  ///
-  /// - Parameter request: The `URLRequest` to modify by adding an App Check token header.
-  func addAppCheckHeader(request: inout URLRequest) async {
-    guard let appCheck else {
-      return
-    }
-
-    let tokenResult = await appCheck.getToken(forcingRefresh: false)
-    request.addValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck")
-  }
 }

Package.swift

@@ -192,10 +192,6 @@ let package = Package(
       "100.0.0" ..< "101.0.0"
     ),
     .package(url: "https://github.com/google/app-check.git", "10.19.0" ..< "11.0.0"),
-    .package(
-      url: "https://github.com/google/generative-ai-swift.git",
-      revision: "c9f2c4913bc65aa267815962c7e91358c2d8463f"
-    ),
   ],
   targets: [
     .target(
@@ -1374,13 +1370,8 @@ let package = Package(
         "FirebaseAppCheckInterop",
         "FirebaseCore",
         "FirebaseCoreExtension",
-        .product(name: "GoogleGenerativeAI", package: "generative-ai-swift"),
       ],
-      path: "FirebaseVertexAI/Sources",
-      sources: [
-        "VertexAI.swift",
-        "VertexAIComponent.swift",
-      ]
+      path: "FirebaseVertexAI/Sources"
     ),
   ] + firestoreTargets(),
   cLanguageStandard: .c99,