Skip to content

Commit fdc2985

Browse files
andrewheardandrewheard
committed
Fix CitationMetadata parsing in Vertex AI (#12626)
1 parent 82d7335 commit fdc2985

File tree

4 files changed

+56
-52
lines changed

4 files changed

+56
-52
lines changed

‎FirebaseVertexAI/Sources/GenerateContentResponse.swift

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,10 @@ extension CandidateResponse: Decodable {
155155
/// A collection of source attributions for a piece of content.
156156
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
157157
public struct CitationMetadata: Decodable {
158+
enum CodingKeys: String, CodingKey {
159+
case citationSources = "citations"
160+
}
161+
158162
/// A list of individual cited sources and the parts of the content to which they apply.
159163
public let citationSources: [Citation]
160164
}
@@ -172,7 +176,7 @@ public struct Citation: Decodable {
172176
public let uri: String
173177

174178
/// The license the cited source work is distributed under.
175-
public let license: String
179+
public let license: String?
176180
}
177181

178182
/// A value enumerating possible reasons for a model to terminate a content generation request.

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

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,9 @@
1-
data: {"candidates": [{"content": {"parts": [{"text": "Some information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}}
1+
data: {"candidates": [{"content": {"role": "model","parts": [{"text": "Some information"}]}}]}
22

3-
data: {"candidates": [{"content": {"parts": [{"text": " More information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]}
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}]}]}
44

5-
data: {"candidates": [{"content": {"parts": [{"text": ", Even more information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]}
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"}]}}]}
66

7-
data: {"candidates": [{"content": {"parts": [{"text": " Some information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.example.com/citation-1","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.example.com/citation-2","license": ""}]}}]}
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"}]}}]}
88

9-
data: {"candidates": [{"content": {"parts": [{"text": "More information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.example.com/citation-3","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.example.com/citation-4","license": ""}]}}]}
10-
11-
data: {"candidates": [{"content": {"parts": [{"text": "Even more information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.example.com/citation-5","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.example.com/citation-6","license": ""}]}}]}
12-
13-
data: {"candidates": [{"content": {"parts": [{"text": "Physics (YouTube Channel)"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.google.com","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.google.com","license": ""}]}}]}
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}}

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

Lines changed: 29 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,63 +2,58 @@
22
"candidates": [
33
{
44
"content": {
5+
"role": "model",
56
"parts": [
67
{
78
"text": "Some information cited from an external source"
89
}
9-
],
10-
"role": "model"
10+
]
1111
},
1212
"finishReason": "STOP",
13-
"index": 0,
1413
"safetyRatings": [
1514
{
16-
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
17-
"probability": "NEGLIGIBLE"
15+
"category": "HARM_CATEGORY_HATE_SPEECH",
16+
"probability": "NEGLIGIBLE",
17+
"probabilityScore": 0.16013464,
18+
"severity": "HARM_SEVERITY_NEGLIGIBLE",
19+
"severityScore": 0.074500255
1820
},
1921
{
20-
"category": "HARM_CATEGORY_HATE_SPEECH",
21-
"probability": "NEGLIGIBLE"
22+
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
23+
"probability": "NEGLIGIBLE",
24+
"probabilityScore": 0.09687653,
25+
"severity": "HARM_SEVERITY_NEGLIGIBLE",
26+
"severityScore": 0.049313594
2227
},
2328
{
2429
"category": "HARM_CATEGORY_HARASSMENT",
25-
"probability": "NEGLIGIBLE"
30+
"probability": "NEGLIGIBLE",
31+
"probabilityScore": 0.16817278,
32+
"severity": "HARM_SEVERITY_NEGLIGIBLE",
33+
"severityScore": 0.09451043
2634
},
2735
{
28-
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
29-
"probability": "NEGLIGIBLE"
36+
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
37+
"probability": "NEGLIGIBLE",
38+
"probabilityScore": 0.05023736,
39+
"severity": "HARM_SEVERITY_NEGLIGIBLE",
40+
"severityScore": 0.034553625
3041
}
3142
],
3243
"citationMetadata": {
33-
"citationSources": [
44+
"citations": [
3445
{
35-
"startIndex": 574,
36-
"endIndex": 705,
37-
"uri": "https://www.example.com/some-citation",
38-
"license": ""
46+
"startIndex": 179,
47+
"endIndex": 366,
48+
"uri": "https://www.example.com/some-citation"
3949
}
4050
]
4151
}
4252
}
4353
],
44-
"promptFeedback": {
45-
"safetyRatings": [
46-
{
47-
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
48-
"probability": "NEGLIGIBLE"
49-
},
50-
{
51-
"category": "HARM_CATEGORY_HATE_SPEECH",
52-
"probability": "NEGLIGIBLE"
53-
},
54-
{
55-
"category": "HARM_CATEGORY_HARASSMENT",
56-
"probability": "NEGLIGIBLE"
57-
},
58-
{
59-
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
60-
"probability": "NEGLIGIBLE"
61-
}
62-
]
54+
"usageMetadata": {
55+
"promptTokenCount": 11,
56+
"candidatesTokenCount": 135,
57+
"totalTokenCount": 146
6358
}
6459
}

‎FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -107,9 +107,9 @@ final class GenerativeModelTests: XCTestCase {
107107
XCTAssertEqual(citationMetadata.citationSources.count, 1)
108108
let citationSource = try XCTUnwrap(citationMetadata.citationSources.first)
109109
XCTAssertEqual(citationSource.uri, "https://www.example.com/some-citation")
110-
XCTAssertEqual(citationSource.startIndex, 574)
111-
XCTAssertEqual(citationSource.endIndex, 705)
112-
XCTAssertEqual(citationSource.license, "")
110+
XCTAssertEqual(citationSource.startIndex, 179)
111+
XCTAssertEqual(citationSource.endIndex, 366)
112+
XCTAssertNil(citationSource.license)
113113
}
114114

115115
func testGenerateContent_success_quoteReply() async throws {
@@ -669,21 +669,30 @@ final class GenerativeModelTests: XCTestCase {
669669
)
670670

671671
let stream = model.generateContentStream("Hi")
672-
var citations: [Citation] = []
672+
var citations = [Citation]()
673+
var responses = [GenerateContentResponse]()
673674
for try await content in stream {
675+
responses.append(content)
674676
XCTAssertNotNil(content.text)
675677
let candidate = try XCTUnwrap(content.candidates.first)
676-
XCTAssertEqual(candidate.finishReason, .stop)
677678
if let sources = candidate.citationMetadata?.citationSources {
678679
citations.append(contentsOf: sources)
679680
}
680681
}
681682

682-
XCTAssertEqual(citations.count, 8)
683+
let lastCandidate = try XCTUnwrap(responses.last?.candidates.first)
684+
XCTAssertEqual(lastCandidate.finishReason, .stop)
685+
XCTAssertEqual(citations.count, 3)
683686
XCTAssertTrue(citations
684-
.contains(where: { $0.startIndex == 574 && $0.endIndex == 705 && !$0.uri.isEmpty }))
687+
.contains(where: {
688+
$0.startIndex == 31 && $0.endIndex == 187 && $0
689+
.uri == "https://www.example.com/citation-1" && $0.license == nil
690+
}))
685691
XCTAssertTrue(citations
686-
.contains(where: { $0.startIndex == 899 && $0.endIndex == 1026 && !$0.uri.isEmpty }))
692+
.contains(where: {
693+
$0.startIndex == 133 && $0.endIndex == 272 && $0
694+
.uri == "https://www.example.com/citation-3" && $0.license == "mit"
695+
}))
687696
}
688697

689698
func testGenerateContentStream_appCheck_validToken() async throws {

0 commit comments

Comments
 (0)