Add google-app-id to Vertex AI requests (#14479)

Author: paulb777

FirebaseVertexAI/Sources/FirebaseInfo.swift

@@ -0,0 +1,44 @@
+// 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
+  let app: FirebaseApp
+
+  init(appCheck: AppCheckInterop? = nil,
+       auth: AuthInterop? = nil,
+       projectID: String,
+       apiKey: 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

@@ -26,23 +26,12 @@ struct GenerativeAIService {
   /// The Firebase SDK version in the format `fire/<version>`.
   static let firebaseVersionTag = "fire/\(FirebaseVersion())"
 
-  private let projectID: String
-
-  /// Gives permission to talk to the backend.
-  private let apiKey: String
-
-  private let appCheck: AppCheckInterop?
-
-  private let auth: AuthInterop?
+  private let firebaseInfo: FirebaseInfo
 
   private let urlSession: URLSession
 
-  init(projectID: String, apiKey: String, appCheck: AppCheckInterop?, auth: AuthInterop?,
-       urlSession: URLSession) {
-    self.projectID = projectID
-    self.apiKey = apiKey
-    self.appCheck = appCheck
-    self.auth = auth
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
+    self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
   }
 
@@ -180,14 +169,14 @@ struct GenerativeAIService {
   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(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 {
@@ -198,10 +187,16 @@ 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 firebaseInfo.app.isDataCollectionDefaultEnabled {
+      urlRequest.setValue(firebaseInfo.googleAppID, forHTTPHeaderField: "X-Firebase-AppId")
+    }
+
     let encoder = JSONEncoder()
     urlRequest.httpBody = try encoder.encode(request)
     urlRequest.timeoutInterval = request.options.timeout
@@ -260,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 \

FirebaseVertexAI/Sources/GenerativeModel.swift

@@ -59,23 +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,
+       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,
-      appCheck: appCheck,
-      auth: auth,
+      firebaseInfo: firebaseInfo,
       urlSession: urlSession
     )
     self.generationConfig = generationConfig

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

@@ -41,20 +41,14 @@ public final class ImagenModel {
   let requestOptions: RequestOptions
 
   init(name: String,
-       projectID: String,
-       apiKey: String,
+       firebaseInfo: FirebaseInfo,
        generationConfig: ImagenGenerationConfig?,
        safetySettings: ImagenSafetySettings?,
        requestOptions: RequestOptions,
-       appCheck: AppCheckInterop?,
-       auth: AuthInterop?,
        urlSession: URLSession = .shared) {
     modelResourceName = name
     generativeAIService = GenerativeAIService(
-      projectID: projectID,
-      apiKey: apiKey,
-      appCheck: appCheck,
-      auth: auth,
+      firebaseInfo: firebaseInfo,
       urlSession: urlSession
     )
     self.generationConfig = generationConfig

FirebaseVertexAI/Sources/VertexAI.swift

@@ -91,16 +91,13 @@ public class VertexAI {
     -> GenerativeModel {
     return GenerativeModel(
       name: modelResourceName(modelName: modelName),
-      projectID: projectID,
-      apiKey: apiKey,
+      firebaseInfo: firebaseInfo,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
       tools: tools,
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
-      requestOptions: requestOptions,
-      appCheck: appCheck,
-      auth: auth
+      requestOptions: requestOptions
     )
   }
 
@@ -126,13 +123,10 @@ public class VertexAI {
                           requestOptions: RequestOptions = RequestOptions()) -> ImagenModel {
     return ImagenModel(
       name: modelResourceName(modelName: modelName),
-      projectID: projectID,
-      apiKey: apiKey,
+      firebaseInfo: firebaseInfo,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      requestOptions: requestOptions,
-      appCheck: appCheck,
-      auth: auth
+      requestOptions: requestOptions
     )
   }
 
@@ -142,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
@@ -165,25 +155,26 @@ public class VertexAI {
     private static var instancesLock: os_unfair_lock = .init()
   #endif
 
-  let projectID: String
-  let apiKey: String
   let location: String
 
   init(app: FirebaseApp, location: String) {
-    self.app = app
-    appCheck = ComponentType<AppCheckInterop>.instance(for: AppCheckInterop.self, in: app.container)
-    auth = ComponentType<AuthInterop>.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
-
+    firebaseInfo = FirebaseInfo(
+      appCheck: ComponentType<AppCheckInterop>.instance(
+        for: AppCheckInterop.self,
+        in: app.container
+      ),
+      auth: ComponentType<AuthInterop>.instance(for: AuthInterop.self, in: app.container),
+      projectID: projectID,
+      apiKey: apiKey,
+      googleAppID: app.options.googleAppID,
+      firebaseApp: app
+    )
     self.location = location
   }
 
@@ -205,6 +196,7 @@ public class VertexAI {
       """)
     }
 
+    let projectID = firebaseInfo.projectID
     return "projects/\(projectID)/locations/\(location)/publishers/google/models/\(modelName)"
   }
 }

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,14 +54,19 @@ 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",
-      projectID: "my-project-id",
-      apiKey: "API_KEY",
+      firebaseInfo: FirebaseInfo(
+        projectID: "my-project-id",
+        apiKey: "API_KEY",
+        googleAppID: "My app ID",
+        firebaseApp: app
+      ),
       tools: nil,
       requestOptions: RequestOptions(),
-      appCheck: nil,
-      auth: nil,
       urlSession: urlSession
     )
     let chat = Chat(model: model, history: [])