Check latest isDataCollectionDefaultEnabled value during request

Author: paulb777

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
   }
 }

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()

FirebaseVertexAI/Sources/VertexAI.swift

@@ -172,7 +172,8 @@ public class VertexAI {
       auth: ComponentType<AuthInterop>.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
   }

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(),

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
     )
   }