Skip to content

Commit d88b035

Browse files
andrewheardandrewheard
committed
[Vertex AI] Update Citation decoding to handle optional startIndex (#12832)
1 parent bcd7d97 commit d88b035

File tree

4 files changed

+72
-32
lines changed

4 files changed

+72
-32
lines changed

‎FirebaseVertexAI/Sources/GenerateContentResponse.swift

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ public struct CitationMetadata: Decodable {
197197

198198
/// A struct describing a source attribution.
199199
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
200-
public struct Citation: Decodable {
200+
public struct Citation {
201201
/// The inclusive beginning of a sequence in a model response that derives from a cited source.
202202
public let startIndex: Int
203203

@@ -207,7 +207,7 @@ public struct Citation: Decodable {
207207
/// A link to the cited source.
208208
public let uri: String
209209

210-
/// The license the cited source work is distributed under.
210+
/// The license the cited source work is distributed under, if specified.
211211
public let license: String?
212212
}
213213

@@ -337,3 +337,23 @@ extension GenerateContentResponse.UsageMetadata: Decodable {
337337
totalTokenCount = try container.decodeIfPresent(Int.self, forKey: .totalTokenCount) ?? 0
338338
}
339339
}
340+
341+
// MARK: - Codable Conformances
342+
343+
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
344+
extension Citation: Decodable {
345+
enum CodingKeys: CodingKey {
346+
case startIndex
347+
case endIndex
348+
case uri
349+
case license
350+
}
351+
352+
public init(from decoder: any Decoder) throws {
353+
let container = try decoder.container(keyedBy: CodingKeys.self)
354+
startIndex = try container.decodeIfPresent(Int.self, forKey: .startIndex) ?? 0
355+
endIndex = try container.decode(Int.self, forKey: .endIndex)
356+
uri = try container.decode(String.self, forKey: .uri)
357+
license = try container.decodeIfPresent(String.self, forKey: .license)
358+
}
359+
}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
1-
data: {"candidates": [{"content": {"role": "model","parts": [{"text": "Some information"}]}}]}
1+
data: {"candidates": [{"content": {"role": "model","parts": [{"text": "Some information"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.043204036,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.082549304},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.046291895,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.071461484},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.100701615,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.06164962},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.13150747,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.040087357}]}]}
22

3-
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " More information"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.06632687,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.03825006},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.07477004,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.048767097},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.13695431,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.059866417},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.046119746,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.036425155}]}]}
3+
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " Some information cited from an external source"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.18982129,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.1337543},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.13637818,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.021906368},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.25404602,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.09073549},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.24202643,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.06548521}],"citationMetadata": {"citations": [{"endIndex": 128,"uri": "https://www.example.com/citation-1"},{"startIndex": 130,"endIndex": 265,"uri": "https://www.example.com/citation-2"}]}}]}
44

5-
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " Some information cited from an external source"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.07850098,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.039416388},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.08035747,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.04885778},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.12273335,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.059646938},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.053206205,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.04099903}],"citationMetadata": {"citations": [{"startIndex": 31,"endIndex": 187,"uri": "https://www.example.com/citation-1"}]}}]}
5+
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " More information cited from an external source"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.07850098,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.039416388},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.08035747,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.04885778},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.12273335,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.059646938},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.053206205,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.04099903}],"citationMetadata": {"citations": [{"startIndex": 272,"endIndex": 431,"uri": "https://www.example.com/citation-3","license": "mit"}]}}]}
66

7-
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " More information cited from an external source"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.08803312,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.044183318},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.094176665,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.0575992},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.13660839,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.08035747},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.060197048,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.046464667}],"citationMetadata": {"citations": [{"startIndex": 73,"endIndex": 248,"uri": "https://www.example.com/citation-2"},{"startIndex": 133,"endIndex": 272,"uri": "https://www.example.com/citation-3", "license": "mit"}]}}]}
8-
9-
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " More information"}]},"finishReason": "STOP","safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.12147716,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.0647717},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.11858909,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.053899158},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.14866412,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.08479541},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.05470151,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.045015533}]}],"usageMetadata": {"promptTokenCount": 9,"candidatesTokenCount": 163,"totalTokenCount": 172}}
7+
data: {"candidates": [{"content": {"role": "model","parts": [{"text": " More information "}]},"finishReason": "STOP","safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.16013464,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.11716747},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.10818896,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.021990221},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.2158462,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.07682221},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.19636348,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.08021325}]}],"usageMetadata": {"promptTokenCount": 9,"candidatesTokenCount": 53,"totalTokenCount": 62}}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,46 +14,56 @@
1414
{
1515
"category": "HARM_CATEGORY_HATE_SPEECH",
1616
"probability": "NEGLIGIBLE",
17-
"probabilityScore": 0.16013464,
17+
"probabilityScore": 0.16926852,
1818
"severity": "HARM_SEVERITY_NEGLIGIBLE",
19-
"severityScore": 0.074500255
19+
"severityScore": 0.08181271
2020
},
2121
{
2222
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
2323
"probability": "NEGLIGIBLE",
24-
"probabilityScore": 0.09687653,
24+
"probabilityScore": 0.15636235,
2525
"severity": "HARM_SEVERITY_NEGLIGIBLE",
26-
"severityScore": 0.049313594
26+
"severityScore": 0.02981654
2727
},
2828
{
2929
"category": "HARM_CATEGORY_HARASSMENT",
3030
"probability": "NEGLIGIBLE",
31-
"probabilityScore": 0.16817278,
31+
"probabilityScore": 0.33133608,
3232
"severity": "HARM_SEVERITY_NEGLIGIBLE",
33-
"severityScore": 0.09451043
33+
"severityScore": 0.10875559
3434
},
3535
{
3636
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
3737
"probability": "NEGLIGIBLE",
38-
"probabilityScore": 0.05023736,
38+
"probabilityScore": 0.17766814,
3939
"severity": "HARM_SEVERITY_NEGLIGIBLE",
40-
"severityScore": 0.034553625
40+
"severityScore": 0.18787657
4141
}
4242
],
4343
"citationMetadata": {
4444
"citations": [
4545
{
46-
"startIndex": 179,
47-
"endIndex": 366,
48-
"uri": "https://www.example.com/some-citation"
46+
"endIndex": 128,
47+
"uri": "https://www.example.com/some-citation-1"
48+
},
49+
{
50+
"startIndex": 130,
51+
"endIndex": 265,
52+
"uri": "https://www.example.com/some-citation-2"
53+
},
54+
{
55+
"startIndex": 272,
56+
"endIndex": 431,
57+
"uri": "https://www.example.com/some-citation-3",
58+
"license": "mit"
4959
}
5060
]
5161
}
5262
}
5363
],
5464
"usageMetadata": {
55-
"promptTokenCount": 11,
56-
"candidatesTokenCount": 135,
57-
"totalTokenCount": 146
65+
"promptTokenCount": 15,
66+
"candidatesTokenCount": 253,
67+
"totalTokenCount": 268
5868
}
5969
}

‎FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -108,12 +108,22 @@ final class GenerativeModelTests: XCTestCase {
108108
XCTAssertEqual(candidate.content.parts.count, 1)
109109
XCTAssertEqual(response.text, "Some information cited from an external source")
110110
let citationMetadata = try XCTUnwrap(candidate.citationMetadata)
111-
XCTAssertEqual(citationMetadata.citationSources.count, 1)
112-
let citationSource = try XCTUnwrap(citationMetadata.citationSources.first)
113-
XCTAssertEqual(citationSource.uri, "https://www.example.com/some-citation")
114-
XCTAssertEqual(citationSource.startIndex, 179)
115-
XCTAssertEqual(citationSource.endIndex, 366)
116-
XCTAssertNil(citationSource.license)
111+
XCTAssertEqual(citationMetadata.citationSources.count, 3)
112+
let citationSource1 = try XCTUnwrap(citationMetadata.citationSources[0])
113+
XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
114+
XCTAssertEqual(citationSource1.startIndex, 0)
115+
XCTAssertEqual(citationSource1.endIndex, 128)
116+
XCTAssertNil(citationSource1.license)
117+
let citationSource2 = try XCTUnwrap(citationMetadata.citationSources[1])
118+
XCTAssertEqual(citationSource2.uri, "https://www.example.com/some-citation-2")
119+
XCTAssertEqual(citationSource2.startIndex, 130)
120+
XCTAssertEqual(citationSource2.endIndex, 265)
121+
XCTAssertNil(citationSource2.license)
122+
let citationSource3 = try XCTUnwrap(citationMetadata.citationSources[2])
123+
XCTAssertEqual(citationSource3.uri, "https://www.example.com/some-citation-3")
124+
XCTAssertEqual(citationSource3.startIndex, 272)
125+
XCTAssertEqual(citationSource3.endIndex, 431)
126+
XCTAssertEqual(citationSource3.license, "mit")
117127
}
118128

119129
func testGenerateContent_success_quoteReply() async throws {
@@ -778,13 +788,15 @@ final class GenerativeModelTests: XCTestCase {
778788
XCTAssertEqual(citations.count, 3)
779789
XCTAssertTrue(citations
780790
.contains(where: {
781-
$0.startIndex == 31 && $0.endIndex == 187 && $0
782-
.uri == "https://www.example.com/citation-1" && $0.license == nil
791+
$0.startIndex == 0 && $0.endIndex == 128 && !$0.uri.isEmpty && $0.license == nil
783792
}))
784793
XCTAssertTrue(citations
785794
.contains(where: {
786-
$0.startIndex == 133 && $0.endIndex == 272 && $0
787-
.uri == "https://www.example.com/citation-3" && $0.license == "mit"
795+
$0.startIndex == 130 && $0.endIndex == 265 && !$0.uri.isEmpty && $0.license == nil
796+
}))
797+
XCTAssertTrue(citations
798+
.contains(where: {
799+
$0.startIndex == 272 && $0.endIndex == 431 && !$0.uri.isEmpty && $0.license == "mit"
788800
}))
789801
}
790802

0 commit comments

Comments
 (0)