Skip to content

Commit 6835193

Browse files
andrewheardandrewheard
authored
[Vertex AI] Remove extraneous GenerativeModel.modelResourceName method (#13544)
1 parent cc4f0cb commit 6835193

File tree

4 files changed

+79
-102
lines changed

4 files changed

+79
-102
lines changed

‎FirebaseVertexAI/Sources/GenerativeModel.swift

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,6 @@ import Foundation
2020
/// content based on various input types.
2121
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2222
public final class GenerativeModel {
23-
// The prefix for a model resource in the Gemini API.
24-
private static let modelResourcePrefix = "models/"
25-
2623
/// The resource name of the model in the backend; has the format "models/model-name".
2724
let modelResourceName: String
2825

@@ -73,7 +70,7 @@ public final class GenerativeModel {
7370
appCheck: AppCheckInterop?,
7471
auth: AuthInterop?,
7572
urlSession: URLSession = .shared) {
76-
modelResourceName = GenerativeModel.modelResourceName(name: name)
73+
modelResourceName = name
7774
generativeAIService = GenerativeAIService(
7875
projectID: projectID,
7976
apiKey: apiKey,
@@ -297,15 +294,6 @@ public final class GenerativeModel {
297294
}
298295
}
299296

300-
/// Returns a model resource name of the form "models/model-name" based on `name`.
301-
private static func modelResourceName(name: String) -> String {
302-
if name.contains("/") {
303-
return name
304-
} else {
305-
return modelResourcePrefix + name
306-
}
307-
}
308-
309297
/// Returns a `GenerateContentError` (for public consumption) from an internal error.
310298
///
311299
/// If `error` is already a `GenerateContentError` the error is returned unchanged.

‎FirebaseVertexAI/Sources/VertexAI.swift

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -88,22 +88,8 @@ public class VertexAI: NSObject {
8888
systemInstruction: ModelContent? = nil,
8989
requestOptions: RequestOptions = RequestOptions())
9090
-> GenerativeModel {
91-
guard let projectID = app.options.projectID else {
92-
fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
93-
}
94-
95-
let modelResourceName = modelResourceName(
96-
modelName: modelName,
97-
projectID: projectID,
98-
location: location
99-
)
100-
101-
guard let apiKey = app.options.apiKey else {
102-
fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.")
103-
}
104-
10591
return GenerativeModel(
106-
name: modelResourceName,
92+
name: modelResourceName(modelName: modelName),
10793
projectID: projectID,
10894
apiKey: apiKey,
10995
generationConfig: generationConfig,
@@ -137,16 +123,29 @@ public class VertexAI: NSObject {
137123
/// Lock to manage access to the `instances` array to avoid race conditions.
138124
private static var instancesLock: os_unfair_lock = .init()
139125

126+
let projectID: String
127+
let apiKey: String
140128
let location: String
141129

142130
init(app: FirebaseApp, location: String) {
143131
self.app = app
144-
self.location = location
145132
appCheck = ComponentType<AppCheckInterop>.instance(for: AppCheckInterop.self, in: app.container)
146133
auth = ComponentType<AuthInterop>.instance(for: AuthInterop.self, in: app.container)
134+
135+
guard let projectID = app.options.projectID else {
136+
fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
137+
}
138+
self.projectID = projectID
139+
140+
guard let apiKey = app.options.apiKey else {
141+
fatalError("The Firebase app named \"\(app.name)\" has no API key in its configuration.")
142+
}
143+
self.apiKey = apiKey
144+
145+
self.location = location
147146
}
148147

149-
private func modelResourceName(modelName: String, projectID: String, location: String) -> String {
148+
func modelResourceName(modelName: String) -> String {
150149
guard !modelName.isEmpty && modelName
151150
.allSatisfy({ !$0.isWhitespace && !$0.isNewline && $0 != "/" }) else {
152151
fatalError("""

‎FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

Lines changed: 12 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ final class GenerativeModelTests: XCTestCase {
2828
.init(category: .harassment, probability: .negligible),
2929
.init(category: .dangerousContent, probability: .negligible),
3030
].sorted()
31+
let testModelResourceName =
32+
"projects/test-project-id/locations/test-location/publishers/google/models/test-model"
3133

3234
var urlSession: URLSession!
3335
var model: GenerativeModel!
@@ -37,7 +39,7 @@ final class GenerativeModelTests: XCTestCase {
3739
configuration.protocolClasses = [MockURLProtocol.self]
3840
urlSession = try XCTUnwrap(URLSession(configuration: configuration))
3941
model = GenerativeModel(
40-
name: "my-model",
42+
name: testModelResourceName,
4143
projectID: "my-project-id",
4244
apiKey: "API_KEY",
4345
tools: nil,
@@ -303,7 +305,7 @@ final class GenerativeModelTests: XCTestCase {
303305
func testGenerateContent_appCheck_validToken() async throws {
304306
let appCheckToken = "test-valid-token"
305307
model = GenerativeModel(
306-
name: "my-model",
308+
name: testModelResourceName,
307309
projectID: "my-project-id",
308310
apiKey: "API_KEY",
309311
tools: nil,
@@ -324,7 +326,7 @@ final class GenerativeModelTests: XCTestCase {
324326

325327
func testGenerateContent_appCheck_tokenRefreshError() async throws {
326328
model = GenerativeModel(
327-
name: "my-model",
329+
name: testModelResourceName,
328330
projectID: "my-project-id",
329331
apiKey: "API_KEY",
330332
tools: nil,
@@ -346,7 +348,7 @@ final class GenerativeModelTests: XCTestCase {
346348
func testGenerateContent_auth_validAuthToken() async throws {
347349
let authToken = "test-valid-token"
348350
model = GenerativeModel(
349-
name: "my-model",
351+
name: testModelResourceName,
350352
projectID: "my-project-id",
351353
apiKey: "API_KEY",
352354
tools: nil,
@@ -367,7 +369,7 @@ final class GenerativeModelTests: XCTestCase {
367369

368370
func testGenerateContent_auth_nilAuthToken() async throws {
369371
model = GenerativeModel(
370-
name: "my-model",
372+
name: testModelResourceName,
371373
projectID: "my-project-id",
372374
apiKey: "API_KEY",
373375
tools: nil,
@@ -733,7 +735,7 @@ final class GenerativeModelTests: XCTestCase {
733735
)
734736
let requestOptions = RequestOptions(timeout: expectedTimeout)
735737
model = GenerativeModel(
736-
name: "my-model",
738+
name: testModelResourceName,
737739
projectID: "my-project-id",
738740
apiKey: "API_KEY",
739741
tools: nil,
@@ -978,7 +980,7 @@ final class GenerativeModelTests: XCTestCase {
978980
func testGenerateContentStream_appCheck_validToken() async throws {
979981
let appCheckToken = "test-valid-token"
980982
model = GenerativeModel(
981-
name: "my-model",
983+
name: testModelResourceName,
982984
projectID: "my-project-id",
983985
apiKey: "API_KEY",
984986
tools: nil,
@@ -1000,7 +1002,7 @@ final class GenerativeModelTests: XCTestCase {
10001002

10011003
func testGenerateContentStream_appCheck_tokenRefreshError() async throws {
10021004
model = GenerativeModel(
1003-
name: "my-model",
1005+
name: testModelResourceName,
10041006
projectID: "my-project-id",
10051007
apiKey: "API_KEY",
10061008
tools: nil,
@@ -1146,7 +1148,7 @@ final class GenerativeModelTests: XCTestCase {
11461148
)
11471149
let requestOptions = RequestOptions(timeout: expectedTimeout)
11481150
model = GenerativeModel(
1149-
name: "my-model",
1151+
name: testModelResourceName,
11501152
projectID: "my-project-id",
11511153
apiKey: "API_KEY",
11521154
tools: nil,
@@ -1224,7 +1226,7 @@ final class GenerativeModelTests: XCTestCase {
12241226
)
12251227
let requestOptions = RequestOptions(timeout: expectedTimeout)
12261228
model = GenerativeModel(
1227-
name: "my-model",
1229+
name: testModelResourceName,
12281230
projectID: "my-project-id",
12291231
apiKey: "API_KEY",
12301232
tools: nil,
@@ -1239,57 +1241,6 @@ final class GenerativeModelTests: XCTestCase {
12391241
XCTAssertEqual(response.totalTokens, 6)
12401242
}
12411243

1242-
// MARK: - Model Resource Name
1243-
1244-
func testModelResourceName_noPrefix() async throws {
1245-
let modelName = "my-model"
1246-
let modelResourceName = "models/\(modelName)"
1247-
1248-
model = GenerativeModel(
1249-
name: modelName,
1250-
projectID: "my-project-id",
1251-
apiKey: "API_KEY",
1252-
tools: nil,
1253-
requestOptions: RequestOptions(),
1254-
appCheck: nil,
1255-
auth: nil
1256-
)
1257-
1258-
XCTAssertEqual(model.modelResourceName, modelResourceName)
1259-
}
1260-
1261-
func testModelResourceName_modelsPrefix() async throws {
1262-
let modelResourceName = "models/my-model"
1263-
1264-
model = GenerativeModel(
1265-
name: modelResourceName,
1266-
projectID: "my-project-id",
1267-
apiKey: "API_KEY",
1268-
tools: nil,
1269-
requestOptions: RequestOptions(),
1270-
appCheck: nil,
1271-
auth: nil
1272-
)
1273-
1274-
XCTAssertEqual(model.modelResourceName, modelResourceName)
1275-
}
1276-
1277-
func testModelResourceName_tunedModelsPrefix() async throws {
1278-
let tunedModelResourceName = "tunedModels/my-model"
1279-
1280-
model = GenerativeModel(
1281-
name: tunedModelResourceName,
1282-
projectID: "my-project-id",
1283-
apiKey: "API_KEY",
1284-
tools: nil,
1285-
requestOptions: RequestOptions(),
1286-
appCheck: nil,
1287-
auth: nil
1288-
)
1289-
1290-
XCTAssertEqual(model.modelResourceName, tunedModelResourceName)
1291-
}
1292-
12931244
// MARK: - Helpers
12941245

12951246
private func nonHTTPRequestHandler() throws -> ((URLRequest) -> (

‎FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift

Lines changed: 50 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,25 +12,29 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
import Foundation
16-
1715
import FirebaseCore
18-
@testable import FirebaseVertexAI
19-
16+
import Foundation
2017
import SharedTestUtilities
21-
2218
import XCTest
2319

20+
@testable import FirebaseVertexAI
21+
2422
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2523
class VertexComponentTests: XCTestCase {
26-
static var app: FirebaseApp!
24+
static let projectID = "test-project-id"
25+
static let apiKey = "test-api-key"
26+
27+
static var app: FirebaseApp?
28+
29+
let location = "test-location"
2730

2831
override class func setUp() {
2932
super.setUp()
3033
if app == nil {
3134
let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
3235
gcmSenderID: "00000000000000000-00000000000-000000000")
33-
options.projectID = "myProjectID"
36+
options.projectID = VertexComponentTests.projectID
37+
options.apiKey = VertexComponentTests.apiKey
3438
FirebaseApp.configure(options: options)
3539
app = FirebaseApp(instanceWithName: "test", options: options)
3640
}
@@ -44,15 +48,20 @@ class VertexComponentTests: XCTestCase {
4448
/// Tests that a vertex instance can be created properly.
4549
func testVertexInstanceCreation() throws {
4650
let app = try XCTUnwrap(VertexComponentTests.app)
47-
let vertex = VertexAI.vertexAI(app: app, location: "my-location")
51+
52+
let vertex = VertexAI.vertexAI(app: app, location: location)
53+
4854
XCTAssertNotNil(vertex)
55+
XCTAssertEqual(vertex.projectID, VertexComponentTests.projectID)
56+
XCTAssertEqual(vertex.apiKey, VertexComponentTests.apiKey)
57+
XCTAssertEqual(vertex.location, location)
4958
}
5059

5160
/// Tests that a vertex instances are reused properly.
5261
func testMultipleComponentInstancesCreated() throws {
5362
let app = try XCTUnwrap(VertexComponentTests.app)
54-
let vertex1 = VertexAI.vertexAI(app: app, location: "my-location")
55-
let vertex2 = VertexAI.vertexAI(app: app, location: "my-location")
63+
let vertex1 = VertexAI.vertexAI(app: app, location: location)
64+
let vertex2 = VertexAI.vertexAI(app: app, location: location)
5665

5766
// Ensure they're the same instance.
5867
XCTAssert(vertex1 === vertex2)
@@ -68,7 +77,8 @@ class VertexComponentTests: XCTestCase {
6877
try autoreleasepool {
6978
let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
7079
gcmSenderID: "00000000000000000-00000000000-000000000")
71-
options.projectID = "myProjectID"
80+
options.projectID = VertexComponentTests.projectID
81+
options.apiKey = VertexComponentTests.apiKey
7282
let app1 = FirebaseApp(instanceWithName: "transitory app", options: options)
7383
weakApp = try XCTUnwrap(app1)
7484
let vertex = VertexAI(app: app1, location: "transitory location")
@@ -78,4 +88,33 @@ class VertexComponentTests: XCTestCase {
7888
XCTAssertNil(weakApp)
7989
XCTAssertNil(weakVertex)
8090
}
91+
92+
func testModelResourceName() throws {
93+
let app = try XCTUnwrap(VertexComponentTests.app)
94+
let vertex = VertexAI.vertexAI(app: app, location: location)
95+
let model = "test-model-name"
96+
97+
let modelResourceName = vertex.modelResourceName(modelName: model)
98+
99+
XCTAssertEqual(
100+
modelResourceName,
101+
"projects/\(vertex.projectID)/locations/\(vertex.location)/publishers/google/models/\(model)"
102+
)
103+
}
104+
105+
func testGenerativeModel() async throws {
106+
let app = try XCTUnwrap(VertexComponentTests.app)
107+
let vertex = VertexAI.vertexAI(app: app, location: location)
108+
let modelName = "test-model-name"
109+
let modelResourceName = vertex.modelResourceName(modelName: modelName)
110+
let systemInstruction = ModelContent(role: "system", parts: "test-system-instruction-prompt")
111+
112+
let generativeModel = vertex.generativeModel(
113+
modelName: modelName,
114+
systemInstruction: systemInstruction
115+
)
116+
117+
XCTAssertEqual(generativeModel.modelResourceName, modelResourceName)
118+
XCTAssertEqual(generativeModel.systemInstruction, systemInstruction)
119+
}
81120
}

0 commit comments

Comments
 (0)