Skip to content

Commit a96d7d9

Browse files
andrewheardandrewheard
authored
[Firebase AI] Add handling for Google AI-formatted CitationMetadata (#14780)
1 parent 42d402f commit a96d7d9

File tree

3 files changed

+223
-2
lines changed

3 files changed

+223
-2
lines changed

‎FirebaseAI/Sources/GenerateContentResponse.swift

+32-2
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ public struct CitationMetadata: Sendable {
145145

146146
/// A struct describing a source attribution.
147147
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
148-
public struct Citation: Sendable {
148+
public struct Citation: Sendable, Equatable {
149149
/// The inclusive beginning of a sequence in a model response that derives from a cited source.
150150
public let startIndex: Int
151151

@@ -165,6 +165,20 @@ public struct Citation: Sendable {
165165
///
166166
/// > Tip: `DateComponents` can be converted to a `Date` using the `date` computed property.
167167
public let publicationDate: DateComponents?
168+
169+
init(startIndex: Int,
170+
endIndex: Int,
171+
uri: String? = nil,
172+
title: String? = nil,
173+
license: String? = nil,
174+
publicationDate: DateComponents? = nil) {
175+
self.startIndex = startIndex
176+
self.endIndex = endIndex
177+
self.uri = uri
178+
self.title = title
179+
self.license = license
180+
self.publicationDate = publicationDate
181+
}
168182
}
169183

170184
/// A value enumerating possible reasons for a model to terminate a content generation request.
@@ -385,7 +399,23 @@ extension Candidate: Decodable {
385399
}
386400

387401
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
388-
extension CitationMetadata: Decodable {}
402+
extension CitationMetadata: Decodable {
403+
enum CodingKeys: CodingKey {
404+
case citations // Vertex AI
405+
case citationSources // Google AI
406+
}
407+
408+
public init(from decoder: any Decoder) throws {
409+
let container = try decoder.container(keyedBy: CodingKeys.self)
410+
411+
// Decode for Google API if `citationSources` key is present.
412+
if container.contains(.citationSources) {
413+
citations = try container.decode([Citation].self, forKey: .citationSources)
414+
} else { // Fallback to default Vertex AI decoding.
415+
citations = try container.decode([Citation].self, forKey: .citations)
416+
}
417+
}
418+
}
389419

390420
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
391421
extension Citation: Decodable {

‎FirebaseAI/Tests/Unit/Types/CitationMetadataTests.swift

+78
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import XCTest
16+
17+
@testable import FirebaseAI
18+
19+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
20+
final class CitationMetadataTests: XCTestCase {
21+
let decoder = JSONDecoder()
22+
23+
let expectedStartIndex = 100
24+
let expectedEndIndex = 200
25+
let expectedURI = "https://example.com/citation-1"
26+
lazy var citationJSON = """
27+
{
28+
"startIndex" : \(expectedStartIndex),
29+
"endIndex" : \(expectedEndIndex),
30+
"uri" : "\(expectedURI)"
31+
}
32+
"""
33+
lazy var expectedCitation = Citation(
34+
startIndex: expectedStartIndex,
35+
endIndex: expectedEndIndex,
36+
uri: expectedURI
37+
)
38+
39+
// MARK: - Google AI Format Decoding
40+
41+
func testDecodeCitationMetadata_googleAIFormat() throws {
42+
let json = """
43+
{
44+
"citationSources": [\(citationJSON)]
45+
}
46+
"""
47+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
48+
49+
let citationMetadata = try decoder.decode(
50+
CitationMetadata.self,
51+
from: jsonData
52+
)
53+
54+
XCTAssertEqual(citationMetadata.citations.count, 1)
55+
let citation = try XCTUnwrap(citationMetadata.citations.first)
56+
XCTAssertEqual(citation, expectedCitation)
57+
}
58+
59+
// MARK: - Vertex AI Format Decoding
60+
61+
func testDecodeCitationMetadata_vertexAIFormat() throws {
62+
let json = """
63+
{
64+
"citations": [\(citationJSON)]
65+
}
66+
"""
67+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
68+
69+
let citationMetadata = try decoder.decode(
70+
CitationMetadata.self,
71+
from: jsonData
72+
)
73+
74+
XCTAssertEqual(citationMetadata.citations.count, 1)
75+
let citation = try XCTUnwrap(citationMetadata.citations.first)
76+
XCTAssertEqual(citation, expectedCitation)
77+
}
78+
}

‎FirebaseAI/Tests/Unit/Types/CitationTests.swift

+113
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import FirebaseAI
16+
import XCTest
17+
18+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
19+
final class CitationTests: XCTestCase {
20+
let decoder = JSONDecoder()
21+
22+
// MARK: - Decoding Tests
23+
24+
func testDecodeCitation_minimalParameters() throws {
25+
let expectedEndIndex = 150
26+
let json = """
27+
{
28+
"endIndex" : \(expectedEndIndex)
29+
}
30+
"""
31+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
32+
33+
let citation = try decoder.decode(Citation.self, from: jsonData)
34+
35+
XCTAssertEqual(citation.startIndex, 0, "Omitted startIndex should be decoded as 0.")
36+
XCTAssertEqual(citation.endIndex, expectedEndIndex)
37+
XCTAssertNil(citation.uri)
38+
XCTAssertNil(citation.title)
39+
XCTAssertNil(citation.license)
40+
XCTAssertNil(citation.publicationDate)
41+
}
42+
43+
func testDecodeCitation_allParameters() throws {
44+
let expectedStartIndex = 100
45+
let expectedEndIndex = 200
46+
let expectedURI = "https://example.com/citation-1"
47+
let expectedTitle = "Example Citation Title"
48+
let expectedLicense = "mit"
49+
let expectedYear = 2023
50+
let expectedMonth = 10
51+
let expectedDay = 26
52+
let json = """
53+
{
54+
"startIndex" : \(expectedStartIndex),
55+
"endIndex" : \(expectedEndIndex),
56+
"uri" : "\(expectedURI)",
57+
"title" : "\(expectedTitle)",
58+
"license" : "\(expectedLicense)",
59+
"publicationDate" : {
60+
"year" : \(expectedYear),
61+
"month" : \(expectedMonth),
62+
"day" : \(expectedDay)
63+
}
64+
}
65+
"""
66+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
67+
68+
let citation = try decoder.decode(Citation.self, from: jsonData)
69+
70+
XCTAssertEqual(citation.startIndex, expectedStartIndex)
71+
XCTAssertEqual(citation.endIndex, expectedEndIndex)
72+
XCTAssertEqual(citation.uri, expectedURI)
73+
XCTAssertEqual(citation.title, expectedTitle)
74+
XCTAssertEqual(citation.license, expectedLicense)
75+
let publicationDate = try XCTUnwrap(citation.publicationDate)
76+
XCTAssertEqual(publicationDate.year, expectedYear)
77+
XCTAssertEqual(publicationDate.month, expectedMonth)
78+
XCTAssertEqual(publicationDate.day, expectedDay)
79+
}
80+
81+
func testDecodeCitation_emptyStringsForOptionals_setsToNil() throws {
82+
let expectedEndIndex = 300
83+
let json = """
84+
{
85+
"endIndex" : \(expectedEndIndex),
86+
"uri" : "",
87+
"title" : "",
88+
"license" : ""
89+
}
90+
"""
91+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
92+
93+
let citation = try decoder.decode(Citation.self, from: jsonData)
94+
95+
XCTAssertEqual(citation.startIndex, 0, "Omitted startIndex should be decoded as 0.")
96+
XCTAssertEqual(citation.endIndex, expectedEndIndex)
97+
XCTAssertNil(citation.uri, "Empty URI string should be decoded as nil.")
98+
XCTAssertNil(citation.title, "Empty title string should be decoded as nil.")
99+
XCTAssertNil(citation.license, "Empty license string should be decoded as nil.")
100+
XCTAssertNil(citation.publicationDate)
101+
}
102+
103+
func testDecodeCitation_missingEndIndex_throws() throws {
104+
let json = """
105+
{
106+
"startIndex" : 10
107+
}
108+
"""
109+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
110+
111+
XCTAssertThrowsError(try decoder.decode(Citation.self, from: jsonData))
112+
}
113+
}

0 commit comments

Comments
 (0)