Skip to content

Commit 7e20da3

Browse files
authored
[VertexAI] Add schema minItems/maxItems test (#14703)
1 parent 3eca6be commit 7e20da3

File tree

2 files changed

+33
-0
lines changed

2 files changed

+33
-0
lines changed

FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,31 @@ struct GenerateContentIntegrationTests {
165165
#endif // canImport(UIKit)
166166
}
167167

168+
@Test(arguments: InstanceConfig.allConfigsExceptDeveloperV1)
169+
func generateContentSchemaItems(_ config: InstanceConfig) async throws {
170+
let model = VertexAI.componentInstance(config).generativeModel(
171+
modelName: ModelNames.gemini2FlashLite,
172+
generationConfig: GenerationConfig(
173+
responseMIMEType: "application/json",
174+
responseSchema:
175+
.array(
176+
items: .string(description: "The name of the city"),
177+
description: "A list of city names",
178+
minItems: 3,
179+
maxItems: 5
180+
)
181+
),
182+
safetySettings: safetySettings
183+
)
184+
let prompt = "What are the biggest cities in Canada?"
185+
let response = try await model.generateContent(prompt)
186+
let text = try #require(response.text).trimmingCharacters(in: .whitespacesAndNewlines)
187+
let jsonData = try #require(text.data(using: .utf8))
188+
let decodedJSON = try JSONDecoder().decode([String].self, from: jsonData)
189+
#expect(decodedJSON.count >= 3, "Expected at least 3 cities, but got \(decodedJSON.count)")
190+
#expect(decodedJSON.count <= 5, "Expected at most 5 cities, but got \(decodedJSON.count)")
191+
}
192+
168193
// MARK: Streaming Tests
169194

170195
@Test(arguments: InstanceConfig.allConfigs)

FirebaseVertexAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,14 @@ struct InstanceConfig {
5353
developerV1BetaSpark,
5454
]
5555

56+
static let allConfigsExceptDeveloperV1 = [
57+
vertexV1,
58+
vertexV1Staging,
59+
vertexV1Beta,
60+
vertexV1BetaStaging,
61+
developerV1BetaSpark,
62+
]
63+
5664
static let vertexV1AppCheckNotConfigured = InstanceConfig(
5765
appName: FirebaseAppNames.appCheckNotConfigured,
5866
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1)

0 commit comments

Comments
 (0)