Skip to content

Commit 6bf9def

Browse files
andrewheardandrewheard
committed
[Vertex AI] Add enum integration test and run on Developer API v1
1 parent 0fcadb9 commit 6bf9def

File tree

9 files changed

+184
-52
lines changed

9 files changed

+184
-52
lines changed

‎FirebaseCore/Sources/Public/FirebaseCore/FIROptions.h

+2-2
Original file line numberDiff line numberDiff line change
@@ -102,12 +102,12 @@ NS_SWIFT_NAME(FirebaseOptions)
102102
* This will read the file synchronously from disk.
103103
* For example:
104104
* ```swift
105-
* if let path = Bundle.main.path(forResource:"GoogleServices-Info", ofType:"plist") {
105+
* if let path = Bundle.main.path(forResource:"GoogleService-Info", ofType:"plist") {
106106
* let options = FirebaseOptions(contentsOfFile: path)
107107
* }
108108
* ```
109109
* Note that it is not possible to customize `FirebaseOptions` for Firebase Analytics which expects
110-
* a static file named `GoogleServices-Info.plist` -
110+
* a static file named `GoogleService-Info.plist` -
111111
* https://github.com/firebase/firebase-ios-sdk/issues/230.
112112
* Returns `nil` if the plist file does not exist or is invalid.
113113
*/

‎FirebaseVertexAI/Tests/TestApp/Sources/Constants.swift

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
public enum FirebaseAppNames {
16+
/// The name, or a substring of the name, of Firebase apps where App Check is not configured.
17+
public static let appCheckNotConfigured = "app-check-not-configured"
18+
19+
/// The name of a Firebase app with no billing account (i.e., the "Spark" plan).
20+
public static let spark = "spark"
21+
}
22+
23+
public enum ModelNames {
24+
public static let gemini2FlashLite = "gemini-2.0-flash-lite-001"
25+
}

‎FirebaseVertexAI/Tests/TestApp/Sources/TestApp.swift

+12
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,19 @@ import SwiftUI
2020
struct TestApp: App {
2121
init() {
2222
AppCheck.setAppCheckProviderFactory(TestAppCheckProviderFactory())
23+
24+
// Configure default Firebase App
2325
FirebaseApp.configure()
26+
27+
// Configure a Firebase App without a billing account (i.e., the "Spark" plan).
28+
guard let plistPath =
29+
Bundle.main.path(forResource: "GoogleService-Info-Spark", ofType: "plist") else {
30+
fatalError("The file 'GoogleService-Info-Spark.plist' was not found.")
31+
}
32+
guard let options = FirebaseOptions(contentsOfFile: plistPath) else {
33+
fatalError("Failed to parse options from 'GoogleService-Info-Spark.plist'.")
34+
}
35+
FirebaseApp.configure(name: FirebaseAppNames.spark, options: options)
2436
}
2537

2638
var body: some Scene {

‎FirebaseVertexAI/Tests/TestApp/Sources/TestAppCheckProviderFactory.swift

+1-4
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,9 @@ import Foundation
2222
/// ``notConfiguredName``, in which case App Check is not configured; this facilitates integration
2323
/// testing of App Check failure cases.
2424
public class TestAppCheckProviderFactory: NSObject, AppCheckProviderFactory {
25-
/// The name, or a substring of the name, of Firebase apps where App Check is not configured.
26-
public static let notConfiguredName = "app-check-not-configured"
27-
2825
/// Returns the `AppCheckDebugProvider` unless `app.name` contains ``notConfiguredName``.
2926
public func createProvider(with app: FirebaseApp) -> (any AppCheckProvider)? {
30-
if app.name.contains(TestAppCheckProviderFactory.notConfiguredName) {
27+
if app.name.contains(FirebaseAppNames.appCheckNotConfigured) {
3128
return nil
3229
}
3330

‎FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

+64-43
Original file line numberDiff line numberDiff line change
@@ -23,25 +23,28 @@ import VertexAITestApp
2323

2424
@Suite(.serialized)
2525
struct GenerateContentIntegrationTests {
26-
static let vertexV1Config = APIConfig(service: .vertexAI, version: .v1)
27-
static let vertexV1BetaConfig = APIConfig(service: .vertexAI, version: .v1beta)
28-
static let developerV1BetaConfig = APIConfig(
29-
service: .developer(endpoint: .generativeLanguage),
30-
version: .v1beta
26+
static let vertexV1Config =
27+
InstanceConfig(apiConfig: APIConfig(service: .vertexAI, version: .v1))
28+
static let vertexV1BetaConfig =
29+
InstanceConfig(apiConfig: APIConfig(service: .vertexAI, version: .v1beta))
30+
static let developerV1Config = InstanceConfig(
31+
appName: FirebaseAppNames.spark,
32+
apiConfig: APIConfig(
33+
service: .developer(endpoint: .generativeLanguage), version: .v1
34+
)
35+
)
36+
static let developerV1BetaConfig = InstanceConfig(
37+
appName: FirebaseAppNames.spark,
38+
apiConfig: APIConfig(
39+
service: .developer(endpoint: .generativeLanguage), version: .v1beta
40+
)
3141
)
42+
static let allConfigs =
43+
[vertexV1Config, vertexV1BetaConfig, developerV1Config, developerV1BetaConfig]
3244

3345
// Set temperature, topP and topK to lowest allowed values to make responses more deterministic.
34-
static let generationConfig = GenerationConfig(
35-
temperature: 0.0,
36-
topP: 0.0,
37-
topK: 1,
38-
responseMIMEType: "text/plain"
39-
)
40-
static let systemInstruction = ModelContent(
41-
role: "system",
42-
parts: "You are a friendly and helpful assistant."
43-
)
44-
static let safetySettings = [
46+
let generationConfig = GenerationConfig(temperature: 0.0, topP: 0.0, topK: 1)
47+
let safetySettings = [
4548
SafetySetting(harmCategory: .harassment, threshold: .blockLowAndAbove),
4649
SafetySetting(harmCategory: .hateSpeech, threshold: .blockLowAndAbove),
4750
SafetySetting(harmCategory: .sexuallyExplicit, threshold: .blockLowAndAbove),
@@ -64,9 +67,13 @@ struct GenerateContentIntegrationTests {
6467
storage = Storage.storage()
6568
}
6669

67-
@Test(arguments: [vertexV1Config, vertexV1BetaConfig, developerV1BetaConfig])
68-
func generateContent(_ apiConfig: APIConfig) async throws {
69-
let model = GenerateContentIntegrationTests.model(apiConfig: apiConfig)
70+
@Test(arguments: allConfigs)
71+
func generateContent(_ config: InstanceConfig) async throws {
72+
let model = VertexAI.componentInstance(config).generativeModel(
73+
modelName: ModelNames.gemini2FlashLite,
74+
generationConfig: generationConfig,
75+
safetySettings: safetySettings
76+
)
7077
let prompt = "Where is Google headquarters located? Answer with the city name only."
7178

7279
let response = try await model.generateContent(prompt)
@@ -75,9 +82,9 @@ struct GenerateContentIntegrationTests {
7582
#expect(text == "Mountain View")
7683

7784
let usageMetadata = try #require(response.usageMetadata)
78-
#expect(usageMetadata.promptTokenCount == 21)
85+
#expect(usageMetadata.promptTokenCount == 13)
7986
#expect(usageMetadata.candidatesTokenCount.isEqual(to: 3, accuracy: tokenCountAccuracy))
80-
#expect(usageMetadata.totalTokenCount.isEqual(to: 24, accuracy: tokenCountAccuracy))
87+
#expect(usageMetadata.totalTokenCount.isEqual(to: 16, accuracy: tokenCountAccuracy))
8188
#expect(usageMetadata.promptTokensDetails.count == 1)
8289
let promptTokensDetails = try #require(usageMetadata.promptTokensDetails.first)
8390
#expect(promptTokensDetails.modality == .text)
@@ -88,31 +95,45 @@ struct GenerateContentIntegrationTests {
8895
#expect(candidatesTokensDetails.tokenCount == usageMetadata.candidatesTokenCount)
8996
}
9097

91-
static func model(apiConfig: APIConfig) -> GenerativeModel {
92-
return instance(apiConfig: apiConfig).generativeModel(
93-
modelName: "gemini-2.0-flash",
94-
generationConfig: generationConfig,
98+
@Test(
99+
"Generate an enum and provide a system instruction",
100+
arguments: [
101+
vertexV1Config,
102+
vertexV1BetaConfig,
103+
/* System instructions are not supported on the v1 Developer API. */
104+
developerV1BetaConfig,
105+
]
106+
)
107+
func generateContentEnum(_ config: InstanceConfig) async throws {
108+
let model = VertexAI.componentInstance(config).generativeModel(
109+
modelName: ModelNames.gemini2FlashLite,
110+
generationConfig: GenerationConfig(
111+
responseMIMEType: "text/x.enum", // Not supported on the v1 Developer API
112+
responseSchema: .enumeration(values: ["Red", "Green", "Blue"])
113+
),
95114
safetySettings: safetySettings,
96-
tools: [],
97-
toolConfig: .init(functionCallingConfig: .none()),
98-
systemInstruction: systemInstruction
115+
tools: [], // Not supported on the v1 Developer API
116+
toolConfig: .init(functionCallingConfig: .none()), // Not supported on the v1 Developer API
117+
systemInstruction: ModelContent(role: "system", parts: "Always pick blue.")
99118
)
100-
}
119+
let prompt = "What is your favourite colour?"
101120

102-
// TODO(andrewheard): Move this helper to a file in the Utilities folder.
103-
static func instance(apiConfig: APIConfig) -> VertexAI {
104-
switch apiConfig.service {
105-
case .vertexAI:
106-
return VertexAI.vertexAI(app: nil, location: "us-central1", apiConfig: apiConfig)
107-
case .developer:
108-
return VertexAI.vertexAI(app: nil, location: nil, apiConfig: apiConfig)
109-
}
110-
}
111-
}
121+
let response = try await model.generateContent(prompt)
122+
123+
let text = try #require(response.text).trimmingCharacters(in: .whitespacesAndNewlines)
124+
#expect(text == "Blue")
112125

113-
// TODO(andrewheard): Move this extension to a file in the Utilities folder.
114-
extension Numeric where Self: Strideable, Self.Stride.Magnitude: Comparable {
115-
func isEqual(to other: Self, accuracy: Self.Stride) -> Bool {
116-
return distance(to: other).magnitude < accuracy.magnitude
126+
let usageMetadata = try #require(response.usageMetadata)
127+
#expect(usageMetadata.promptTokenCount == 14)
128+
#expect(usageMetadata.candidatesTokenCount.isEqual(to: 1, accuracy: tokenCountAccuracy))
129+
#expect(usageMetadata.totalTokenCount.isEqual(to: 15, accuracy: tokenCountAccuracy))
130+
#expect(usageMetadata.promptTokensDetails.count == 1)
131+
let promptTokensDetails = try #require(usageMetadata.promptTokensDetails.first)
132+
#expect(promptTokensDetails.modality == .text)
133+
#expect(promptTokensDetails.tokenCount == usageMetadata.promptTokenCount)
134+
#expect(usageMetadata.candidatesTokensDetails.count == 1)
135+
let candidatesTokensDetails = try #require(usageMetadata.candidatesTokensDetails.first)
136+
#expect(candidatesTokensDetails.modality == .text)
137+
#expect(candidatesTokensDetails.tokenCount == usageMetadata.candidatesTokenCount)
117138
}
118139
}

‎FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

+2-2
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ final class IntegrationTests: XCTestCase {
109109
}
110110

111111
func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
112-
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
112+
let app = try FirebaseApp.defaultNamedCopy(name: FirebaseAppNames.appCheckNotConfigured)
113113
addTeardownBlock { await app.delete() }
114114
let vertex = VertexAI.vertexAI(app: app)
115115
let model = vertex.generativeModel(modelName: "gemini-2.0-flash")
@@ -285,7 +285,7 @@ final class IntegrationTests: XCTestCase {
285285
}
286286

287287
func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
288-
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
288+
let app = try FirebaseApp.defaultNamedCopy(name: FirebaseAppNames.appCheckNotConfigured)
289289
addTeardownBlock { await app.delete() }
290290
let vertex = VertexAI.vertexAI(app: app)
291291
let model = vertex.generativeModel(modelName: "gemini-2.0-flash")

‎FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTestUtils.swift ‎FirebaseVertexAI/Tests/TestApp/Tests/Utilities/IntegrationTestUtils.swift

+6
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,9 @@ enum IntegrationTestUtils {
3737
}
3838
}
3939
}
40+
41+
extension Numeric where Self: Strideable, Self.Stride.Magnitude: Comparable {
42+
func isEqual(to other: Self, accuracy: Self.Stride) -> Bool {
43+
return distance(to: other).magnitude < accuracy.magnitude
44+
}
45+
}

‎FirebaseVertexAI/Tests/TestApp/Tests/Utilities/VertexAITestUtils.swift

+59
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import FirebaseCore
16+
import VertexAITestApp
17+
18+
@testable import struct FirebaseVertexAI.APIConfig
19+
@testable import class FirebaseVertexAI.VertexAI
20+
21+
struct InstanceConfig {
22+
let appName: String?
23+
let location: String?
24+
let apiConfig: APIConfig
25+
26+
init(appName: String? = nil, location: String? = nil, apiConfig: APIConfig) {
27+
self.appName = appName
28+
self.location = location
29+
self.apiConfig = apiConfig
30+
}
31+
32+
var app: FirebaseApp? {
33+
return appName.map { FirebaseApp.app(name: $0) } ?? FirebaseApp.app()
34+
}
35+
}
36+
37+
extension VertexAI {
38+
static func componentInstance(_ instanceConfig: InstanceConfig) -> VertexAI {
39+
switch instanceConfig.apiConfig.service {
40+
case .vertexAI:
41+
let location = instanceConfig.location ?? "us-central1"
42+
return VertexAI.vertexAI(
43+
app: instanceConfig.app,
44+
location: location,
45+
apiConfig: instanceConfig.apiConfig
46+
)
47+
case .developer:
48+
assert(
49+
instanceConfig.location == nil,
50+
"The Developer API is global and does not support `location`."
51+
)
52+
return VertexAI.vertexAI(
53+
app: instanceConfig.app,
54+
location: nil,
55+
apiConfig: instanceConfig.apiConfig
56+
)
57+
}
58+
}
59+
}

‎FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj

+13-1
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@
2323
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */; };
2424
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */; };
2525
86D77DFC2D7A5340003D155D /* GenerateContentIntegrationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */; };
26+
86D77DFE2D7B5C86003D155D /* GoogleService-Info-Spark.plist in Resources */ = {isa = PBXBuildFile; fileRef = 86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */; };
27+
86D77E022D7B63AF003D155D /* Constants.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77E012D7B63AC003D155D /* Constants.swift */; };
28+
86D77E042D7B6C9D003D155D /* VertexAITestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77E032D7B6C95003D155D /* VertexAITestUtils.swift */; };
2629
/* End PBXBuildFile section */
2730

2831
/* Begin PBXContainerItemProxy section */
@@ -51,6 +54,9 @@
5154
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirebaseAppTestUtils.swift; sourceTree = "<group>"; };
5255
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TestAppCheckProviderFactory.swift; sourceTree = "<group>"; };
5356
86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = GenerateContentIntegrationTests.swift; sourceTree = "<group>"; };
57+
86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = "GoogleService-Info-Spark.plist"; sourceTree = "<group>"; };
58+
86D77E012D7B63AC003D155D /* Constants.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Constants.swift; sourceTree = "<group>"; };
59+
86D77E032D7B6C95003D155D /* VertexAITestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = VertexAITestUtils.swift; sourceTree = "<group>"; };
5460
/* End PBXFileReference section */
5561

5662
/* Begin PBXFrameworksBuildPhase section */
@@ -101,6 +107,7 @@
101107
868A7C532CCC26B500E449DD /* Assets.xcassets */,
102108
868A7C552CCC271300E449DD /* TestApp.entitlements */,
103109
868A7C462CCA931B00E449DD /* GoogleService-Info.plist */,
110+
86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */,
104111
);
105112
path = Resources;
106113
sourceTree = "<group>";
@@ -119,6 +126,7 @@
119126
8661385B2CC943DD00F4B78E /* TestApp.swift */,
120127
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */,
121128
8661385D2CC943DD00F4B78E /* ContentView.swift */,
129+
86D77E012D7B63AC003D155D /* Constants.swift */,
122130
);
123131
path = Sources;
124132
sourceTree = "<group>";
@@ -130,7 +138,6 @@
130138
8661386D2CC943DE00F4B78E /* IntegrationTests.swift */,
131139
86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */,
132140
864F8F702D4980D60002EA7E /* ImagenIntegrationTests.swift */,
133-
862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */,
134141
);
135142
path = Integration;
136143
sourceTree = "<group>";
@@ -147,7 +154,9 @@
147154
8698D7442CD3CEF700ABA833 /* Utilities */ = {
148155
isa = PBXGroup;
149156
children = (
157+
86D77E032D7B6C95003D155D /* VertexAITestUtils.swift */,
150158
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */,
159+
862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */,
151160
);
152161
path = Utilities;
153162
sourceTree = "<group>";
@@ -243,6 +252,7 @@
243252
buildActionMask = 2147483647;
244253
files = (
245254
868A7C522CCC263300E449DD /* Preview Assets.xcassets in Resources */,
255+
86D77DFE2D7B5C86003D155D /* GoogleService-Info-Spark.plist in Resources */,
246256
868A7C542CCC26B500E449DD /* Assets.xcassets in Resources */,
247257
868A7C482CCA931B00E449DD /* GoogleService-Info.plist in Resources */,
248258
);
@@ -265,13 +275,15 @@
265275
8661385E2CC943DD00F4B78E /* ContentView.swift in Sources */,
266276
8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */,
267277
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */,
278+
86D77E022D7B63AF003D155D /* Constants.swift in Sources */,
268279
);
269280
runOnlyForDeploymentPostprocessing = 0;
270281
};
271282
866138652CC943DE00F4B78E /* Sources */ = {
272283
isa = PBXSourcesBuildPhase;
273284
buildActionMask = 2147483647;
274285
files = (
286+
86D77E042D7B6C9D003D155D /* VertexAITestUtils.swift in Sources */,
275287
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */,
276288
868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */,
277289
864F8F712D4980DD0002EA7E /* ImagenIntegrationTests.swift in Sources */,

0 commit comments

Comments
 (0)