Skip to content

Commit

Permalink
Add google-app-id to Vertex AI requests (#14479)
Browse files Browse the repository at this point in the history
  • Loading branch information
paulb777 authored Feb 25, 2025
1 parent df200d7 commit 3c7b987
Show file tree
Hide file tree
Showing 8 changed files with 148 additions and 120 deletions.
44 changes: 44 additions & 0 deletions FirebaseVertexAI/Sources/FirebaseInfo.swift
Original file line number Diff line number Diff line change
@@ -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
}
}
30 changes: 13 additions & 17 deletions FirebaseVertexAI/Sources/GenerativeAIService.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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 \
Expand Down
10 changes: 2 additions & 8 deletions FirebaseVertexAI/Sources/GenerativeModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 2 additions & 8 deletions FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
44 changes: 18 additions & 26 deletions FirebaseVertexAI/Sources/VertexAI.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
}

Expand All @@ -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
)
}

Expand All @@ -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
Expand All @@ -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
}

Expand All @@ -205,6 +196,7 @@ public class VertexAI {
""")
}

let projectID = firebaseInfo.projectID
return "projects/\(projectID)/locations/\(location)/publishers/google/models/\(modelName)"
}
}
14 changes: 10 additions & 4 deletions FirebaseVertexAI/Tests/Unit/ChatTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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, *)
Expand Down Expand Up @@ -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: [])
Expand Down
Loading

0 comments on commit 3c7b987

Please sign in to comment.