Skip to content

[Firebase AI] Implement new public API surface and tests #14765

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

Merged
merged 4 commits into from
Apr 24, 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
28 changes: 12 additions & 16 deletions FirebaseAI/Sources/FirebaseAI.swift
Original file line number Diff line number Diff line change
Expand Up @@ -25,26 +25,22 @@ internal import FirebaseCoreExtension
public final class FirebaseAI: Sendable {
// MARK: - Public APIs

/// Creates an instance of `VertexAI`.
/// Creates an instance of `FirebaseAI`.
///
/// - Parameters:
/// - app: A custom `FirebaseApp` used for initialization; if not specified, uses the default
/// ``FirebaseApp``.
/// - location: The region identifier, defaulting to `us-central1`; see
/// [Vertex AI locations]
/// (https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations)
/// for a list of supported locations.
/// - Returns: A `VertexAI` instance, configured with the custom `FirebaseApp`.
public static func vertexAI(app: FirebaseApp? = nil,
location: String = "us-central1") -> FirebaseAI {
let vertexInstance = vertexAI(app: app, location: location, apiConfig: defaultVertexAIAPIConfig)
// Verify that the `VertexAI` instance is always configured with the production endpoint since
/// - backend: The backend API for the Firebase AI SDK; if not specified, uses the default
/// ``Backend/googleAI()`` (Gemini Developer API).
/// - Returns: A `FirebaseAI` instance, configured with the custom `FirebaseApp`.
public static func firebaseAI(app: FirebaseApp? = nil,
backend: Backend = .googleAI()) -> FirebaseAI {
let instance = firebaseAI(app: app, location: backend.location, apiConfig: backend.apiConfig)
// Verify that the `FirebaseAI` instance is always configured with the production endpoint since
// this is the public API surface for creating an instance.
assert(vertexInstance.apiConfig.service == .vertexAI(endpoint: .firebaseVertexAIProd))
assert(vertexInstance.apiConfig.service.endpoint == .firebaseVertexAIProd)
assert(vertexInstance.apiConfig.version == .v1beta)

return vertexInstance
assert(instance.apiConfig.service.endpoint == .firebaseVertexAIProd)
assert(instance.apiConfig.version == .v1beta)
return instance
}

/// Initializes a generative model with the given parameters.
Expand Down Expand Up @@ -163,7 +159,7 @@ public final class FirebaseAI: Sendable {
version: .v1beta
)

static func vertexAI(app: FirebaseApp?, location: String?, apiConfig: APIConfig) -> FirebaseAI {
static func firebaseAI(app: FirebaseApp?, location: String?, apiConfig: APIConfig) -> FirebaseAI {
guard let app = app ?? FirebaseApp.app() else {
fatalError("No instance of the default Firebase app was found.")
}
Expand Down
50 changes: 50 additions & 0 deletions FirebaseAI/Sources/Types/Public/Backend.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
// 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.

/// Represents available backend APIs for the Firebase AI SDK.
public struct Backend {
// MARK: - Public API

/// Initializes a `Backend` configured for the Gemini API in Vertex AI.
///
/// - Parameters:
/// - location: The region identifier, defaulting to `us-central1`; see
/// [Vertex AI locations]
/// (https://firebase.google.com/docs/vertex-ai/locations?platform=ios#available-locations)
/// for a list of supported locations.
public static func vertexAI(location: String = "us-central1") -> Backend {
return Backend(
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1beta),
location: location
)
}

/// Initializes a `Backend` configured for the Google Developer API.
public static func googleAI() -> Backend {
return Backend(
apiConfig: APIConfig(service: .developer(endpoint: .firebaseVertexAIProd), version: .v1beta),
location: nil
)
}

// MARK: - Internal

let apiConfig: APIConfig
let location: String?

init(apiConfig: APIConfig, location: String?) {
self.apiConfig = apiConfig
self.location = location
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ struct ImagenIntegrationTests {

init() async throws {
userID1 = try await TestHelpers.getUserID()
vertex = FirebaseAI.vertexAI()
vertex = FirebaseAI.firebaseAI(backend: .vertexAI())
storage = Storage.storage()
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ final class IntegrationTests: XCTestCase {

override func setUp() async throws {
userID1 = try await TestHelpers.getUserID()
vertex = FirebaseAI.vertexAI()
vertex = FirebaseAI.firebaseAI(backend: .vertexAI())
model = vertex.generativeModel(
modelName: "gemini-2.0-flash",
generationConfig: generationConfig,
Expand Down Expand Up @@ -200,7 +200,7 @@ final class IntegrationTests: XCTestCase {

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

Expand Down
4 changes: 2 additions & 2 deletions FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ extension FirebaseAI {
switch instanceConfig.apiConfig.service {
case .vertexAI:
let location = instanceConfig.location ?? "us-central1"
return FirebaseAI.vertexAI(
return FirebaseAI.firebaseAI(
app: instanceConfig.app,
location: location,
apiConfig: instanceConfig.apiConfig
Expand All @@ -137,7 +137,7 @@ extension FirebaseAI {
instanceConfig.location == nil,
"The Developer API is global and does not support `location`."
)
return FirebaseAI.vertexAI(
return FirebaseAI.firebaseAI(
app: instanceConfig.app,
location: nil,
apiConfig: instanceConfig.apiConfig
Expand Down
109 changes: 41 additions & 68 deletions FirebaseAI/Tests/Unit/APITests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -41,36 +41,40 @@ final class APITests: XCTestCase {
let requestOptions = RequestOptions()
let _ = RequestOptions(timeout: 30.0)

// Instantiate Vertex AI SDK - Default App
let vertexAI = FirebaseAI.vertexAI()
let _ = FirebaseAI.vertexAI(location: "my-location")

// Instantiate Vertex AI SDK - Custom App
let _ = FirebaseAI.vertexAI(app: app!)
let _ = FirebaseAI.vertexAI(app: app!, location: "my-location")
// Instantiate Firebase AI SDK - Default App
let firebaseAI = FirebaseAI.firebaseAI()
let _ = FirebaseAI.firebaseAI(backend: .googleAI())
let _ = FirebaseAI.firebaseAI(backend: .vertexAI())
let _ = FirebaseAI.firebaseAI(backend: .vertexAI(location: "my-location"))

// Instantiate Firebase AI SDK - Custom App
let _ = FirebaseAI.firebaseAI(app: app!)
let _ = FirebaseAI.firebaseAI(app: app!, backend: .googleAI())
let _ = FirebaseAI.firebaseAI(app: app!, backend: .vertexAI())
let _ = FirebaseAI.firebaseAI(app: app!, backend: .vertexAI(location: "my-location"))

// Permutations without optional arguments.

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

let _ = vertexAI.generativeModel(
modelName: "gemini-1.0-pro",
let _ = firebaseAI.generativeModel(
modelName: "gemini-2.0-flash",
safetySettings: filters
)

let _ = vertexAI.generativeModel(
modelName: "gemini-1.0-pro",
let _ = firebaseAI.generativeModel(
modelName: "gemini-2.0-flash",
generationConfig: config
)

let _ = vertexAI.generativeModel(
modelName: "gemini-1.0-pro",
let _ = firebaseAI.generativeModel(
modelName: "gemini-2.0-flash",
systemInstruction: systemInstruction
)

// All arguments passed.
let genAI = vertexAI.generativeModel(
modelName: "gemini-1.0-pro",
let model = firebaseAI.generativeModel(
modelName: "gemini-2.0-flash",
generationConfig: config, // Optional
safetySettings: filters, // Optional
systemInstruction: systemInstruction, // Optional
Expand All @@ -88,35 +92,35 @@ final class APITests: XCTestCase {
)]

do {
let response = try await genAI.generateContent(contents)
let response = try await model.generateContent(contents)
print(response.text ?? "Couldn't get text... check status")
} catch {
print("Error generating content: \(error)")
}

// Content input combinations.
let _ = try await genAI.generateContent("Constant String")
let _ = try await model.generateContent("Constant String")
let str = "String Variable"
let _ = try await genAI.generateContent(str)
let _ = try await genAI.generateContent([str])
let _ = try await genAI.generateContent(str, "abc", "def")
let _ = try await genAI.generateContent(
let _ = try await model.generateContent(str)
let _ = try await model.generateContent([str])
let _ = try await model.generateContent(str, "abc", "def")
let _ = try await model.generateContent(
str,
FileDataPart(uri: "gs://test-bucket/image.jpg", mimeType: "image/jpeg")
)
#if canImport(UIKit)
_ = try await genAI.generateContent(UIImage())
_ = try await genAI.generateContent([UIImage()])
_ = try await genAI.generateContent([str, UIImage(), TextPart(str)])
_ = try await genAI.generateContent(str, UIImage(), "def", UIImage())
_ = try await genAI.generateContent([str, UIImage(), "def", UIImage()])
_ = try await genAI.generateContent([ModelContent(parts: "def", UIImage()),
_ = try await model.generateContent(UIImage())
_ = try await model.generateContent([UIImage()])
_ = try await model.generateContent([str, UIImage(), TextPart(str)])
_ = try await model.generateContent(str, UIImage(), "def", UIImage())
_ = try await model.generateContent([str, UIImage(), "def", UIImage()])
_ = try await model.generateContent([ModelContent(parts: "def", UIImage()),
ModelContent(parts: "def", UIImage())])
#elseif canImport(AppKit)
_ = try await genAI.generateContent(NSImage())
_ = try await genAI.generateContent([NSImage()])
_ = try await genAI.generateContent(str, NSImage(), "def", NSImage())
_ = try await genAI.generateContent([str, NSImage(), "def", NSImage()])
_ = try await model.generateContent(NSImage())
_ = try await model.generateContent([NSImage()])
_ = try await model.generateContent(str, NSImage(), "def", NSImage())
_ = try await model.generateContent([str, NSImage(), "def", NSImage()])
#endif

// PartsRepresentable combinations.
Expand Down Expand Up @@ -147,19 +151,19 @@ final class APITests: XCTestCase {
#endif

// countTokens API
let _: CountTokensResponse = try await genAI.countTokens("What color is the Sky?")
let _: CountTokensResponse = try await model.countTokens("What color is the Sky?")
#if canImport(UIKit)
let _: CountTokensResponse = try await genAI.countTokens("What color is the Sky?",
let _: CountTokensResponse = try await model.countTokens("What color is the Sky?",
UIImage())
let _: CountTokensResponse = try await genAI.countTokens([
let _: CountTokensResponse = try await model.countTokens([
ModelContent(parts: "What color is the Sky?", UIImage()),
ModelContent(parts: UIImage(), "What color is the Sky?", UIImage()),
])
#endif

// Chat
_ = genAI.startChat()
_ = genAI.startChat(history: [ModelContent(parts: "abc")])
_ = model.startChat()
_ = model.startChat(history: [ModelContent(parts: "abc")])
}

// Public API tests for GenerateContentResponse.
Expand All @@ -179,35 +183,4 @@ final class APITests: XCTestCase {
let _: String? = response.text
let _: [FunctionCallPart] = response.functionCalls
}

// Result builder alternative

/*
let pngData = Data() // ....
let contents = [GenAIContent(role: "user",
parts: [
.text("Is it a cat?"),
.png(pngData)
])]

// Turns into...

let contents = GenAIContent {
Role("user") {
Text("Is this a cat?")
Image(png: pngData)
}
}

GenAIContent {
ForEach(myInput) { input in
Role(input.role) {
input.contents
}
}
}

// Thoughts: this looks great from a code demo, but since I assume most content will be
// user generated, the result builder may not be the best API.
*/
}
2 changes: 1 addition & 1 deletion FirebaseAI/Tests/Unit/Snippets/ChatSnippets.swift
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import XCTest

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

override func setUpWithError() throws {
try FirebaseApp.configureDefaultAppForSnippets()
Expand Down
4 changes: 2 additions & 2 deletions FirebaseAI/Tests/Unit/Snippets/FunctionCallingSnippets.swift
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,8 @@ final class FunctionCallingSnippets: XCTestCase {

// Initialize the Vertex AI service and the generative model.
// Use a model that supports function calling, like a Gemini 1.5 model.
let model = FirebaseAI.vertexAI().generativeModel(
modelName: "gemini-1.5-flash",
let model = FirebaseAI.firebaseAI().generativeModel(
modelName: "gemini-2.0-flash",
// Provide the function declaration to the model.
tools: [.functionDeclarations([fetchWeatherTool])]
)
Expand Down
4 changes: 3 additions & 1 deletion FirebaseAI/Tests/Unit/Snippets/MultimodalSnippets.swift
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,9 @@ import XCTest
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
final class MultimodalSnippets: XCTestCase {
let bundle = BundleTestUtil.bundle()
lazy var model = FirebaseAI.vertexAI().generativeModel(modelName: "gemini-2.0-flash")
lazy var model = FirebaseAI.firebaseAI(backend: .vertexAI()).generativeModel(
modelName: "gemini-2.0-flash"
)
lazy var videoURL = {
guard let url = bundle.url(forResource: "animals", withExtension: "mp4") else {
fatalError("Video file animals.mp4 not found in Resources.")
Expand Down
8 changes: 4 additions & 4 deletions FirebaseAI/Tests/Unit/Snippets/StructuredOutputSnippets.swift
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@ final class StructuredOutputSnippets: XCTestCase {

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

// Initialize the Vertex AI service and the generative model.
// Use a model that supports `responseSchema`, like one of the Gemini 1.5 models.
let model = FirebaseAI.vertexAI().generativeModel(
modelName: "gemini-1.5-flash",
let model = FirebaseAI.firebaseAI().generativeModel(
modelName: "gemini-2.0-flash",
// In the generation config, set the `responseMimeType` to `text/x.enum`
// and pass the enum schema object into `responseSchema`.
generationConfig: GenerationConfig(
Expand Down
2 changes: 1 addition & 1 deletion FirebaseAI/Tests/Unit/Snippets/TextSnippets.swift
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import XCTest

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

override func setUpWithError() throws {
try FirebaseApp.configureDefaultAppForSnippets()
Expand Down
Loading
Loading