Skip to content

Commit f0e7743

Browse files
andrewheardandrewheard
authored
Update Citation decoding to handle optional values (#135)
1 parent fad5c49 commit f0e7743

File tree

4 files changed

+77
-15
lines changed

4 files changed

+77
-15
lines changed

‎Sources/GoogleAI/GenerateContentResponse.swift

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ public struct CitationMetadata: Decodable {
174174

175175
/// A struct describing a source attribution.
176176
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
177-
public struct Citation: Decodable {
177+
public struct Citation {
178178
/// The inclusive beginning of a sequence in a model response that derives from a cited source.
179179
public let startIndex: Int
180180

@@ -297,3 +297,23 @@ extension PromptFeedback: Decodable {
297297
}
298298
}
299299
}
300+
301+
// MARK: - Codable Conformances
302+
303+
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
304+
extension Citation: Decodable {
305+
enum CodingKeys: CodingKey {
306+
case startIndex
307+
case endIndex
308+
case uri
309+
case license
310+
}
311+
312+
public init(from decoder: any Decoder) throws {
313+
let container = try decoder.container(keyedBy: CodingKeys.self)
314+
startIndex = try container.decodeIfPresent(Int.self, forKey: .startIndex) ?? 0
315+
endIndex = try container.decode(Int.self, forKey: .endIndex)
316+
uri = try container.decode(String.self, forKey: .uri)
317+
license = try container.decodeIfPresent(String.self, forKey: .license) ?? ""
318+
}
319+
}

‎Tests/GoogleAITests/GenerateContentResponses/streaming-success-citations.txt

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@ data: {"candidates": [{"content": {"parts": [{"text": " More information"}],"rol
44

55
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"}]}]}
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": {"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": [{"endIndex": 128,"uri": "https://www.example.com/citation-1"},{"startIndex": 130,"endIndex": 265,"uri": "https://www.example.com/citation-2"}]}}]}
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": ""}]}}]}
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": 272,"endIndex": 431,"uri": "https://www.example.com/citation-3","license": ""},{"startIndex": 444,"endIndex": 630,"uri": "https://www.example.com/citation-4","license": "mit"}]}}]}
1010

1111
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": ""}]}}]}
1212

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": ""}]}}]}
13+
data: {"candidates": [{"content": {"parts": [{"text": "More text"}],"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": ""}]}}]}

‎Tests/GoogleAITests/GenerateContentResponses/unary-success-citations.json

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,25 @@
3232
"citationMetadata": {
3333
"citationSources": [
3434
{
35-
"startIndex": 574,
36-
"endIndex": 705,
37-
"uri": "https://www.example.com/some-citation",
35+
"endIndex": 128,
36+
"uri": "https://www.example.com/some-citation-1",
37+
},
38+
{
39+
"startIndex": 130,
40+
"endIndex": 265,
41+
"uri": "https://www.example.com/some-citation-2",
42+
},
43+
{
44+
"startIndex": 272,
45+
"endIndex": 431,
46+
"uri": "https://www.example.com/some-citation-3",
3847
"license": ""
48+
},
49+
{
50+
"startIndex": 444,
51+
"endIndex": 630,
52+
"uri": "https://www.example.com/some-citation-4",
53+
"license": "mit"
3954
}
4055
]
4156
}

‎Tests/GoogleAITests/GenerativeModelTests.swift

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -104,12 +104,27 @@ final class GenerativeModelTests: XCTestCase {
104104
XCTAssertEqual(candidate.content.parts.count, 1)
105105
XCTAssertEqual(response.text, "Some information cited from an external source")
106106
let citationMetadata = try XCTUnwrap(candidate.citationMetadata)
107-
XCTAssertEqual(citationMetadata.citationSources.count, 1)
108-
let citationSource = try XCTUnwrap(citationMetadata.citationSources.first)
109-
XCTAssertEqual(citationSource.uri, "https://www.example.com/some-citation")
110-
XCTAssertEqual(citationSource.startIndex, 574)
111-
XCTAssertEqual(citationSource.endIndex, 705)
112-
XCTAssertEqual(citationSource.license, "")
107+
XCTAssertEqual(citationMetadata.citationSources.count, 4)
108+
let citationSource1 = try XCTUnwrap(citationMetadata.citationSources[0])
109+
XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
110+
XCTAssertEqual(citationSource1.startIndex, 0)
111+
XCTAssertEqual(citationSource1.endIndex, 128)
112+
XCTAssertEqual(citationSource1.license, "")
113+
let citationSource2 = try XCTUnwrap(citationMetadata.citationSources[1])
114+
XCTAssertEqual(citationSource2.uri, "https://www.example.com/some-citation-2")
115+
XCTAssertEqual(citationSource2.startIndex, 130)
116+
XCTAssertEqual(citationSource2.endIndex, 265)
117+
XCTAssertEqual(citationSource2.license, "")
118+
let citationSource3 = try XCTUnwrap(citationMetadata.citationSources[2])
119+
XCTAssertEqual(citationSource3.uri, "https://www.example.com/some-citation-3")
120+
XCTAssertEqual(citationSource3.startIndex, 272)
121+
XCTAssertEqual(citationSource3.endIndex, 431)
122+
XCTAssertEqual(citationSource3.license, "")
123+
let citationSource4 = try XCTUnwrap(citationMetadata.citationSources[3])
124+
XCTAssertEqual(citationSource4.uri, "https://www.example.com/some-citation-4")
125+
XCTAssertEqual(citationSource4.startIndex, 444)
126+
XCTAssertEqual(citationSource4.endIndex, 630)
127+
XCTAssertEqual(citationSource4.license, "mit")
113128
}
114129

115130
func testGenerateContent_success_quoteReply() async throws {
@@ -724,9 +739,21 @@ final class GenerativeModelTests: XCTestCase {
724739

725740
XCTAssertEqual(citations.count, 8)
726741
XCTAssertTrue(citations
727-
.contains(where: { $0.startIndex == 574 && $0.endIndex == 705 && !$0.uri.isEmpty }))
742+
.contains(where: {
743+
$0.startIndex == 0 && $0.endIndex == 128 && !$0.uri.isEmpty && $0.license.isEmpty
744+
}))
728745
XCTAssertTrue(citations
729-
.contains(where: { $0.startIndex == 899 && $0.endIndex == 1026 && !$0.uri.isEmpty }))
746+
.contains(where: {
747+
$0.startIndex == 130 && $0.endIndex == 265 && !$0.uri.isEmpty && $0.license.isEmpty
748+
}))
749+
XCTAssertTrue(citations
750+
.contains(where: {
751+
$0.startIndex == 272 && $0.endIndex == 431 && !$0.uri.isEmpty && $0.license.isEmpty
752+
}))
753+
XCTAssertTrue(citations
754+
.contains(where: {
755+
$0.startIndex == 444 && $0.endIndex == 630 && !$0.uri.isEmpty && $0.license == "mit"
756+
}))
730757
}
731758

732759
func testGenerateContentStream_errorMidStream() async throws {

0 commit comments

Comments
 (0)