Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 6 commits into from
Feb 25, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@ -43,7 +44,7 @@
#if os(watchOS)
throw XCTSkip("Custom URL protocols are unsupported in watchOS 2 and later.")
#endif // os(watchOS)
MockURLProtocol.requestHandler = { request in

Check warning on line 47 in FirebaseVertexAI/Tests/Unit/ChatTests.swift

View workflow job for this annotation

GitHub Actions / spm-unit (macos-15, Xcode_16.2, watchOS)

code after 'throw' will never be executed

Check warning on line 47 in FirebaseVertexAI/Tests/Unit/ChatTests.swift

View workflow job for this annotation

GitHub Actions / spm-unit (macos-15, Xcode_16.2, watchOS)

code after 'throw' will never be executed
let response = HTTPURLResponse(
url: request.url!,
statusCode: 200,
Expand All @@ -53,14 +54,19 @@
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
Loading