Skip to content

Commit f0555c6

Browse files
andrewheardandrewheard
authored
[Firebase AI] Add support for direct Vertex AI integration testing (#15498)
1 parent 07bdccf commit f0555c6

File tree

7 files changed

+92
-26
lines changed

7 files changed

+92
-26
lines changed

‎FirebaseAI/Sources/Constants.swift‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,15 @@ enum Constants {
2121
/// - Important: A suffix must be appended to produce an error domain (e.g.,
2222
/// "com.google.firebase.firebaseai.ExampleError").
2323
static let baseErrorDomain = "com.google.firebase.firebaseai"
24+
25+
#if DEBUG
26+
/// The key for an environment variable containing a Google Cloud Access Token.
27+
///
28+
/// This should only be used for SDK development and testing with the Vertex AI direct backend
29+
/// that bypasses the Firebase proxy..
30+
///
31+
/// The value should is typically obtained from the gcloud CLI by calling
32+
/// `gcloud auth print-access-token`.
33+
static let gCloudAccessTokenEnvVarKey = "FIRGCloudAuthAccessToken"
34+
#endif // DEBUG
2435
}

‎FirebaseAI/Sources/FirebaseAI.swift‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -299,11 +299,18 @@ public final class FirebaseAI: Sendable {
299299

300300
private func developerModelResourceName(modelName: String) -> String {
301301
switch apiConfig.service.endpoint {
302-
case .firebaseProxyStaging, .firebaseProxyProd:
303-
let projectID = firebaseInfo.projectID
304-
return "projects/\(projectID)/models/\(modelName)"
305-
case .googleAIBypassProxy:
306-
return "models/\(modelName)"
302+
case .firebaseProxyProd:
303+
return "projects/\(firebaseInfo.projectID)/models/\(modelName)"
304+
#if DEBUG
305+
case .googleAIBypassProxy:
306+
return "models/\(modelName)"
307+
case .firebaseProxyStaging:
308+
return "projects/\(firebaseInfo.projectID)/models/\(modelName)"
309+
case .vertexAIStagingBypassProxy:
310+
fatalError(
311+
"The Vertex AI staging endpoint does not support the Gemini Developer API (Google AI)."
312+
)
313+
#endif // DEBUG
307314
}
308315
}
309316

‎FirebaseAI/Sources/GenerativeAIService.swift‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,16 @@ struct GenerativeAIService {
169169
private func urlRequest<T: GenerativeAIRequest>(request: T) async throws -> URLRequest {
170170
var urlRequest = try URLRequest(url: request.getURL())
171171
urlRequest.httpMethod = "POST"
172-
urlRequest.setValue(firebaseInfo.apiKey, forHTTPHeaderField: "x-goog-api-key")
172+
#if DEBUG
173+
let accessToken = ProcessInfo.processInfo.environment[Constants.gCloudAccessTokenEnvVarKey]
174+
#else
175+
let accessToken: String? = nil
176+
#endif // DEBUG
177+
if let accessToken {
178+
urlRequest.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
179+
} else {
180+
urlRequest.setValue(firebaseInfo.apiKey, forHTTPHeaderField: "x-goog-api-key")
181+
}
173182
urlRequest.setValue(
174183
"\(GenerativeAIService.languageTag) \(GenerativeAIService.firebaseVersionTag)",
175184
forHTTPHeaderField: "x-goog-api-client"
@@ -190,9 +199,8 @@ struct GenerativeAIService {
190199
}
191200
}
192201

193-
if let auth = firebaseInfo.auth, let authToken = try await auth.getToken(
194-
forcingRefresh: false
195-
) {
202+
if let auth = firebaseInfo.auth, let authToken = try await auth.getToken(forcingRefresh: false),
203+
accessToken == nil {
196204
urlRequest.setValue("Firebase \(authToken)", forHTTPHeaderField: "Authorization")
197205
}
198206

‎FirebaseAI/Sources/GenerativeModel.swift‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -322,11 +322,20 @@ public final class GenerativeModel: Sendable {
322322
// "models/model-name". This field is unaltered by the Firebase backend before forwarding the
323323
// request to the Generative Language backend, which expects the form "models/model-name".
324324
let generateContentRequestModelResourceName = switch apiConfig.service {
325-
case .vertexAI, .googleAI(endpoint: .googleAIBypassProxy):
325+
case .vertexAI:
326326
modelResourceName
327-
case .googleAI(endpoint: .firebaseProxyProd),
328-
.googleAI(endpoint: .firebaseProxyStaging):
327+
case .googleAI(endpoint: .firebaseProxyProd):
329328
"models/\(modelName)"
329+
#if DEBUG
330+
case .googleAI(endpoint: .firebaseProxyStaging):
331+
"models/\(modelName)"
332+
case .googleAI(endpoint: .googleAIBypassProxy):
333+
modelResourceName
334+
case .googleAI(endpoint: .vertexAIStagingBypassProxy):
335+
fatalError(
336+
"The Vertex AI staging endpoint does not support the Gemini Developer API (Google AI)."
337+
)
338+
#endif // DEBUG
330339
}
331340

332341
let generateContentRequest = GenerateContentRequest(

‎FirebaseAI/Sources/Types/Internal/APIConfig.swift‎

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -75,28 +75,43 @@ extension APIConfig.Service {
7575
/// This endpoint supports both Google AI and Vertex AI.
7676
case firebaseProxyProd = "https://firebasevertexai.googleapis.com"
7777

78-
/// The Firebase proxy staging endpoint; for SDK development and testing only.
79-
///
80-
/// This endpoint supports both the Gemini Developer API (commonly referred to as Google AI)
81-
/// and the Gemini API in Vertex AI (commonly referred to simply as Vertex AI).
82-
case firebaseProxyStaging = "https://staging-firebasevertexai.sandbox.googleapis.com"
78+
#if DEBUG
79+
/// The Firebase proxy staging endpoint; for SDK development and testing only.
80+
///
81+
/// This endpoint supports both the Gemini Developer API (commonly referred to as Google AI)
82+
/// and the Gemini API in Vertex AI (commonly referred to simply as Vertex AI).
83+
case firebaseProxyStaging = "https://staging-firebasevertexai.sandbox.googleapis.com"
8384

84-
/// The Gemini Developer API (Google AI) direct production endpoint; for SDK development and
85-
/// testing only.
86-
///
87-
/// This bypasses the Firebase proxy and directly connects to the Gemini Developer API
88-
/// (Google AI) backend. This endpoint only supports the Gemini Developer API, not Vertex AI.
89-
case googleAIBypassProxy = "https://generativelanguage.googleapis.com"
85+
/// The Gemini Developer API (Google AI) direct production endpoint; for SDK development and
86+
/// testing only.
87+
///
88+
/// This bypasses the Firebase proxy and directly connects to the Gemini Developer API
89+
/// (Google AI) backend. This endpoint only supports the Gemini Developer API, not Vertex AI.
90+
case googleAIBypassProxy = "https://generativelanguage.googleapis.com"
91+
92+
/// The Vertex AI direct staging endpoint; for SDK development and testing only.
93+
///
94+
/// This bypasses the Firebase proxy and directly connects to the Vertex AI backend. This
95+
/// endpoint only supports the Gemini API in Vertex AI, not the Gemini Developer API.
96+
case vertexAIStagingBypassProxy = "https://staging-aiplatform.sandbox.googleapis.com"
97+
#endif // DEBUG
9098
}
9199
}
92100

93101
extension APIConfig {
94102
/// Versions of the configured API service (`APIConfig.Service`).
95103
enum Version: String, Encodable {
96-
/// The stable channel for version 1 of the API.
97-
case v1
98-
99104
/// The beta channel for version 1 of the API.
100105
case v1beta
106+
107+
#if DEBUG
108+
/// The stable channel for version 1 of the API; currently for SDK development and testing
109+
/// only.
110+
case v1
111+
112+
/// The beta channel for version 1 of the direct Vertex AI API, when bypassing the Firebase
113+
/// proxy; for SDK development and testing only.
114+
case v1beta1
115+
#endif // DEBUG
101116
}
102117
}

‎FirebaseAI/Tests/TestApp/FirebaseAITestApp.xcodeproj/xcshareddata/xcschemes/FirebaseAITestApp-SPM.xcscheme‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,13 @@
6969
isEnabled = "YES">
7070
</CommandLineArgument>
7171
</CommandLineArguments>
72+
<EnvironmentVariables>
73+
<EnvironmentVariable
74+
key = "FIRGCloudAuthAccessToken"
75+
value = "Run `gcloud auth print-access-token` to obtain as access token for direct Vertex AI testing."
76+
isEnabled = "NO">
77+
</EnvironmentVariable>
78+
</EnvironmentVariables>
7279
</LaunchAction>
7380
<ProfileAction
7481
buildConfiguration = "Release"

‎FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,12 @@ struct InstanceConfig: Equatable, Encodable {
5252
version: .v1beta
5353
)
5454
)
55+
static let vertexAI_v1beta_staging_global_bypassProxy = InstanceConfig(
56+
apiConfig: APIConfig(
57+
service: .vertexAI(endpoint: .vertexAIStagingBypassProxy, location: "global"),
58+
version: .v1beta1
59+
)
60+
)
5561
static let googleAI_v1beta = InstanceConfig(
5662
apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta)
5763
)
@@ -80,6 +86,7 @@ struct InstanceConfig: Equatable, Encodable {
8086
googleAI_v1beta_freeTier,
8187
// Note: The following configs are commented out for easy one-off manual testing.
8288
// vertexAI_v1beta_staging,
89+
// vertexAI_v1beta_staging_global_bypassProxy,
8390
// googleAI_v1beta_staging,
8491
// googleAI_v1beta_freeTier_bypassProxy,
8592
]
@@ -162,6 +169,8 @@ extension InstanceConfig: CustomTestStringConvertible {
162169
" - Staging"
163170
case .googleAIBypassProxy:
164171
" - Bypass Proxy"
172+
case .vertexAIStagingBypassProxy:
173+
" - Staging - Bypass Proxy"
165174
}
166175
let locationSuffix: String
167176
if case let .vertexAI(_, location: location) = apiConfig.service {

0 commit comments

Comments
 (0)