diff --git a/FirebaseAI/Sources/FirebaseAI.swift b/FirebaseAI/Sources/FirebaseAI.swift index d4c2e9d545d..da954ab7b28 100644 --- a/FirebaseAI/Sources/FirebaseAI.swift +++ b/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 { diff --git a/FirebaseAI/Sources/FirebaseInfo.swift b/FirebaseAI/Sources/FirebaseInfo.swift index c1f27aa7fe3..6b10dec4e5f 100644 --- a/FirebaseAI/Sources/FirebaseInfo.swift +++ b/FirebaseAI/Sources/FirebaseInfo.swift @@ -27,6 +27,7 @@ struct FirebaseInfo: Sendable { let projectID: String let apiKey: String let firebaseAppID: String + let useLimitedUseAppCheckTokens: Bool let app: FirebaseApp init(appCheck: AppCheckInterop? = nil, @@ -34,12 +35,14 @@ struct FirebaseInfo: Sendable { 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 } } diff --git a/FirebaseAI/Sources/GenerativeAIService.swift b/FirebaseAI/Sources/GenerativeAIService.swift index 5b676162c94..8056d4172b8 100644 --- a/FirebaseAI/Sources/GenerativeAIService.swift +++ b/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(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 diff --git a/FirebaseAI/Sources/GenerativeModel.swift b/FirebaseAI/Sources/GenerativeModel.swift index 2807022de59..8d3f5e043a7 100644 --- a/FirebaseAI/Sources/GenerativeModel.swift +++ b/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 diff --git a/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift index 542f842362c..729fad9f28d 100644 --- a/FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift +++ b/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 diff --git a/FirebaseAI/Tests/Unit/ChatTests.swift b/FirebaseAI/Tests/Unit/ChatTests.swift index 4bf89f8cb51..7ecebf42e28 100644 --- a/FirebaseAI/Tests/Unit/ChatTests.swift +++ b/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, diff --git a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift index 80a383eb461..94589577568 100644 --- a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift +++ b/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( diff --git a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift b/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift index 9b00f0b0c87..ee4f47bc5b0 100644 --- a/FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift +++ b/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 ) } }