Skip to content

Commit 4d2543a

Browse files
andrewheardandrewheard
committed
[Firebase AI] Implement new public API surface and tests (#14765)
1 parent 8a854c7 commit 4d2543a

14 files changed

+212
-117
lines changed

‎FirebaseAI/Sources/FirebaseAI.swift

+12-16
Original file line numberDiff line numberDiff line change
@@ -25,26 +25,22 @@ internal import FirebaseCoreExtension
2525
public final class FirebaseAI: Sendable {
2626
// MARK: - Public APIs
2727

28-
/// Creates an instance of `VertexAI`.
28+
/// Creates an instance of `FirebaseAI`.
2929
///
3030
/// - Parameters:
3131
/// - app: A custom `FirebaseApp` used for initialization; if not specified, uses the default
3232
/// ``FirebaseApp``.
33-
/// - location: The region identifier, defaulting to `us-central1`; see
34-
/// [Vertex AI locations]
35-
/// (https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations)
36-
/// for a list of supported locations.
37-
/// - Returns: A `VertexAI` instance, configured with the custom `FirebaseApp`.
38-
public static func vertexAI(app: FirebaseApp? = nil,
39-
location: String = "us-central1") -> FirebaseAI {
40-
let vertexInstance = vertexAI(app: app, location: location, apiConfig: defaultVertexAIAPIConfig)
41-
// Verify that the `VertexAI` instance is always configured with the production endpoint since
33+
/// - backend: The backend API for the Firebase AI SDK; if not specified, uses the default
34+
/// ``Backend/googleAI()`` (Gemini Developer API).
35+
/// - Returns: A `FirebaseAI` instance, configured with the custom `FirebaseApp`.
36+
public static func firebaseAI(app: FirebaseApp? = nil,
37+
backend: Backend = .googleAI()) -> FirebaseAI {
38+
let instance = firebaseAI(app: app, location: backend.location, apiConfig: backend.apiConfig)
39+
// Verify that the `FirebaseAI` instance is always configured with the production endpoint since
4240
// this is the public API surface for creating an instance.
43-
assert(vertexInstance.apiConfig.service == .vertexAI(endpoint: .firebaseVertexAIProd))
44-
assert(vertexInstance.apiConfig.service.endpoint == .firebaseVertexAIProd)
45-
assert(vertexInstance.apiConfig.version == .v1beta)
46-
47-
return vertexInstance
41+
assert(instance.apiConfig.service.endpoint == .firebaseVertexAIProd)
42+
assert(instance.apiConfig.version == .v1beta)
43+
return instance
4844
}
4945

5046
/// Initializes a generative model with the given parameters.
@@ -163,7 +159,7 @@ public final class FirebaseAI: Sendable {
163159
version: .v1beta
164160
)
165161

166-
static func vertexAI(app: FirebaseApp?, location: String?, apiConfig: APIConfig) -> FirebaseAI {
162+
static func firebaseAI(app: FirebaseApp?, location: String?, apiConfig: APIConfig) -> FirebaseAI {
167163
guard let app = app ?? FirebaseApp.app() else {
168164
fatalError("No instance of the default Firebase app was found.")
169165
}

‎FirebaseAI/Sources/Types/Public/Backend.swift

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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+
/// Represents available backend APIs for the Firebase AI SDK.
16+
public struct Backend {
17+
// MARK: - Public API
18+
19+
/// Initializes a `Backend` configured for the Gemini API in Vertex AI.
20+
///
21+
/// - Parameters:
22+
/// - location: The region identifier, defaulting to `us-central1`; see
23+
/// [Vertex AI locations]
24+
/// (https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations)
25+
/// for a list of supported locations.
26+
public static func vertexAI(location: String = "us-central1") -> Backend {
27+
return Backend(
28+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1beta),
29+
location: location
30+
)
31+
}
32+
33+
/// Initializes a `Backend` configured for the Google Developer API.
34+
public static func googleAI() -> Backend {
35+
return Backend(
36+
apiConfig: APIConfig(service: .developer(endpoint: .firebaseVertexAIProd), version: .v1beta),
37+
location: nil
38+
)
39+
}
40+
41+
// MARK: - Internal
42+
43+
let apiConfig: APIConfig
44+
let location: String?
45+
46+
init(apiConfig: APIConfig, location: String?) {
47+
self.apiConfig = apiConfig
48+
self.location = location
49+
}
50+
}

‎FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ struct ImagenIntegrationTests {
4040

4141
init() async throws {
4242
userID1 = try await TestHelpers.getUserID()
43-
vertex = FirebaseAI.vertexAI()
43+
vertex = FirebaseAI.firebaseAI(backend: .vertexAI())
4444
storage = Storage.storage()
4545
}
4646

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ final class IntegrationTests: XCTestCase {
4949

5050
override func setUp() async throws {
5151
userID1 = try await TestHelpers.getUserID()
52-
vertex = FirebaseAI.vertexAI()
52+
vertex = FirebaseAI.firebaseAI(backend: .vertexAI())
5353
model = vertex.generativeModel(
5454
modelName: "gemini-2.0-flash",
5555
generationConfig: generationConfig,
@@ -200,7 +200,7 @@ final class IntegrationTests: XCTestCase {
200200

201201
func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
202202
let app = try XCTUnwrap(FirebaseApp.app(name: FirebaseAppNames.appCheckNotConfigured))
203-
let vertex = FirebaseAI.vertexAI(app: app)
203+
let vertex = FirebaseAI.firebaseAI(app: app, backend: .vertexAI())
204204
let model = vertex.generativeModel(modelName: "gemini-2.0-flash")
205205
let prompt = "Why is the sky blue?"
206206

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ extension FirebaseAI {
127127
switch instanceConfig.apiConfig.service {
128128
case .vertexAI:
129129
let location = instanceConfig.location ?? "us-central1"
130-
return FirebaseAI.vertexAI(
130+
return FirebaseAI.firebaseAI(
131131
app: instanceConfig.app,
132132
location: location,
133133
apiConfig: instanceConfig.apiConfig
@@ -137,7 +137,7 @@ extension FirebaseAI {
137137
instanceConfig.location == nil,
138138
"The Developer API is global and does not support `location`."
139139
)
140-
return FirebaseAI.vertexAI(
140+
return FirebaseAI.firebaseAI(
141141
app: instanceConfig.app,
142142
location: nil,
143143
apiConfig: instanceConfig.apiConfig

‎FirebaseAI/Tests/Unit/APITests.swift

+41-68
Original file line numberDiff line numberDiff line change
@@ -41,36 +41,40 @@ final class APITests: XCTestCase {
4141
let requestOptions = RequestOptions()
4242
let _ = RequestOptions(timeout: 30.0)
4343

44-
// Instantiate Vertex AI SDK - Default App
45-
let vertexAI = FirebaseAI.vertexAI()
46-
let _ = FirebaseAI.vertexAI(location: "my-location")
47-
48-
// Instantiate Vertex AI SDK - Custom App
49-
let _ = FirebaseAI.vertexAI(app: app!)
50-
let _ = FirebaseAI.vertexAI(app: app!, location: "my-location")
44+
// Instantiate Firebase AI SDK - Default App
45+
let firebaseAI = FirebaseAI.firebaseAI()
46+
let _ = FirebaseAI.firebaseAI(backend: .googleAI())
47+
let _ = FirebaseAI.firebaseAI(backend: .vertexAI())
48+
let _ = FirebaseAI.firebaseAI(backend: .vertexAI(location: "my-location"))
49+
50+
// Instantiate Firebase AI SDK - Custom App
51+
let _ = FirebaseAI.firebaseAI(app: app!)
52+
let _ = FirebaseAI.firebaseAI(app: app!, backend: .googleAI())
53+
let _ = FirebaseAI.firebaseAI(app: app!, backend: .vertexAI())
54+
let _ = FirebaseAI.firebaseAI(app: app!, backend: .vertexAI(location: "my-location"))
5155

5256
// Permutations without optional arguments.
5357

54-
let _ = vertexAI.generativeModel(modelName: "gemini-1.0-pro")
58+
let _ = firebaseAI.generativeModel(modelName: "gemini-2.0-flash")
5559

56-
let _ = vertexAI.generativeModel(
57-
modelName: "gemini-1.0-pro",
60+
let _ = firebaseAI.generativeModel(
61+
modelName: "gemini-2.0-flash",
5862
safetySettings: filters
5963
)
6064

61-
let _ = vertexAI.generativeModel(
62-
modelName: "gemini-1.0-pro",
65+
let _ = firebaseAI.generativeModel(
66+
modelName: "gemini-2.0-flash",
6367
generationConfig: config
6468
)
6569

66-
let _ = vertexAI.generativeModel(
67-
modelName: "gemini-1.0-pro",
70+
let _ = firebaseAI.generativeModel(
71+
modelName: "gemini-2.0-flash",
6872
systemInstruction: systemInstruction
6973
)
7074

7175
// All arguments passed.
72-
let genAI = vertexAI.generativeModel(
73-
modelName: "gemini-1.0-pro",
76+
let model = firebaseAI.generativeModel(
77+
modelName: "gemini-2.0-flash",
7478
generationConfig: config, // Optional
7579
safetySettings: filters, // Optional
7680
systemInstruction: systemInstruction, // Optional
@@ -88,35 +92,35 @@ final class APITests: XCTestCase {
8892
)]
8993

9094
do {
91-
let response = try await genAI.generateContent(contents)
95+
let response = try await model.generateContent(contents)
9296
print(response.text ?? "Couldn't get text... check status")
9397
} catch {
9498
print("Error generating content: \(error)")
9599
}
96100

97101
// Content input combinations.
98-
let _ = try await genAI.generateContent("Constant String")
102+
let _ = try await model.generateContent("Constant String")
99103
let str = "String Variable"
100-
let _ = try await genAI.generateContent(str)
101-
let _ = try await genAI.generateContent([str])
102-
let _ = try await genAI.generateContent(str, "abc", "def")
103-
let _ = try await genAI.generateContent(
104+
let _ = try await model.generateContent(str)
105+
let _ = try await model.generateContent([str])
106+
let _ = try await model.generateContent(str, "abc", "def")
107+
let _ = try await model.generateContent(
104108
str,
105109
FileDataPart(uri: "gs://test-bucket/image.jpg", mimeType: "image/jpeg")
106110
)
107111
#if canImport(UIKit)
108-
_ = try await genAI.generateContent(UIImage())
109-
_ = try await genAI.generateContent([UIImage()])
110-
_ = try await genAI.generateContent([str, UIImage(), TextPart(str)])
111-
_ = try await genAI.generateContent(str, UIImage(), "def", UIImage())
112-
_ = try await genAI.generateContent([str, UIImage(), "def", UIImage()])
113-
_ = try await genAI.generateContent([ModelContent(parts: "def", UIImage()),
112+
_ = try await model.generateContent(UIImage())
113+
_ = try await model.generateContent([UIImage()])
114+
_ = try await model.generateContent([str, UIImage(), TextPart(str)])
115+
_ = try await model.generateContent(str, UIImage(), "def", UIImage())
116+
_ = try await model.generateContent([str, UIImage(), "def", UIImage()])
117+
_ = try await model.generateContent([ModelContent(parts: "def", UIImage()),
114118
ModelContent(parts: "def", UIImage())])
115119
#elseif canImport(AppKit)
116-
_ = try await genAI.generateContent(NSImage())
117-
_ = try await genAI.generateContent([NSImage()])
118-
_ = try await genAI.generateContent(str, NSImage(), "def", NSImage())
119-
_ = try await genAI.generateContent([str, NSImage(), "def", NSImage()])
120+
_ = try await model.generateContent(NSImage())
121+
_ = try await model.generateContent([NSImage()])
122+
_ = try await model.generateContent(str, NSImage(), "def", NSImage())
123+
_ = try await model.generateContent([str, NSImage(), "def", NSImage()])
120124
#endif
121125

122126
// PartsRepresentable combinations.
@@ -147,19 +151,19 @@ final class APITests: XCTestCase {
147151
#endif
148152

149153
// countTokens API
150-
let _: CountTokensResponse = try await genAI.countTokens("What color is the Sky?")
154+
let _: CountTokensResponse = try await model.countTokens("What color is the Sky?")
151155
#if canImport(UIKit)
152-
let _: CountTokensResponse = try await genAI.countTokens("What color is the Sky?",
156+
let _: CountTokensResponse = try await model.countTokens("What color is the Sky?",
153157
UIImage())
154-
let _: CountTokensResponse = try await genAI.countTokens([
158+
let _: CountTokensResponse = try await model.countTokens([
155159
ModelContent(parts: "What color is the Sky?", UIImage()),
156160
ModelContent(parts: UIImage(), "What color is the Sky?", UIImage()),
157161
])
158162
#endif
159163

160164
// Chat
161-
_ = genAI.startChat()
162-
_ = genAI.startChat(history: [ModelContent(parts: "abc")])
165+
_ = model.startChat()
166+
_ = model.startChat(history: [ModelContent(parts: "abc")])
163167
}
164168

165169
// Public API tests for GenerateContentResponse.
@@ -179,35 +183,4 @@ final class APITests: XCTestCase {
179183
let _: String? = response.text
180184
let _: [FunctionCallPart] = response.functionCalls
181185
}
182-
183-
// Result builder alternative
184-
185-
/*
186-
let pngData = Data() // ....
187-
let contents = [GenAIContent(role: "user",
188-
parts: [
189-
.text("Is it a cat?"),
190-
.png(pngData)
191-
])]
192-
193-
// Turns into...
194-
195-
let contents = GenAIContent {
196-
Role("user") {
197-
Text("Is this a cat?")
198-
Image(png: pngData)
199-
}
200-
}
201-
202-
GenAIContent {
203-
ForEach(myInput) { input in
204-
Role(input.role) {
205-
input.contents
206-
}
207-
}
208-
}
209-
210-
// Thoughts: this looks great from a code demo, but since I assume most content will be
211-
// user generated, the result builder may not be the best API.
212-
*/
213186
}

‎FirebaseAI/Tests/Unit/Snippets/ChatSnippets.swift

+1-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import XCTest
2121

2222
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2323
final class ChatSnippets: XCTestCase {
24-
lazy var model = FirebaseAI.vertexAI().generativeModel(modelName: "gemini-1.5-flash")
24+
lazy var model = FirebaseAI.firebaseAI().generativeModel(modelName: "gemini-2.0-flash")
2525

2626
override func setUpWithError() throws {
2727
try FirebaseApp.configureDefaultAppForSnippets()

‎FirebaseAI/Tests/Unit/Snippets/FunctionCallingSnippets.swift

+2-2
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,8 @@ final class FunctionCallingSnippets: XCTestCase {
6767

6868
// Initialize the Vertex AI service and the generative model.
6969
// Use a model that supports function calling, like a Gemini 1.5 model.
70-
let model = FirebaseAI.vertexAI().generativeModel(
71-
modelName: "gemini-1.5-flash",
70+
let model = FirebaseAI.firebaseAI().generativeModel(
71+
modelName: "gemini-2.0-flash",
7272
// Provide the function declaration to the model.
7373
tools: [.functionDeclarations([fetchWeatherTool])]
7474
)

‎FirebaseAI/Tests/Unit/Snippets/MultimodalSnippets.swift

+3-1
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@ import XCTest
2626
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2727
final class MultimodalSnippets: XCTestCase {
2828
let bundle = BundleTestUtil.bundle()
29-
lazy var model = FirebaseAI.vertexAI().generativeModel(modelName: "gemini-2.0-flash")
29+
lazy var model = FirebaseAI.firebaseAI(backend: .vertexAI()).generativeModel(
30+
modelName: "gemini-2.0-flash"
31+
)
3032
lazy var videoURL = {
3133
guard let url = bundle.url(forResource: "animals", withExtension: "mp4") else {
3234
fatalError("Video file animals.mp4 not found in Resources.")

‎FirebaseAI/Tests/Unit/Snippets/StructuredOutputSnippets.swift

+4-4
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@ final class StructuredOutputSnippets: XCTestCase {
5050

5151
// Initialize the Vertex AI service and the generative model.
5252
// Use a model that supports `responseSchema`, like one of the Gemini 1.5 models.
53-
let model = FirebaseAI.vertexAI().generativeModel(
54-
modelName: "gemini-1.5-flash",
53+
let model = FirebaseAI.firebaseAI().generativeModel(
54+
modelName: "gemini-2.0-flash",
5555
// In the generation config, set the `responseMimeType` to `application/json`
5656
// and pass the JSON schema object into `responseSchema`.
5757
generationConfig: GenerationConfig(
@@ -73,8 +73,8 @@ final class StructuredOutputSnippets: XCTestCase {
7373

7474
// Initialize the Vertex AI service and the generative model.
7575
// Use a model that supports `responseSchema`, like one of the Gemini 1.5 models.
76-
let model = FirebaseAI.vertexAI().generativeModel(
77-
modelName: "gemini-1.5-flash",
76+
let model = FirebaseAI.firebaseAI().generativeModel(
77+
modelName: "gemini-2.0-flash",
7878
// In the generation config, set the `responseMimeType` to `text/x.enum`
7979
// and pass the enum schema object into `responseSchema`.
8080
generationConfig: GenerationConfig(

‎FirebaseAI/Tests/Unit/Snippets/TextSnippets.swift

+1-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import XCTest
2121

2222
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2323
final class TextSnippets: XCTestCase {
24-
lazy var model = FirebaseAI.vertexAI().generativeModel(modelName: "gemini-1.5-flash")
24+
lazy var model = FirebaseAI.firebaseAI().generativeModel(modelName: "gemini-2.0-flash")
2525

2626
override func setUpWithError() throws {
2727
try FirebaseApp.configureDefaultAppForSnippets()

0 commit comments

Comments
 (0)