diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index 648ed0d15c0..9f29599c692 100644 --- a/FirebaseAI/Sources/GenerateContentResponse.swift +++ b/FirebaseAI/Sources/GenerateContentResponse.swift @@ -253,6 +253,16 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable { case prohibitedContent = "PROHIBITED_CONTENT" case spii = "SPII" case malformedFunctionCall = "MALFORMED_FUNCTION_CALL" + case imageSafety = "IMAGE_SAFETY" + case imageProhibitedContent = "IMAGE_PROHIBITED_CONTENT" + case imageOther = "IMAGE_OTHER" + case noImage = "NO_IMAGE" + case imageRecitation = "IMAGE_RECITATION" + case language = "LANGUAGE" + case unexpectedToolCall = "UNEXPECTED_TOOL_CALL" + case tooManyToolCalls = "TOO_MANY_TOOL_CALLS" + case missingThoughtSignature = "MISSING_THOUGHT_SIGNATURE" + case malformedResponse = "MALFORMED_RESPONSE" } /// Natural stop point of the model or provided stop sequence. @@ -285,6 +295,36 @@ public struct FinishReason: DecodableProtoEnum, Hashable, Sendable { /// Token generation was stopped because the function call generated by the model was invalid. public static let malformedFunctionCall = FinishReason(kind: .malformedFunctionCall) + /// Token generation stopped because generated images contain safety violations. + public static let imageSafety = FinishReason(kind: .imageSafety) + + /// Image generation stopped because generated images have other prohibited content. + public static let imageProhibitedContent = FinishReason(kind: .imageProhibitedContent) + + /// Image generation stopped because of other miscellaneous issue. + public static let imageOther = FinishReason(kind: .imageOther) + + /// The model was expected to generate an image, but none was generated. + public static let noImage = FinishReason(kind: .noImage) + + /// Image generation stopped due to recitation. + public static let imageRecitation = FinishReason(kind: .imageRecitation) + + /// The response candidate content was flagged for using an unsupported language. + public static let language = FinishReason(kind: .language) + + /// Model generated a tool call but no tools were enabled in the request. + public static let unexpectedToolCall = FinishReason(kind: .unexpectedToolCall) + + /// Model called too many tools consecutively, thus the system exited execution. + public static let tooManyToolCalls = FinishReason(kind: .tooManyToolCalls) + + /// Request has at least one thought signature missing. + public static let missingThoughtSignature = FinishReason(kind: .missingThoughtSignature) + + /// Finished due to malformed response. + public static let malformedResponse = FinishReason(kind: .malformedResponse) + /// Returns the raw string representation of the `FinishReason` value. /// /// > Note: This value directly corresponds to the values in the [REST diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift index 455db0226fe..bfbd0512fa8 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift @@ -472,6 +472,44 @@ struct GenerateContentIntegrationTests { #endif // canImport(UIKit) } + @Test(arguments: [ + (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2_5_FlashImage), + (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2_5_FlashImage), + (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_FlashImage), + (InstanceConfig.googleAI_v1beta, ModelNames.gemini3_1_FlashImagePreview), + (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini3_1_FlashImagePreview), + ]) + func generateContent_finishReason_imageSafety(_ config: InstanceConfig, + modelName: String) async throws { + let generationConfig = GenerationConfig( + responseModalities: [.image] + ) + let model = FirebaseAI.componentInstance(config).generativeModel( + modelName: modelName, + generationConfig: generationConfig, + ) + let prompt = "A graphic image of violence" // This prompt should trigger safety violation + + do { + let response = try await model.generateContent(prompt) + + // vertexAI gemini3_1_FlashImagePreview doesn't throw. + let candidate = try #require(response.candidates.first) + #expect(candidate.finishReason == .stop) + } catch { + guard let error = error as? GenerateContentError else { + Issue.record("Expected a \(GenerateContentError.self); got \(error.self).") + throw error + } + guard case let .responseStoppedEarly(reason, response) = error else { + Issue.record("Expected a GenerateContentError.responseStoppedEarly; got \(error.self).") + throw error + } + #expect(reason == .imageSafety || reason == .noImage) + #expect(response.candidates.first?.content.parts.isEmpty == true) // Ensure no content + } + } + @Test(arguments: [ (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2_5_FlashImage), (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2_5_FlashImage), diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift index 4d9fce956e2..623b757ce57 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift @@ -168,6 +168,27 @@ struct ImagenIntegrationTests { } } + @Test func generateImage_finishReason_imageSafety() async throws { + let generationConfig = ImagenGenerationConfig( + numberOfImages: 1, + imageFormat: .jpeg() + ) + let model = vertex.imagenModel( + modelName: "imagen-3.0-fast-generate-001", + generationConfig: generationConfig, + safetySettings: ImagenSafetySettings( + safetyFilterLevel: .blockLowAndAbove, + personFilterLevel: .blockAll + ) + ) + let imagePrompt = "A graphic image of violence" // This prompt should trigger safety violation + + let response = try await model.generateImages(prompt: imagePrompt) + + #expect(response.filteredReason == "SAFETY") // Assuming filteredReason is set + #expect(response.images.isEmpty) // No images should be generated + } + // TODO(#14221): Add an integration test for the prompt being blocked. // TODO(#14452): Add integration tests for validating that Storage Rules are enforced. diff --git a/FirebaseAI/Tests/Unit/APITests.swift b/FirebaseAI/Tests/Unit/APITests.swift index d6e870f42b7..96f40d0d000 100644 --- a/FirebaseAI/Tests/Unit/APITests.swift +++ b/FirebaseAI/Tests/Unit/APITests.swift @@ -233,4 +233,88 @@ final class APITests: XCTestCase { XCTAssertEqual(cacheDetail?.modality, .text) XCTAssertEqual(cacheDetail?.tokenCount, 50) } + + func testFinishReason_decoding() throws { + let decoder = JSONDecoder() + + // Test LANGUAGE + var json = createResponseJSON(finishReason: "LANGUAGE") + var response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .language) + + // Test UNEXPECTED_TOOL_CALL + json = createResponseJSON(finishReason: "UNEXPECTED_TOOL_CALL") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .unexpectedToolCall) + + // Test TOO_MANY_TOOL_CALLS + json = createResponseJSON(finishReason: "TOO_MANY_TOOL_CALLS") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .tooManyToolCalls) + + // Test MISSING_THOUGHT_SIGNATURE + json = createResponseJSON(finishReason: "MISSING_THOUGHT_SIGNATURE") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .missingThoughtSignature) + + // Test MALFORMED_RESPONSE + json = createResponseJSON(finishReason: "MALFORMED_RESPONSE") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .malformedResponse) + + // Test IMAGE_SAFETY + json = createResponseJSON(finishReason: "IMAGE_SAFETY") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .imageSafety) + + // Test IMAGE_PROHIBITED_CONTENT + json = createResponseJSON(finishReason: "IMAGE_PROHIBITED_CONTENT") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .imageProhibitedContent) + + // Test IMAGE_OTHER + json = createResponseJSON(finishReason: "IMAGE_OTHER") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .imageOther) + + // Test NO_IMAGE + json = createResponseJSON(finishReason: "NO_IMAGE") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .noImage) + + // Test IMAGE_RECITATION + json = createResponseJSON(finishReason: "IMAGE_RECITATION") + response = try decoder.decode(GenerateContentResponse.self, from: json) + XCTAssertEqual(response.candidates.first?.finishReason, .imageRecitation) + } + + // Helper function to create JSON for GenerateContentResponse with a specific finish reason + private func createResponseJSON(finishReason: String) -> Data { + return """ + { + "candidates": [ + { + "content": { + "parts": [ + { "text": "Test reply" } + ], + "role": "model" + }, + "finishReason": "\(finishReason)", + "index": 0, + "safetyRatings": [] + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "cachedContentTokenCount": 0, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + "promptTokensDetails": [], + "cacheTokensDetails": [], + "candidatesTokensDetails": [] + } + } + """.data(using: .utf8)! + } }