diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index 8f8026ef376..79f6a2b0193 100644 --- a/FirebaseAI/Sources/GenerateContentResponse.swift +++ b/FirebaseAI/Sources/GenerateContentResponse.swift @@ -145,7 +145,7 @@ public struct CitationMetadata: Sendable { /// A struct describing a source attribution. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -public struct Citation: Sendable { +public struct Citation: Sendable, Equatable { /// The inclusive beginning of a sequence in a model response that derives from a cited source. public let startIndex: Int @@ -165,6 +165,20 @@ public struct Citation: Sendable { /// /// > Tip: `DateComponents` can be converted to a `Date` using the `date` computed property. public let publicationDate: DateComponents? + + init(startIndex: Int, + endIndex: Int, + uri: String? = nil, + title: String? = nil, + license: String? = nil, + publicationDate: DateComponents? = nil) { + self.startIndex = startIndex + self.endIndex = endIndex + self.uri = uri + self.title = title + self.license = license + self.publicationDate = publicationDate + } } /// A value enumerating possible reasons for a model to terminate a content generation request. @@ -385,7 +399,23 @@ extension Candidate: Decodable { } @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) -extension CitationMetadata: Decodable {} +extension CitationMetadata: Decodable { + enum CodingKeys: CodingKey { + case citations // Vertex AI + case citationSources // Google AI + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + + // Decode for Google API if `citationSources` key is present. + if container.contains(.citationSources) { + citations = try container.decode([Citation].self, forKey: .citationSources) + } else { // Fallback to default Vertex AI decoding. + citations = try container.decode([Citation].self, forKey: .citations) + } + } +} @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) extension Citation: Decodable { diff --git a/FirebaseAI/Tests/Unit/Types/CitationMetadataTests.swift b/FirebaseAI/Tests/Unit/Types/CitationMetadataTests.swift new file mode 100644 index 00000000000..d75325f1a88 --- /dev/null +++ b/FirebaseAI/Tests/Unit/Types/CitationMetadataTests.swift @@ -0,0 +1,78 @@ +// 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 XCTest + +@testable import FirebaseAI + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +final class CitationMetadataTests: XCTestCase { + let decoder = JSONDecoder() + + let expectedStartIndex = 100 + let expectedEndIndex = 200 + let expectedURI = "https://example.com/citation-1" + lazy var citationJSON = """ + { + "startIndex" : \(expectedStartIndex), + "endIndex" : \(expectedEndIndex), + "uri" : "\(expectedURI)" + } + """ + lazy var expectedCitation = Citation( + startIndex: expectedStartIndex, + endIndex: expectedEndIndex, + uri: expectedURI + ) + + // MARK: - Google AI Format Decoding + + func testDecodeCitationMetadata_googleAIFormat() throws { + let json = """ + { + "citationSources": [\(citationJSON)] + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let citationMetadata = try decoder.decode( + CitationMetadata.self, + from: jsonData + ) + + XCTAssertEqual(citationMetadata.citations.count, 1) + let citation = try XCTUnwrap(citationMetadata.citations.first) + XCTAssertEqual(citation, expectedCitation) + } + + // MARK: - Vertex AI Format Decoding + + func testDecodeCitationMetadata_vertexAIFormat() throws { + let json = """ + { + "citations": [\(citationJSON)] + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let citationMetadata = try decoder.decode( + CitationMetadata.self, + from: jsonData + ) + + XCTAssertEqual(citationMetadata.citations.count, 1) + let citation = try XCTUnwrap(citationMetadata.citations.first) + XCTAssertEqual(citation, expectedCitation) + } +} diff --git a/FirebaseAI/Tests/Unit/Types/CitationTests.swift b/FirebaseAI/Tests/Unit/Types/CitationTests.swift new file mode 100644 index 00000000000..ced45526721 --- /dev/null +++ b/FirebaseAI/Tests/Unit/Types/CitationTests.swift @@ -0,0 +1,113 @@ +// 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 FirebaseAI +import XCTest + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +final class CitationTests: XCTestCase { + let decoder = JSONDecoder() + + // MARK: - Decoding Tests + + func testDecodeCitation_minimalParameters() throws { + let expectedEndIndex = 150 + let json = """ + { + "endIndex" : \(expectedEndIndex) + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let citation = try decoder.decode(Citation.self, from: jsonData) + + XCTAssertEqual(citation.startIndex, 0, "Omitted startIndex should be decoded as 0.") + XCTAssertEqual(citation.endIndex, expectedEndIndex) + XCTAssertNil(citation.uri) + XCTAssertNil(citation.title) + XCTAssertNil(citation.license) + XCTAssertNil(citation.publicationDate) + } + + func testDecodeCitation_allParameters() throws { + let expectedStartIndex = 100 + let expectedEndIndex = 200 + let expectedURI = "https://example.com/citation-1" + let expectedTitle = "Example Citation Title" + let expectedLicense = "mit" + let expectedYear = 2023 + let expectedMonth = 10 + let expectedDay = 26 + let json = """ + { + "startIndex" : \(expectedStartIndex), + "endIndex" : \(expectedEndIndex), + "uri" : "\(expectedURI)", + "title" : "\(expectedTitle)", + "license" : "\(expectedLicense)", + "publicationDate" : { + "year" : \(expectedYear), + "month" : \(expectedMonth), + "day" : \(expectedDay) + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let citation = try decoder.decode(Citation.self, from: jsonData) + + XCTAssertEqual(citation.startIndex, expectedStartIndex) + XCTAssertEqual(citation.endIndex, expectedEndIndex) + XCTAssertEqual(citation.uri, expectedURI) + XCTAssertEqual(citation.title, expectedTitle) + XCTAssertEqual(citation.license, expectedLicense) + let publicationDate = try XCTUnwrap(citation.publicationDate) + XCTAssertEqual(publicationDate.year, expectedYear) + XCTAssertEqual(publicationDate.month, expectedMonth) + XCTAssertEqual(publicationDate.day, expectedDay) + } + + func testDecodeCitation_emptyStringsForOptionals_setsToNil() throws { + let expectedEndIndex = 300 + let json = """ + { + "endIndex" : \(expectedEndIndex), + "uri" : "", + "title" : "", + "license" : "" + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let citation = try decoder.decode(Citation.self, from: jsonData) + + XCTAssertEqual(citation.startIndex, 0, "Omitted startIndex should be decoded as 0.") + XCTAssertEqual(citation.endIndex, expectedEndIndex) + XCTAssertNil(citation.uri, "Empty URI string should be decoded as nil.") + XCTAssertNil(citation.title, "Empty title string should be decoded as nil.") + XCTAssertNil(citation.license, "Empty license string should be decoded as nil.") + XCTAssertNil(citation.publicationDate) + } + + func testDecodeCitation_missingEndIndex_throws() throws { + let json = """ + { + "startIndex" : 10 + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + XCTAssertThrowsError(try decoder.decode(Citation.self, from: jsonData)) + } +}