Move useLimitedUseTokens to FirebaseInfo

Author: daymxn

FirebaseAI/Sources/FirebaseAI.swift

@@ -110,8 +110,7 @@ public final class FirebaseAI: Sendable {
       tools: tools,
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
-      requestOptions: requestOptions,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      requestOptions: requestOptions
     )
   }
 
@@ -147,8 +146,7 @@ public final class FirebaseAI: Sendable {
       apiConfig: apiConfig,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      requestOptions: requestOptions,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      requestOptions: requestOptions
     )
   }
 
@@ -163,8 +161,6 @@ public final class FirebaseAI: Sendable {
 
   let apiConfig: APIConfig
 
-  let useLimitedUseAppCheckTokens: Bool
-
   /// A map of active `FirebaseAI` instances keyed by the `FirebaseApp` name and the `location`,
   /// in the format `appName:location`.
   private nonisolated(unsafe) static var instances: [InstanceKey: FirebaseAI] = [:]
@@ -227,11 +223,11 @@ public final class FirebaseAI: Sendable {
       projectID: projectID,
       apiKey: apiKey,
       firebaseAppID: app.options.googleAppID,
-      firebaseApp: app
+      firebaseApp: app,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     self.apiConfig = apiConfig
     self.location = location
-    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func modelResourceName(modelName: String) -> String {

FirebaseAI/Sources/FirebaseInfo.swift

@@ -27,19 +27,22 @@ struct FirebaseInfo: Sendable {
   let projectID: String
   let apiKey: String
   let firebaseAppID: String
+  let useLimitedUseAppCheckTokens: Bool
   let app: FirebaseApp
 
   init(appCheck: AppCheckInterop? = nil,
        auth: AuthInterop? = nil,
        projectID: String,
        apiKey: String,
        firebaseAppID: String,
-       firebaseApp: FirebaseApp) {
+       firebaseApp: FirebaseApp,
+       useLimitedUseAppCheckTokens: Bool) {
     self.appCheck = appCheck
     self.auth = auth
     self.projectID = projectID
     self.apiKey = apiKey
     self.firebaseAppID = firebaseAppID
+    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
     app = firebaseApp
   }
 }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -30,12 +30,9 @@ struct GenerativeAIService {
 
   private let urlSession: URLSession
 
-  private let useLimitedUseAppCheckTokens: Bool
-
-  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, useLimitedUseAppCheckTokens: Bool) {
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
     self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
-    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
@@ -212,7 +209,7 @@ struct GenerativeAIService {
 
   private func fetchAppCheckToken(appCheck: AppCheckInterop) async throws
     -> FIRAppCheckTokenResultInterop {
-    if useLimitedUseAppCheckTokens {
+      if firebaseInfo.useLimitedUseAppCheckTokens {
       if let token = await getLimitedUseAppCheckToken(appCheck: appCheck) {
         return token
       }
@@ -242,7 +239,7 @@ struct GenerativeAIService {
       Never
     >) in
       guard
-        useLimitedUseAppCheckTokens,
+        firebaseInfo.useLimitedUseAppCheckTokens,
         // `getLimitedUseToken(completion:)` is an optional protocol method. Optional binding
         // is performed to make sure `continuation` is called even if the method’s not implemented.
         let limitedUseTokenClosure = appCheck.getLimitedUseToken

FirebaseAI/Sources/GenerativeModel.swift

@@ -76,8 +76,6 @@ public final class GenerativeModel: Sendable {
   ///     only text content is supported.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
-  ///   - useLimitedUseAppCheckTokens: Use App Check's limited-use tokens instead of the standard
-  /// cached tokens.
   init(modelName: String,
        modelResourceName: String,
        firebaseInfo: FirebaseInfo,
@@ -88,15 +86,13 @@ public final class GenerativeModel: Sendable {
        toolConfig: ToolConfig? = nil,
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default,
-       useLimitedUseAppCheckTokens: Bool = false) {
+       urlSession: URLSession = GenAIURLSession.default) {
     self.modelName = modelName
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      urlSession: urlSession
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift

@@ -53,14 +53,12 @@ public final class ImagenModel {
        generationConfig: ImagenGenerationConfig?,
        safetySettings: ImagenSafetySettings?,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default,
-       useLimitedUseAppCheckTokens: Bool = false) {
+       urlSession: URLSession = GenAIURLSession.default) {
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      urlSession: urlSession
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Tests/Unit/ChatTests.swift

@@ -68,7 +68,8 @@ final class ChatTests: XCTestCase {
           projectID: "my-project-id",
           apiKey: "API_KEY",
           firebaseAppID: "My app ID",
-          firebaseApp: app
+          firebaseApp: app,
+          useLimitedUseAppCheckTokens: false
         ),
         apiConfig: FirebaseAI.defaultVertexAIAPIConfig,
         tools: nil,

FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift

@@ -507,13 +507,13 @@ final class GenerativeModelVertexAITests: XCTestCase {
       modelName: testModelName,
       modelResourceName: testModelResourceName,
       firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(
-        appCheck: AppCheckInteropFake(token: appCheckToken)
+        appCheck: AppCheckInteropFake(token: appCheckToken),
+        useLimitedUseAppCheckTokens: true
       ),
       apiConfig: apiConfig,
       tools: nil,
       requestOptions: RequestOptions(),
       urlSession: urlSession,
-      useLimitedUseAppCheckTokens: true
     )
     MockURLProtocol
       .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(

FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift

@@ -103,7 +103,8 @@ enum GenerativeModelTestUtil {
 
   static func testFirebaseInfo(appCheck: AppCheckInterop? = nil,
                                auth: AuthInterop? = nil,
-                               privateAppID: Bool = false) -> FirebaseInfo {
+                               privateAppID: Bool = false,
+                               useLimitedUseAppCheckTokens: Bool = false) -> FirebaseInfo {
     let app = FirebaseApp(instanceWithName: "testApp",
                           options: FirebaseOptions(googleAppID: "ignore",
                                                    gcmSenderID: "ignore"))
@@ -114,7 +115,8 @@ enum GenerativeModelTestUtil {
       projectID: "my-project-id",
       apiKey: "API_KEY",
       firebaseAppID: "My app ID",
-      firebaseApp: app
+      firebaseApp: app,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
   }
 }