[Vertex AI] Add VertexAI.Preview nested type for preview APIs

firebase · Jan 10, 2025 · 4ff1cdf · 4ff1cdf
1 parent a3a026a
commit 4ff1cdf
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 16 deletions.
diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift
@@ -104,21 +104,6 @@ public class VertexAI {
     )
   }
 
-  public func imagenModel(modelName: String, modelConfig: ImagenModelConfig? = nil,
-                          safetySettings: ImagenSafetySettings? = nil,
-                          requestOptions: RequestOptions = RequestOptions()) -> ImagenModel {
-    return ImagenModel(
-      name: modelResourceName(modelName: modelName),
-      projectID: projectID,
-      apiKey: apiKey,
-      modelConfig: modelConfig,
-      safetySettings: safetySettings,
-      requestOptions: requestOptions,
-      appCheck: appCheck,
-      auth: auth
-    )
-  }
-
   /// Class to enable VertexAI to register via the Objective-C based Firebase component system
   /// to include VertexAI in the userAgent.
   @objc(FIRVertexAIComponent) class FirebaseVertexAIComponent: NSObject {}
@@ -182,3 +167,38 @@ public class VertexAI {
     return "projects/\(projectID)/locations/\(location)/publishers/google/models/\(modelName)"
   }
 }
+
+// MARK: - Preview APIs
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public extension VertexAI {
+  struct Preview {
+    let vertexAI: VertexAI
+
+    init(vertexAI: VertexAI) {
+      self.vertexAI = vertexAI
+    }
+  }
+
+  func preview() -> Preview {
+    return Preview(vertexAI: self)
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public extension VertexAI.Preview {
+  func imagenModel(modelName: String, modelConfig: ImagenModelConfig? = nil,
+                   safetySettings: ImagenSafetySettings? = nil,
+                   requestOptions: RequestOptions = RequestOptions()) -> ImagenModel {
+    return ImagenModel(
+      name: vertexAI.modelResourceName(modelName: modelName),
+      projectID: vertexAI.projectID,
+      apiKey: vertexAI.apiKey,
+      modelConfig: modelConfig,
+      safetySettings: safetySettings,
+      requestOptions: requestOptions,
+      appCheck: vertexAI.appCheck,
+      auth: vertexAI.auth
+    )
+  }
+}
diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift
@@ -61,7 +61,7 @@ final class IntegrationTests: XCTestCase {
       toolConfig: .init(functionCallingConfig: .none()),
       systemInstruction: systemInstruction
     )
-    imagenModel = vertex.imagenModel(
+    imagenModel = vertex.preview().imagenModel(
       modelName: "imagen-3.0-fast-generate-001",
       modelConfig: ImagenModelConfig(imageFormat: .jpeg(compressionQuality: 70)),
       safetySettings: ImagenSafetySettings(