diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift index 95835a4ea7d..f1650f26471 100644 --- a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift +++ b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift @@ -165,31 +165,6 @@ struct GenerateContentIntegrationTests { #endif // canImport(UIKit) } - @Test(arguments: InstanceConfig.allConfigsExceptDeveloperV1) - func generateContentSchemaItems(_ config: InstanceConfig) async throws { - let model = VertexAI.componentInstance(config).generativeModel( - modelName: ModelNames.gemini2FlashLite, - generationConfig: GenerationConfig( - responseMIMEType: "application/json", - responseSchema: - .array( - items: .string(description: "The name of the city"), - description: "A list of city names", - minItems: 3, - maxItems: 5 - ) - ), - safetySettings: safetySettings - ) - let prompt = "What are the biggest cities in Canada?" - let response = try await model.generateContent(prompt) - let text = try #require(response.text).trimmingCharacters(in: .whitespacesAndNewlines) - let jsonData = try #require(text.data(using: .utf8)) - let decodedJSON = try JSONDecoder().decode([String].self, from: jsonData) - #expect(decodedJSON.count >= 3, "Expected at least 3 cities, but got \(decodedJSON.count)") - #expect(decodedJSON.count <= 5, "Expected at most 5 cities, but got \(decodedJSON.count)") - } - // MARK: Streaming Tests @Test(arguments: InstanceConfig.allConfigs) diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/SchemaTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/SchemaTests.swift new file mode 100644 index 00000000000..35eb6442d12 --- /dev/null +++ b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/SchemaTests.swift @@ -0,0 +1,75 @@ +// 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. + +import FirebaseAuth +import FirebaseCore +import FirebaseStorage +import FirebaseVertexAI +import Testing +import VertexAITestApp + +#if canImport(UIKit) + import UIKit +#endif // canImport(UIKit) + +@testable import struct FirebaseVertexAI.BackendError + +@Suite(.serialized) +/// Test the schema fields. +struct SchemaTests { + // Set temperature, topP and topK to lowest allowed values to make responses more deterministic. + let generationConfig = GenerationConfig(temperature: 0.0, topP: 0.0, topK: 1) + let safetySettings = [ + SafetySetting(harmCategory: .harassment, threshold: .blockLowAndAbove), + SafetySetting(harmCategory: .hateSpeech, threshold: .blockLowAndAbove), + SafetySetting(harmCategory: .sexuallyExplicit, threshold: .blockLowAndAbove), + SafetySetting(harmCategory: .dangerousContent, threshold: .blockLowAndAbove), + SafetySetting(harmCategory: .civicIntegrity, threshold: .blockLowAndAbove), + ] + // Candidates and total token counts may differ slightly between runs due to whitespace tokens. + let tokenCountAccuracy = 1 + + let storage: Storage + let userID1: String + + init() async throws { + userID1 = try await TestHelpers.getUserID() + storage = Storage.storage() + } + + @Test(arguments: InstanceConfig.allConfigsExceptDeveloperV1) + func generateContentSchemaItems(_ config: InstanceConfig) async throws { + let model = VertexAI.componentInstance(config).generativeModel( + modelName: ModelNames.gemini2FlashLite, + generationConfig: GenerationConfig( + responseMIMEType: "application/json", + responseSchema: + .array( + items: .string(description: "The name of the city"), + description: "A list of city names", + minItems: 3, + maxItems: 5 + ) + ), + safetySettings: safetySettings + ) + let prompt = "What are the biggest cities in Canada?" + let response = try await model.generateContent(prompt) + let text = try #require(response.text).trimmingCharacters(in: .whitespacesAndNewlines) + let jsonData = try #require(text.data(using: .utf8)) + let decodedJSON = try JSONDecoder().decode([String].self, from: jsonData) + #expect(decodedJSON.count >= 3, "Expected at least 3 cities, but got \(decodedJSON.count)") + #expect(decodedJSON.count <= 5, "Expected at most 5 cities, but got \(decodedJSON.count)") + } +} diff --git a/FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj b/FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj index a7d2bf023b2..350f121c394 100644 --- a/FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj +++ b/FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj @@ -28,6 +28,7 @@ 86D77E022D7B63AF003D155D /* Constants.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77E012D7B63AC003D155D /* Constants.swift */; }; 86D77E042D7B6C9D003D155D /* InstanceConfig.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77E032D7B6C95003D155D /* InstanceConfig.swift */; }; DEF0BB4F2DA74F680093E9F4 /* TestHelpers.swift in Sources */ = {isa = PBXBuildFile; fileRef = DEF0BB4E2DA74F460093E9F4 /* TestHelpers.swift */; }; + DEF0BB512DA9B7450093E9F4 /* SchemaTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = DEF0BB502DA9B7400093E9F4 /* SchemaTests.swift */; }; /* End PBXBuildFile section */ /* Begin PBXContainerItemProxy section */ @@ -61,6 +62,7 @@ 86D77E012D7B63AC003D155D /* Constants.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Constants.swift; sourceTree = ""; }; 86D77E032D7B6C95003D155D /* InstanceConfig.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = InstanceConfig.swift; sourceTree = ""; }; DEF0BB4E2DA74F460093E9F4 /* TestHelpers.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TestHelpers.swift; sourceTree = ""; }; + DEF0BB502DA9B7400093E9F4 /* SchemaTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SchemaTests.swift; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFrameworksBuildPhase section */ @@ -139,6 +141,7 @@ 868A7C572CCC27AF00E449DD /* Integration */ = { isa = PBXGroup; children = ( + DEF0BB502DA9B7400093E9F4 /* SchemaTests.swift */, DEF0BB4E2DA74F460093E9F4 /* TestHelpers.swift */, 8689CDCB2D7F8BCF00BF426B /* CountTokensIntegrationTests.swift */, 868A7C4D2CCC1F4700E449DD /* Credentials.swift */, @@ -292,6 +295,7 @@ files = ( 8689CDCC2D7F8BD700BF426B /* CountTokensIntegrationTests.swift in Sources */, 86D77E042D7B6C9D003D155D /* InstanceConfig.swift in Sources */, + DEF0BB512DA9B7450093E9F4 /* SchemaTests.swift in Sources */, DEF0BB4F2DA74F680093E9F4 /* TestHelpers.swift in Sources */, 868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */, 864F8F712D4980DD0002EA7E /* ImagenIntegrationTests.swift in Sources */,