From a23ee1e0bbaf313473fe8912fb20bb70edac0b87 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 21 Feb 2025 15:32:28 -0800 Subject: [PATCH 1/6] Add google-app-id to Vertex AI requests --- .../Sources/GenerativeAIService.swift | 11 ++++- .../Sources/GenerativeModel.swift | 2 + .../Types/Public/Imagen/ImagenModel.swift | 2 + FirebaseVertexAI/Sources/VertexAI.swift | 5 ++- FirebaseVertexAI/Tests/Unit/ChatTests.swift | 1 + .../Tests/Unit/GenerativeModelTests.swift | 42 ++++++++++++++++++- 6 files changed, 59 insertions(+), 4 deletions(-) diff --git a/FirebaseVertexAI/Sources/GenerativeAIService.swift b/FirebaseVertexAI/Sources/GenerativeAIService.swift index fc35c2b258a..8aad4d59601 100644 --- a/FirebaseVertexAI/Sources/GenerativeAIService.swift +++ b/FirebaseVertexAI/Sources/GenerativeAIService.swift @@ -31,16 +31,19 @@ struct GenerativeAIService { /// Gives permission to talk to the backend. private let apiKey: String + private let googleAppID: String? + private let appCheck: AppCheckInterop? private let auth: AuthInterop? private let urlSession: URLSession - init(projectID: String, apiKey: String, appCheck: AppCheckInterop?, auth: AuthInterop?, - urlSession: URLSession) { + init(projectID: String, apiKey: String, googleAppID: String?, + appCheck: AppCheckInterop?, auth: AuthInterop?, urlSession: URLSession) { self.projectID = projectID self.apiKey = apiKey + self.googleAppID = googleAppID self.appCheck = appCheck self.auth = auth self.urlSession = urlSession @@ -202,6 +205,10 @@ struct GenerativeAIService { urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization") } + if let googleAppID = self.googleAppID { + urlRequest.setValue(googleAppID, forHTTPHeaderField: "google-app-id") + } + let encoder = JSONEncoder() urlRequest.httpBody = try encoder.encode(request) urlRequest.timeoutInterval = request.options.timeout diff --git a/FirebaseVertexAI/Sources/GenerativeModel.swift b/FirebaseVertexAI/Sources/GenerativeModel.swift index 0d2ea829f55..c3940ab4548 100644 --- a/FirebaseVertexAI/Sources/GenerativeModel.swift +++ b/FirebaseVertexAI/Sources/GenerativeModel.swift @@ -61,6 +61,7 @@ public final class GenerativeModel { init(name: String, projectID: String, apiKey: String, + googleAppID: String?, generationConfig: GenerationConfig? = nil, safetySettings: [SafetySetting]? = nil, tools: [Tool]?, @@ -74,6 +75,7 @@ public final class GenerativeModel { generativeAIService = GenerativeAIService( projectID: projectID, apiKey: apiKey, + googleAppID: googleAppID, appCheck: appCheck, auth: auth, urlSession: urlSession diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift index 2196d4c6040..9a4e311d964 100644 --- a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift +++ b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift @@ -43,6 +43,7 @@ public final class ImagenModel { init(name: String, projectID: String, apiKey: String, + googleAppID: String?, generationConfig: ImagenGenerationConfig?, safetySettings: ImagenSafetySettings?, requestOptions: RequestOptions, @@ -53,6 +54,7 @@ public final class ImagenModel { generativeAIService = GenerativeAIService( projectID: projectID, apiKey: apiKey, + googleAppID: googleAppID, appCheck: appCheck, auth: auth, urlSession: urlSession diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift index bb6415988a5..fc192176214 100644 --- a/FirebaseVertexAI/Sources/VertexAI.swift +++ b/FirebaseVertexAI/Sources/VertexAI.swift @@ -93,6 +93,7 @@ public class VertexAI { name: modelResourceName(modelName: modelName), projectID: projectID, apiKey: apiKey, + googleAppID: googleAppID, generationConfig: generationConfig, safetySettings: safetySettings, tools: tools, @@ -128,6 +129,7 @@ public class VertexAI { name: modelResourceName(modelName: modelName), projectID: projectID, apiKey: apiKey, + googleAppID: googleAppID, generationConfig: generationConfig, safetySettings: safetySettings, requestOptions: requestOptions, @@ -167,6 +169,7 @@ public class VertexAI { let projectID: String let apiKey: String + let googleAppID: String? let location: String init(app: FirebaseApp, location: String) { @@ -183,7 +186,7 @@ public class VertexAI { fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.") } self.apiKey = apiKey - + self.googleAppID = app.isDataCollectionDefaultEnabled ? app.options.googleAppID : nil self.location = location } diff --git a/FirebaseVertexAI/Tests/Unit/ChatTests.swift b/FirebaseVertexAI/Tests/Unit/ChatTests.swift index 1c4988faf7c..2d82c13b838 100644 --- a/FirebaseVertexAI/Tests/Unit/ChatTests.swift +++ b/FirebaseVertexAI/Tests/Unit/ChatTests.swift @@ -57,6 +57,7 @@ final class ChatTests: XCTestCase { name: "my-model", projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: nil, diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index 3ed40ce2530..6b0a08bbda0 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -70,6 +70,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: nil, @@ -271,6 +272,7 @@ final class GenerativeModelTests: XCTestCase { name: "models/test-model", projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: nil, @@ -391,6 +393,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: AppCheckInteropFake(token: appCheckToken), @@ -407,11 +410,36 @@ final class GenerativeModelTests: XCTestCase { _ = try await model.generateContent(testPrompt) } + func testGenerateContent_dataCollectionOff() async throws { + let appCheckToken = "test-valid-token" + model = GenerativeModel( + name: testModelResourceName, + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: nil, + tools: nil, + requestOptions: RequestOptions(), + appCheck: AppCheckInteropFake(token: appCheckToken), + auth: nil, + urlSession: urlSession + ) + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-basic-reply-short", + withExtension: "json", + appCheckToken: appCheckToken, + dataCollection: false + ) + + _ = try await model.generateContent(testPrompt) + } + func testGenerateContent_appCheck_tokenRefreshError() async throws { model = GenerativeModel( name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), @@ -434,6 +462,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: nil, @@ -455,6 +484,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: nil, @@ -476,6 +506,7 @@ final class GenerativeModelTests: XCTestCase { name: "my-model", projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: nil, @@ -858,6 +889,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: requestOptions, appCheck: nil, @@ -1153,6 +1185,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: AppCheckInteropFake(token: appCheckToken), @@ -1175,6 +1208,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: RequestOptions(), appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), @@ -1321,6 +1355,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: requestOptions, appCheck: nil, @@ -1396,6 +1431,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", generationConfig: generationConfig, tools: [Tool(functionDeclarations: [sumFunction])], systemInstruction: systemInstruction, @@ -1455,6 +1491,7 @@ final class GenerativeModelTests: XCTestCase { name: testModelResourceName, projectID: "my-project-id", apiKey: "API_KEY", + googleAppID: "My app ID", tools: nil, requestOptions: requestOptions, appCheck: nil, @@ -1495,7 +1532,8 @@ final class GenerativeModelTests: XCTestCase { statusCode: Int = 200, timeout: TimeInterval = RequestOptions().timeout, appCheckToken: String? = nil, - authToken: String? = nil) throws -> ((URLRequest) throws -> ( + authToken: String? = nil, + dataCollection: Bool = true) throws -> ((URLRequest) throws -> ( URLResponse, AsyncLineSequence? )) { @@ -1515,6 +1553,8 @@ final class GenerativeModelTests: XCTestCase { XCTAssert(apiClientTags.contains(GenerativeAIService.languageTag)) XCTAssert(apiClientTags.contains(GenerativeAIService.firebaseVersionTag)) XCTAssertEqual(request.value(forHTTPHeaderField: "X-Firebase-AppCheck"), appCheckToken) + let googleAppID = request.value(forHTTPHeaderField: "google-app-id") + XCTAssertEqual(googleAppID, dataCollection ? "My app ID" : nil) if let authToken { XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Firebase \(authToken)") } else { From 117160c0af597b458b89d8c69a0c55609966b855 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 21 Feb 2025 15:41:52 -0800 Subject: [PATCH 2/6] style --- FirebaseVertexAI/Sources/GenerativeAIService.swift | 2 +- FirebaseVertexAI/Sources/VertexAI.swift | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/FirebaseVertexAI/Sources/GenerativeAIService.swift b/FirebaseVertexAI/Sources/GenerativeAIService.swift index 8aad4d59601..7bacf842af1 100644 --- a/FirebaseVertexAI/Sources/GenerativeAIService.swift +++ b/FirebaseVertexAI/Sources/GenerativeAIService.swift @@ -205,7 +205,7 @@ struct GenerativeAIService { urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization") } - if let googleAppID = self.googleAppID { + if let googleAppID = googleAppID { urlRequest.setValue(googleAppID, forHTTPHeaderField: "google-app-id") } diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift index fc192176214..8d4c2f88ea2 100644 --- a/FirebaseVertexAI/Sources/VertexAI.swift +++ b/FirebaseVertexAI/Sources/VertexAI.swift @@ -186,7 +186,7 @@ public class VertexAI { fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.") } self.apiKey = apiKey - self.googleAppID = app.isDataCollectionDefaultEnabled ? app.options.googleAppID : nil + googleAppID = app.isDataCollectionDefaultEnabled ? app.options.googleAppID : nil self.location = location } From 7bdede4b126a5ffa39169dd3b0e0cebd83457388 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 23 Feb 2025 07:51:56 -0800 Subject: [PATCH 3/6] Refactor with FirebaseInfo --- FirebaseVertexAI/Sources/FirebaseInfo.swift | 41 +++++ .../Sources/GenerativeAIService.swift | 31 ++-- .../Sources/GenerativeModel.swift | 12 +- .../Types/Public/Imagen/ImagenModel.swift | 12 +- FirebaseVertexAI/Sources/VertexAI.swift | 46 ++---- FirebaseVertexAI/Tests/Unit/ChatTests.swift | 10 +- .../Tests/Unit/GenerativeModelTests.swift | 148 +++++++++--------- .../Tests/Unit/VertexComponentTests.swift | 8 +- 8 files changed, 159 insertions(+), 149 deletions(-) create mode 100644 FirebaseVertexAI/Sources/FirebaseInfo.swift diff --git a/FirebaseVertexAI/Sources/FirebaseInfo.swift b/FirebaseVertexAI/Sources/FirebaseInfo.swift new file mode 100644 index 00000000000..a6677fcd0f8 --- /dev/null +++ b/FirebaseVertexAI/Sources/FirebaseInfo.swift @@ -0,0 +1,41 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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 + +import FirebaseAppCheckInterop +import FirebaseAuthInterop +import FirebaseCore + +/// Firebase data used by VertexAI +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +struct FirebaseInfo { + let appCheck: AppCheckInterop? + let auth: AuthInterop? + let projectID: String + let apiKey: String + let googleAppID: String? + + init(appCheck: AppCheckInterop? = nil, + auth: AuthInterop? = nil, + projectID: String, + apiKey: String, + googleAppID: String?) { + self.appCheck = appCheck + self.auth = auth + self.projectID = projectID + self.apiKey = apiKey + self.googleAppID = googleAppID + } +} diff --git a/FirebaseVertexAI/Sources/GenerativeAIService.swift b/FirebaseVertexAI/Sources/GenerativeAIService.swift index 7bacf842af1..15c878dc33c 100644 --- a/FirebaseVertexAI/Sources/GenerativeAIService.swift +++ b/FirebaseVertexAI/Sources/GenerativeAIService.swift @@ -26,26 +26,12 @@ struct GenerativeAIService { /// The Firebase SDK version in the format `fire/`. static let firebaseVersionTag = "fire/\(FirebaseVersion())" - private let projectID: String - - /// Gives permission to talk to the backend. - private let apiKey: String - - private let googleAppID: String? - - private let appCheck: AppCheckInterop? - - private let auth: AuthInterop? + private let firebaseInfo: FirebaseInfo private let urlSession: URLSession - init(projectID: String, apiKey: String, googleAppID: String?, - appCheck: AppCheckInterop?, auth: AuthInterop?, urlSession: URLSession) { - self.projectID = projectID - self.apiKey = apiKey - self.googleAppID = googleAppID - self.appCheck = appCheck - self.auth = auth + init(firebaseInfo: FirebaseInfo, urlSession: URLSession) { + self.firebaseInfo = firebaseInfo self.urlSession = urlSession } @@ -183,14 +169,14 @@ struct GenerativeAIService { private func urlRequest(request: T) async throws -> URLRequest { var urlRequest = URLRequest(url: request.url) urlRequest.httpMethod = "POST" - urlRequest.setValue(apiKey, forHTTPHeaderField: "x-goog-api-key") + urlRequest.setValue(firebaseInfo.apiKey, forHTTPHeaderField: "x-goog-api-key") urlRequest.setValue( "\(GenerativeAIService.languageTag) \(GenerativeAIService.firebaseVersionTag)", forHTTPHeaderField: "x-goog-api-client" ) urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type") - if let appCheck { + if let appCheck = firebaseInfo.appCheck { let tokenResult = await appCheck.getToken(forcingRefresh: false) urlRequest.setValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck") if let error = tokenResult.error { @@ -201,11 +187,13 @@ struct GenerativeAIService { } } - if let auth, let authToken = try await auth.getToken(forcingRefresh: false) { + if let auth = firebaseInfo.auth, let authToken = try await auth.getToken( + forcingRefresh: false + ) { urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization") } - if let googleAppID = googleAppID { + if let googleAppID = firebaseInfo.googleAppID { urlRequest.setValue(googleAppID, forHTTPHeaderField: "google-app-id") } @@ -267,6 +255,7 @@ struct GenerativeAIService { // Log specific RPC errors that cannot be mitigated or handled by user code. // These errors do not produce specific GenerateContentError or CountTokensError cases. private func logRPCError(_ error: BackendError) { + let projectID = firebaseInfo.projectID if error.isVertexAIInFirebaseServiceDisabledError() { VertexLog.error(code: .vertexAIInFirebaseAPIDisabled, """ The Vertex AI in Firebase SDK requires the Vertex AI in Firebase API \ diff --git a/FirebaseVertexAI/Sources/GenerativeModel.swift b/FirebaseVertexAI/Sources/GenerativeModel.swift index c3940ab4548..ef104cbc8de 100644 --- a/FirebaseVertexAI/Sources/GenerativeModel.swift +++ b/FirebaseVertexAI/Sources/GenerativeModel.swift @@ -59,25 +59,17 @@ public final class GenerativeModel { /// - requestOptions: Configuration parameters for sending requests to the backend. /// - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`. init(name: String, - projectID: String, - apiKey: String, - googleAppID: String?, + firebaseInfo: FirebaseInfo, generationConfig: GenerationConfig? = nil, safetySettings: [SafetySetting]? = nil, tools: [Tool]?, toolConfig: ToolConfig? = nil, systemInstruction: ModelContent? = nil, requestOptions: RequestOptions, - appCheck: AppCheckInterop?, - auth: AuthInterop?, urlSession: URLSession = .shared) { modelResourceName = name generativeAIService = GenerativeAIService( - projectID: projectID, - apiKey: apiKey, - googleAppID: googleAppID, - appCheck: appCheck, - auth: auth, + firebaseInfo: firebaseInfo, urlSession: urlSession ) self.generationConfig = generationConfig diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift index 9a4e311d964..8f894a52488 100644 --- a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift +++ b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift @@ -41,22 +41,14 @@ public final class ImagenModel { let requestOptions: RequestOptions init(name: String, - projectID: String, - apiKey: String, - googleAppID: String?, + firebaseInfo: FirebaseInfo, generationConfig: ImagenGenerationConfig?, safetySettings: ImagenSafetySettings?, requestOptions: RequestOptions, - appCheck: AppCheckInterop?, - auth: AuthInterop?, urlSession: URLSession = .shared) { modelResourceName = name generativeAIService = GenerativeAIService( - projectID: projectID, - apiKey: apiKey, - googleAppID: googleAppID, - appCheck: appCheck, - auth: auth, + firebaseInfo: firebaseInfo, urlSession: urlSession ) self.generationConfig = generationConfig diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift index 8d4c2f88ea2..fd8931ecf27 100644 --- a/FirebaseVertexAI/Sources/VertexAI.swift +++ b/FirebaseVertexAI/Sources/VertexAI.swift @@ -91,17 +91,13 @@ public class VertexAI { -> GenerativeModel { return GenerativeModel( name: modelResourceName(modelName: modelName), - projectID: projectID, - apiKey: apiKey, - googleAppID: googleAppID, + firebaseInfo: firebaseInfo, generationConfig: generationConfig, safetySettings: safetySettings, tools: tools, toolConfig: toolConfig, systemInstruction: systemInstruction, - requestOptions: requestOptions, - appCheck: appCheck, - auth: auth + requestOptions: requestOptions ) } @@ -127,14 +123,10 @@ public class VertexAI { requestOptions: RequestOptions = RequestOptions()) -> ImagenModel { return ImagenModel( name: modelResourceName(modelName: modelName), - projectID: projectID, - apiKey: apiKey, - googleAppID: googleAppID, + firebaseInfo: firebaseInfo, generationConfig: generationConfig, safetySettings: safetySettings, - requestOptions: requestOptions, - appCheck: appCheck, - auth: auth + requestOptions: requestOptions ) } @@ -144,12 +136,8 @@ public class VertexAI { // MARK: - Private - /// The `FirebaseApp` associated with this `VertexAI` instance. - private let app: FirebaseApp - - private let appCheck: AppCheckInterop? - - private let auth: AuthInterop? + /// Firebase data relevant to Vertex AI. + let firebaseInfo: FirebaseInfo #if compiler(>=6) /// A map of active `VertexAI` instances keyed by the `FirebaseApp` name and the `location`, in @@ -167,26 +155,25 @@ public class VertexAI { private static var instancesLock: os_unfair_lock = .init() #endif - let projectID: String - let apiKey: String - let googleAppID: String? let location: String init(app: FirebaseApp, location: String) { - self.app = app - appCheck = ComponentType.instance(for: AppCheckInterop.self, in: app.container) - auth = ComponentType.instance(for: AuthInterop.self, in: app.container) - guard let projectID = app.options.projectID else { fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.") } - self.projectID = projectID - guard let apiKey = app.options.apiKey else { fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.") } - self.apiKey = apiKey - googleAppID = app.isDataCollectionDefaultEnabled ? app.options.googleAppID : nil + firebaseInfo = FirebaseInfo( + appCheck: ComponentType.instance( + for: AppCheckInterop.self, + in: app.container + ), + auth: ComponentType.instance(for: AuthInterop.self, in: app.container), + projectID: projectID, + apiKey: apiKey, + googleAppID: app.isDataCollectionDefaultEnabled ? app.options.googleAppID : nil + ) self.location = location } @@ -208,6 +195,7 @@ public class VertexAI { """) } + let projectID = firebaseInfo.projectID return "projects/\(projectID)/locations/\(location)/publishers/google/models/\(modelName)" } } diff --git a/FirebaseVertexAI/Tests/Unit/ChatTests.swift b/FirebaseVertexAI/Tests/Unit/ChatTests.swift index 2d82c13b838..b0640bd5c5d 100644 --- a/FirebaseVertexAI/Tests/Unit/ChatTests.swift +++ b/FirebaseVertexAI/Tests/Unit/ChatTests.swift @@ -55,13 +55,13 @@ final class ChatTests: XCTestCase { let model = GenerativeModel( name: "my-model", - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: nil, - auth: nil, urlSession: urlSession ) let chat = Chat(model: model, history: []) diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index 6b0a08bbda0..d990c5d586a 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -68,13 +68,13 @@ final class GenerativeModelTests: XCTestCase { urlSession = try XCTUnwrap(URLSession(configuration: configuration)) model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: nil, - auth: nil, urlSession: urlSession ) } @@ -270,13 +270,13 @@ final class GenerativeModelTests: XCTestCase { let model = GenerativeModel( // Model name is prefixed with "models/". name: "models/test-model", - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: nil, - auth: nil, urlSession: urlSession ) @@ -391,13 +391,14 @@ final class GenerativeModelTests: XCTestCase { let appCheckToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + appCheck: AppCheckInteropFake(token: appCheckToken), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: AppCheckInteropFake(token: appCheckToken), - auth: nil, urlSession: urlSession ) MockURLProtocol @@ -414,13 +415,14 @@ final class GenerativeModelTests: XCTestCase { let appCheckToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: nil, + firebaseInfo: FirebaseInfo( + appCheck: AppCheckInteropFake(token: appCheckToken), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: nil + ), tools: nil, requestOptions: RequestOptions(), - appCheck: AppCheckInteropFake(token: appCheckToken), - auth: nil, urlSession: urlSession ) MockURLProtocol @@ -437,13 +439,14 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContent_appCheck_tokenRefreshError() async throws { model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), - auth: nil, urlSession: urlSession ) MockURLProtocol @@ -460,13 +463,14 @@ final class GenerativeModelTests: XCTestCase { let authToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + auth: AuthInteropFake(token: authToken), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: nil, - auth: AuthInteropFake(token: authToken), urlSession: urlSession ) MockURLProtocol @@ -482,13 +486,14 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContent_auth_nilAuthToken() async throws { model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + auth: AuthInteropFake(token: nil), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: nil, - auth: AuthInteropFake(token: nil), urlSession: urlSession ) MockURLProtocol @@ -504,13 +509,14 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContent_auth_authTokenRefreshError() async throws { model = GenerativeModel( name: "my-model", - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + auth: AuthInteropFake(error: AuthErrorFake()), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: nil, - auth: AuthInteropFake(error: AuthErrorFake()), urlSession: urlSession ) MockURLProtocol @@ -887,13 +893,13 @@ final class GenerativeModelTests: XCTestCase { let requestOptions = RequestOptions(timeout: expectedTimeout) model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: requestOptions, - appCheck: nil, - auth: nil, urlSession: urlSession ) @@ -1183,13 +1189,14 @@ final class GenerativeModelTests: XCTestCase { let appCheckToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + appCheck: AppCheckInteropFake(token: appCheckToken), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: AppCheckInteropFake(token: appCheckToken), - auth: nil, urlSession: urlSession ) MockURLProtocol @@ -1206,13 +1213,14 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContentStream_appCheck_tokenRefreshError() async throws { model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: RequestOptions(), - appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), - auth: nil, urlSession: urlSession ) MockURLProtocol @@ -1353,13 +1361,13 @@ final class GenerativeModelTests: XCTestCase { let requestOptions = RequestOptions(timeout: expectedTimeout) model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: requestOptions, - appCheck: nil, - auth: nil, urlSession: urlSession ) @@ -1429,15 +1437,15 @@ final class GenerativeModelTests: XCTestCase { ) model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), generationConfig: generationConfig, tools: [Tool(functionDeclarations: [sumFunction])], systemInstruction: systemInstruction, requestOptions: RequestOptions(), - appCheck: nil, - auth: nil, urlSession: urlSession ) @@ -1489,13 +1497,13 @@ final class GenerativeModelTests: XCTestCase { let requestOptions = RequestOptions(timeout: expectedTimeout) model = GenerativeModel( name: testModelResourceName, - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID", + firebaseInfo: FirebaseInfo( + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: "My app ID" + ), tools: nil, requestOptions: requestOptions, - appCheck: nil, - auth: nil, urlSession: urlSession ) diff --git a/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift b/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift index 5e685dd98bd..832d56f9cb4 100644 --- a/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift +++ b/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift @@ -50,8 +50,8 @@ class VertexComponentTests: XCTestCase { let vertex = VertexAI.vertexAI(app: VertexComponentTests.app, location: location) XCTAssertNotNil(vertex) - XCTAssertEqual(vertex.projectID, VertexComponentTests.projectID) - XCTAssertEqual(vertex.apiKey, VertexComponentTests.apiKey) + XCTAssertEqual(vertex.firebaseInfo.projectID, VertexComponentTests.projectID) + XCTAssertEqual(vertex.firebaseInfo.apiKey, VertexComponentTests.apiKey) XCTAssertEqual(vertex.location, location) } @@ -121,12 +121,12 @@ class VertexComponentTests: XCTestCase { let app = try XCTUnwrap(VertexComponentTests.app) let vertex = VertexAI.vertexAI(app: app, location: location) let model = "test-model-name" - let modelResourceName = vertex.modelResourceName(modelName: model) + let projectID = vertex.firebaseInfo.projectID XCTAssertEqual( modelResourceName, - "projects/\(vertex.projectID)/locations/\(vertex.location)/publishers/google/models/\(model)" + "projects/\(projectID)/locations/\(vertex.location)/publishers/google/models/\(model)" ) } From 36c02355c9e8e3a29fe99e65a8fb9dc674ad819d Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 23 Feb 2025 09:50:25 -0800 Subject: [PATCH 4/6] Add test helper function --- .../Tests/Unit/GenerativeModelTests.swift | 105 +++++------------- 1 file changed, 27 insertions(+), 78 deletions(-) diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index d990c5d586a..ae10051673a 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -68,11 +68,7 @@ final class GenerativeModelTests: XCTestCase { urlSession = try XCTUnwrap(URLSession(configuration: configuration)) model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -270,11 +266,7 @@ final class GenerativeModelTests: XCTestCase { let model = GenerativeModel( // Model name is prefixed with "models/". name: "models/test-model", - firebaseInfo: FirebaseInfo( - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -391,12 +383,7 @@ final class GenerativeModelTests: XCTestCase { let appCheckToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - appCheck: AppCheckInteropFake(token: appCheckToken), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(appCheck: AppCheckInteropFake(token: appCheckToken)), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -415,12 +402,8 @@ final class GenerativeModelTests: XCTestCase { let appCheckToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - appCheck: AppCheckInteropFake(token: appCheckToken), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: nil - ), + firebaseInfo: testFirebaseInfo(appCheck: AppCheckInteropFake(token: appCheckToken), + privateAppID: true), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -439,12 +422,7 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContent_appCheck_tokenRefreshError() async throws { model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(appCheck: AppCheckInteropFake(error: AppCheckErrorFake())), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -463,12 +441,7 @@ final class GenerativeModelTests: XCTestCase { let authToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - auth: AuthInteropFake(token: authToken), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(auth: AuthInteropFake(token: authToken)), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -486,12 +459,7 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContent_auth_nilAuthToken() async throws { model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - auth: AuthInteropFake(token: nil), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(auth: AuthInteropFake(token: nil)), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -509,12 +477,7 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContent_auth_authTokenRefreshError() async throws { model = GenerativeModel( name: "my-model", - firebaseInfo: FirebaseInfo( - auth: AuthInteropFake(error: AuthErrorFake()), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(auth: AuthInteropFake(error: AuthErrorFake())), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -893,11 +856,7 @@ final class GenerativeModelTests: XCTestCase { let requestOptions = RequestOptions(timeout: expectedTimeout) model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(), tools: nil, requestOptions: requestOptions, urlSession: urlSession @@ -1189,12 +1148,7 @@ final class GenerativeModelTests: XCTestCase { let appCheckToken = "test-valid-token" model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - appCheck: AppCheckInteropFake(token: appCheckToken), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(appCheck: AppCheckInteropFake(token: appCheckToken)), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -1213,12 +1167,7 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContentStream_appCheck_tokenRefreshError() async throws { model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - appCheck: AppCheckInteropFake(error: AppCheckErrorFake()), - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(appCheck: AppCheckInteropFake(error: AppCheckErrorFake())), tools: nil, requestOptions: RequestOptions(), urlSession: urlSession @@ -1361,11 +1310,7 @@ final class GenerativeModelTests: XCTestCase { let requestOptions = RequestOptions(timeout: expectedTimeout) model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(), tools: nil, requestOptions: requestOptions, urlSession: urlSession @@ -1437,11 +1382,7 @@ final class GenerativeModelTests: XCTestCase { ) model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(), generationConfig: generationConfig, tools: [Tool(functionDeclarations: [sumFunction])], systemInstruction: systemInstruction, @@ -1497,11 +1438,7 @@ final class GenerativeModelTests: XCTestCase { let requestOptions = RequestOptions(timeout: expectedTimeout) model = GenerativeModel( name: testModelResourceName, - firebaseInfo: FirebaseInfo( - projectID: "my-project-id", - apiKey: "API_KEY", - googleAppID: "My app ID" - ), + firebaseInfo: testFirebaseInfo(), tools: nil, requestOptions: requestOptions, urlSession: urlSession @@ -1514,6 +1451,18 @@ final class GenerativeModelTests: XCTestCase { // MARK: - Helpers + private func testFirebaseInfo(appCheck: AppCheckInterop? = nil, + auth: AuthInterop? = nil, + privateAppID: Bool = false) -> FirebaseInfo { + FirebaseInfo( + appCheck: appCheck, + auth: auth, + projectID: "my-project-id", + apiKey: "API_KEY", + googleAppID: privateAppID ? nil : "My app ID" + ) + } + private func nonHTTPRequestHandler() throws -> ((URLRequest) -> ( URLResponse, AsyncLineSequence? From c9edcb4e6f757a419102d6a6ee2f31454d7b988e Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Sun, 23 Feb 2025 14:18:29 -0800 Subject: [PATCH 5/6] review --- FirebaseVertexAI/Sources/GenerativeAIService.swift | 2 +- FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/FirebaseVertexAI/Sources/GenerativeAIService.swift b/FirebaseVertexAI/Sources/GenerativeAIService.swift index 15c878dc33c..a0aedd7dad9 100644 --- a/FirebaseVertexAI/Sources/GenerativeAIService.swift +++ b/FirebaseVertexAI/Sources/GenerativeAIService.swift @@ -194,7 +194,7 @@ struct GenerativeAIService { } if let googleAppID = firebaseInfo.googleAppID { - urlRequest.setValue(googleAppID, forHTTPHeaderField: "google-app-id") + urlRequest.setValue(googleAppID, forHTTPHeaderField: "X-Firebase-AppId") } let encoder = JSONEncoder() diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index ae10051673a..39e4ce6d26f 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -1510,7 +1510,7 @@ final class GenerativeModelTests: XCTestCase { XCTAssert(apiClientTags.contains(GenerativeAIService.languageTag)) XCTAssert(apiClientTags.contains(GenerativeAIService.firebaseVersionTag)) XCTAssertEqual(request.value(forHTTPHeaderField: "X-Firebase-AppCheck"), appCheckToken) - let googleAppID = request.value(forHTTPHeaderField: "google-app-id") + let googleAppID = request.value(forHTTPHeaderField: "X-Firebase-AppId") XCTAssertEqual(googleAppID, dataCollection ? "My app ID" : nil) if let authToken { XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Firebase \(authToken)") From b189b5c96294ed9f4af16ed4359977038278dd42 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Mon, 24 Feb 2025 19:50:35 -0800 Subject: [PATCH 6/6] Check latest isDataCollectionDefaultEnabled value during request --- FirebaseVertexAI/Sources/FirebaseInfo.swift | 7 +++++-- FirebaseVertexAI/Sources/GenerativeAIService.swift | 4 ++-- FirebaseVertexAI/Sources/VertexAI.swift | 3 ++- FirebaseVertexAI/Tests/Unit/ChatTests.swift | 7 ++++++- FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift | 9 +++++++-- 5 files changed, 22 insertions(+), 8 deletions(-) diff --git a/FirebaseVertexAI/Sources/FirebaseInfo.swift b/FirebaseVertexAI/Sources/FirebaseInfo.swift index a6677fcd0f8..cddebffadde 100644 --- a/FirebaseVertexAI/Sources/FirebaseInfo.swift +++ b/FirebaseVertexAI/Sources/FirebaseInfo.swift @@ -25,17 +25,20 @@ struct FirebaseInfo { let auth: AuthInterop? let projectID: String let apiKey: String - let googleAppID: String? + let googleAppID: String + let app: FirebaseApp init(appCheck: AppCheckInterop? = nil, auth: AuthInterop? = nil, projectID: String, apiKey: String, - googleAppID: String?) { + googleAppID: String, + firebaseApp: FirebaseApp) { self.appCheck = appCheck self.auth = auth self.projectID = projectID self.apiKey = apiKey self.googleAppID = googleAppID + app = firebaseApp } } diff --git a/FirebaseVertexAI/Sources/GenerativeAIService.swift b/FirebaseVertexAI/Sources/GenerativeAIService.swift index a0aedd7dad9..b3f150b1acb 100644 --- a/FirebaseVertexAI/Sources/GenerativeAIService.swift +++ b/FirebaseVertexAI/Sources/GenerativeAIService.swift @@ -193,8 +193,8 @@ struct GenerativeAIService { urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization") } - if let googleAppID = firebaseInfo.googleAppID { - urlRequest.setValue(googleAppID, forHTTPHeaderField: "X-Firebase-AppId") + if firebaseInfo.app.isDataCollectionDefaultEnabled { + urlRequest.setValue(firebaseInfo.googleAppID, forHTTPHeaderField: "X-Firebase-AppId") } let encoder = JSONEncoder() diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift index fd8931ecf27..097f3230ec4 100644 --- a/FirebaseVertexAI/Sources/VertexAI.swift +++ b/FirebaseVertexAI/Sources/VertexAI.swift @@ -172,7 +172,8 @@ public class VertexAI { auth: ComponentType.instance(for: AuthInterop.self, in: app.container), projectID: projectID, apiKey: apiKey, - googleAppID: app.isDataCollectionDefaultEnabled ? app.options.googleAppID : nil + googleAppID: app.options.googleAppID, + firebaseApp: app ) self.location = location } diff --git a/FirebaseVertexAI/Tests/Unit/ChatTests.swift b/FirebaseVertexAI/Tests/Unit/ChatTests.swift index b0640bd5c5d..a0525880da1 100644 --- a/FirebaseVertexAI/Tests/Unit/ChatTests.swift +++ b/FirebaseVertexAI/Tests/Unit/ChatTests.swift @@ -15,6 +15,7 @@ import Foundation import XCTest +import FirebaseCore @testable import FirebaseVertexAI @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) @@ -53,12 +54,16 @@ final class ChatTests: XCTestCase { return (response, fileURL.lines) } + let app = FirebaseApp(instanceWithName: "testApp", + options: FirebaseOptions(googleAppID: "ignore", + gcmSenderID: "ignore")) let model = GenerativeModel( name: "my-model", firebaseInfo: FirebaseInfo( projectID: "my-project-id", apiKey: "API_KEY", - googleAppID: "My app ID" + googleAppID: "My app ID", + firebaseApp: app ), tools: nil, requestOptions: RequestOptions(), diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index 39e4ce6d26f..10a7d41d793 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -1454,12 +1454,17 @@ final class GenerativeModelTests: XCTestCase { private func testFirebaseInfo(appCheck: AppCheckInterop? = nil, auth: AuthInterop? = nil, privateAppID: Bool = false) -> FirebaseInfo { - FirebaseInfo( + let app = FirebaseApp(instanceWithName: "testApp", + options: FirebaseOptions(googleAppID: "ignore", + gcmSenderID: "ignore")) + app.isDataCollectionDefaultEnabled = !privateAppID + return FirebaseInfo( appCheck: appCheck, auth: auth, projectID: "my-project-id", apiKey: "API_KEY", - googleAppID: privateAppID ? nil : "My app ID" + googleAppID: "My app ID", + firebaseApp: app ) }