[Vertex AI] Remove extraneous GenerativeModel.modelResourceName method (#13544)

Author: andrewheard

FirebaseVertexAI/Sources/GenerativeModel.swift

@@ -20,9 +20,6 @@ import Foundation
 /// content based on various input types.
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public final class GenerativeModel {
-  // The prefix for a model resource in the Gemini API.
-  private static let modelResourcePrefix = "models/"
-
   /// The resource name of the model in the backend; has the format "models/model-name".
   let modelResourceName: String
 
@@ -73,7 +70,7 @@ public final class GenerativeModel {
        appCheck: AppCheckInterop?,
        auth: AuthInterop?,
        urlSession: URLSession = .shared) {
-    modelResourceName = GenerativeModel.modelResourceName(name: name)
+    modelResourceName = name
     generativeAIService = GenerativeAIService(
       projectID: projectID,
       apiKey: apiKey,
@@ -297,15 +294,6 @@ public final class GenerativeModel {
     }
   }
 
-  /// Returns a model resource name of the form "models/model-name" based on `name`.
-  private static func modelResourceName(name: String) -> String {
-    if name.contains("/") {
-      return name
-    } else {
-      return modelResourcePrefix + name
-    }
-  }
-
   /// Returns a `GenerateContentError` (for public consumption) from an internal error.
   ///
   /// If `error` is already a `GenerateContentError` the error is returned unchanged.

FirebaseVertexAI/Sources/VertexAI.swift

@@ -88,22 +88,8 @@ public class VertexAI: NSObject {
                               systemInstruction: ModelContent? = nil,
                               requestOptions: RequestOptions = RequestOptions())
     -> GenerativeModel {
-    guard let projectID = app.options.projectID else {
-      fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
-    }
-
-    let modelResourceName = modelResourceName(
-      modelName: modelName,
-      projectID: projectID,
-      location: location
-    )
-
-    guard let apiKey = app.options.apiKey else {
-      fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.")
-    }
-
     return GenerativeModel(
-      name: modelResourceName,
+      name: modelResourceName(modelName: modelName),
       projectID: projectID,
       apiKey: apiKey,
       generationConfig: generationConfig,
@@ -137,16 +123,29 @@ public class VertexAI: NSObject {
   /// Lock to manage access to the `instances` array to avoid race conditions.
   private static var instancesLock: os_unfair_lock = .init()
 
+  let projectID: String
+  let apiKey: String
   let location: String
 
   init(app: FirebaseApp, location: String) {
     self.app = app
-    self.location = location
     appCheck = ComponentType<AppCheckInterop>.instance(for: AppCheckInterop.self, in: app.container)
     auth = ComponentType<AuthInterop>.instance(for: AuthInterop.self, in: app.container)
+
+    guard let projectID = app.options.projectID else {
+      fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
+    }
+    self.projectID = projectID
+
+    guard let apiKey = app.options.apiKey else {
+      fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.")
+    }
+    self.apiKey = apiKey
+
+    self.location = location
   }
 
-  private func modelResourceName(modelName: String, projectID: String, location: String) -> String {
+  func modelResourceName(modelName: String) -> String {
     guard !modelName.isEmpty && modelName
       .allSatisfy({ !$0.isWhitespace && !$0.isNewline && $0 != "/" }) else {
       fatalError("""

FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

@@ -28,6 +28,8 @@ final class GenerativeModelTests: XCTestCase {
     .init(category: .harassment, probability: .negligible),
     .init(category: .dangerousContent, probability: .negligible),
   ].sorted()
+  let testModelResourceName =
+    "projects/test-project-id/locations/test-location/publishers/google/models/test-model"
 
   var urlSession: URLSession!
   var model: GenerativeModel!
@@ -37,7 +39,7 @@ final class GenerativeModelTests: XCTestCase {
     configuration.protocolClasses = [MockURLProtocol.self]
     urlSession = try XCTUnwrap(URLSession(configuration: configuration))
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -303,7 +305,7 @@ final class GenerativeModelTests: XCTestCase {
   func testGenerateContent_appCheck_validToken() async throws {
     let appCheckToken = "test-valid-token"
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -324,7 +326,7 @@ final class GenerativeModelTests: XCTestCase {
 
   func testGenerateContent_appCheck_tokenRefreshError() async throws {
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -346,7 +348,7 @@ final class GenerativeModelTests: XCTestCase {
   func testGenerateContent_auth_validAuthToken() async throws {
     let authToken = "test-valid-token"
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -367,7 +369,7 @@ final class GenerativeModelTests: XCTestCase {
 
   func testGenerateContent_auth_nilAuthToken() async throws {
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -733,7 +735,7 @@ final class GenerativeModelTests: XCTestCase {
       )
     let requestOptions = RequestOptions(timeout: expectedTimeout)
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -978,7 +980,7 @@ final class GenerativeModelTests: XCTestCase {
   func testGenerateContentStream_appCheck_validToken() async throws {
     let appCheckToken = "test-valid-token"
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -1000,7 +1002,7 @@ final class GenerativeModelTests: XCTestCase {
 
   func testGenerateContentStream_appCheck_tokenRefreshError() async throws {
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -1146,7 +1148,7 @@ final class GenerativeModelTests: XCTestCase {
       )
     let requestOptions = RequestOptions(timeout: expectedTimeout)
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -1224,7 +1226,7 @@ final class GenerativeModelTests: XCTestCase {
       )
     let requestOptions = RequestOptions(timeout: expectedTimeout)
     model = GenerativeModel(
-      name: "my-model",
+      name: testModelResourceName,
       projectID: "my-project-id",
       apiKey: "API_KEY",
       tools: nil,
@@ -1239,57 +1241,6 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(response.totalTokens, 6)
   }
 
-  // MARK: - Model Resource Name
-
-  func testModelResourceName_noPrefix() async throws {
-    let modelName = "my-model"
-    let modelResourceName = "models/\(modelName)"
-
-    model = GenerativeModel(
-      name: modelName,
-      projectID: "my-project-id",
-      apiKey: "API_KEY",
-      tools: nil,
-      requestOptions: RequestOptions(),
-      appCheck: nil,
-      auth: nil
-    )
-
-    XCTAssertEqual(model.modelResourceName, modelResourceName)
-  }
-
-  func testModelResourceName_modelsPrefix() async throws {
-    let modelResourceName = "models/my-model"
-
-    model = GenerativeModel(
-      name: modelResourceName,
-      projectID: "my-project-id",
-      apiKey: "API_KEY",
-      tools: nil,
-      requestOptions: RequestOptions(),
-      appCheck: nil,
-      auth: nil
-    )
-
-    XCTAssertEqual(model.modelResourceName, modelResourceName)
-  }
-
-  func testModelResourceName_tunedModelsPrefix() async throws {
-    let tunedModelResourceName = "tunedModels/my-model"
-
-    model = GenerativeModel(
-      name: tunedModelResourceName,
-      projectID: "my-project-id",
-      apiKey: "API_KEY",
-      tools: nil,
-      requestOptions: RequestOptions(),
-      appCheck: nil,
-      auth: nil
-    )
-
-    XCTAssertEqual(model.modelResourceName, tunedModelResourceName)
-  }
-
   // MARK: - Helpers
 
   private func nonHTTPRequestHandler() throws -> ((URLRequest) -> (

FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift

@@ -12,25 +12,29 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import Foundation
-
 import FirebaseCore
-@testable import FirebaseVertexAI
-
+import Foundation
 import SharedTestUtilities
-
 import XCTest
 
+@testable import FirebaseVertexAI
+
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 class VertexComponentTests: XCTestCase {
-  static var app: FirebaseApp!
+  static let projectID = "test-project-id"
+  static let apiKey = "test-api-key"
+
+  static var app: FirebaseApp?
+
+  let location = "test-location"
 
   override class func setUp() {
     super.setUp()
     if app == nil {
       let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
                                     gcmSenderID: "00000000000000000-00000000000-000000000")
-      options.projectID = "myProjectID"
+      options.projectID = VertexComponentTests.projectID
+      options.apiKey = VertexComponentTests.apiKey
       FirebaseApp.configure(options: options)
       app = FirebaseApp(instanceWithName: "test", options: options)
     }
@@ -44,15 +48,20 @@ class VertexComponentTests: XCTestCase {
   /// Tests that a vertex instance can be created properly.
   func testVertexInstanceCreation() throws {
     let app = try XCTUnwrap(VertexComponentTests.app)
-    let vertex = VertexAI.vertexAI(app: app, location: "my-location")
+
+    let vertex = VertexAI.vertexAI(app: app, location: location)
+
     XCTAssertNotNil(vertex)
+    XCTAssertEqual(vertex.projectID, VertexComponentTests.projectID)
+    XCTAssertEqual(vertex.apiKey, VertexComponentTests.apiKey)
+    XCTAssertEqual(vertex.location, location)
   }
 
   /// Tests that a vertex instances are reused properly.
   func testMultipleComponentInstancesCreated() throws {
     let app = try XCTUnwrap(VertexComponentTests.app)
-    let vertex1 = VertexAI.vertexAI(app: app, location: "my-location")
-    let vertex2 = VertexAI.vertexAI(app: app, location: "my-location")
+    let vertex1 = VertexAI.vertexAI(app: app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: app, location: location)
 
     // Ensure they're the same instance.
     XCTAssert(vertex1 === vertex2)
@@ -68,7 +77,8 @@ class VertexComponentTests: XCTestCase {
     try autoreleasepool {
       let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
                                     gcmSenderID: "00000000000000000-00000000000-000000000")
-      options.projectID = "myProjectID"
+      options.projectID = VertexComponentTests.projectID
+      options.apiKey = VertexComponentTests.apiKey
       let app1 = FirebaseApp(instanceWithName: "transitory app", options: options)
       weakApp = try XCTUnwrap(app1)
       let vertex = VertexAI(app: app1, location: "transitory location")
@@ -78,4 +88,33 @@ class VertexComponentTests: XCTestCase {
     XCTAssertNil(weakApp)
     XCTAssertNil(weakVertex)
   }
+
+  func testModelResourceName() throws {
+    let app = try XCTUnwrap(VertexComponentTests.app)
+    let vertex = VertexAI.vertexAI(app: app, location: location)
+    let model = "test-model-name"
+
+    let modelResourceName = vertex.modelResourceName(modelName: model)
+
+    XCTAssertEqual(
+      modelResourceName,
+      "projects/\(vertex.projectID)/locations/\(vertex.location)/publishers/google/models/\(model)"
+    )
+  }
+
+  func testGenerativeModel() async throws {
+    let app = try XCTUnwrap(VertexComponentTests.app)
+    let vertex = VertexAI.vertexAI(app: app, location: location)
+    let modelName = "test-model-name"
+    let modelResourceName = vertex.modelResourceName(modelName: modelName)
+    let systemInstruction = ModelContent(role: "system", parts: "test-system-instruction-prompt")
+
+    let generativeModel = vertex.generativeModel(
+      modelName: modelName,
+      systemInstruction: systemInstruction
+    )
+
+    XCTAssertEqual(generativeModel.modelResourceName, modelResourceName)
+    XCTAssertEqual(generativeModel.systemInstruction, systemInstruction)
+  }
 }