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

andrewheard · andrewheard · commit 4d2543ad587a · 2025-04-25T10:34:12.000-04:00
diff --git a/FirebaseAI/Sources/FirebaseAI.swift b/FirebaseAI/Sources/FirebaseAI.swift
@@ -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.
@@ -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.")
     }
diff --git a/FirebaseAI/Sources/Types/Public/Backend.swift b/FirebaseAI/Sources/Types/Public/Backend.swift
@@ -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
+  }
+}
diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift
@@ -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()
   }
 
diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift
@@ -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,
@@ -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?"
 
diff --git a/FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift b/FirebaseAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift
@@ -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
@@ -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
diff --git a/FirebaseAI/Tests/Unit/APITests.swift b/FirebaseAI/Tests/Unit/APITests.swift
@@ -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
@@ -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.
@@ -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.
@@ -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.
-   */
 }
diff --git a/FirebaseAI/Tests/Unit/Snippets/ChatSnippets.swift b/FirebaseAI/Tests/Unit/Snippets/ChatSnippets.swift
@@ -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()
diff --git a/FirebaseAI/Tests/Unit/Snippets/FunctionCallingSnippets.swift b/FirebaseAI/Tests/Unit/Snippets/FunctionCallingSnippets.swift
@@ -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])]
     )
diff --git a/FirebaseAI/Tests/Unit/Snippets/MultimodalSnippets.swift b/FirebaseAI/Tests/Unit/Snippets/MultimodalSnippets.swift
@@ -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.")
diff --git a/FirebaseAI/Tests/Unit/Snippets/StructuredOutputSnippets.swift b/FirebaseAI/Tests/Unit/Snippets/StructuredOutputSnippets.swift
@@ -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(
@@ -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(
diff --git a/FirebaseAI/Tests/Unit/Snippets/TextSnippets.swift b/FirebaseAI/Tests/Unit/Snippets/TextSnippets.swift
@@ -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()
diff --git a/FirebaseAI/Tests/Unit/Types/BackendTests.swift b/FirebaseAI/Tests/Unit/Types/BackendTests.swift
diff --git a/FirebaseAI/Tests/Unit/VertexComponentTests.swift b/FirebaseAI/Tests/Unit/VertexComponentTests.swift
diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift

Original file line number	Diff line number	Diff line change
`@@ -40,7 +40,7 @@ struct ImagenIntegrationTests {`
`40`	`40`
`41`	`41`	`init() async throws {`
`42`	`42`	`userID1 = try await TestHelpers.getUserID()`
`43`		`- vertex = FirebaseAI.vertexAI()`
	`43`	`+ vertex = FirebaseAI.firebaseAI(backend: .vertexAI())`
`44`	`44`	`storage = Storage.storage()`
`45`	`45`	`}`
`46`	`46`