Skip to content

Commit

Permalink
Update Citation decoding to handle optional values (#135)
Browse files Browse the repository at this point in the history
  • Loading branch information
@andrewheard
andrewheard authored Apr 18, 2024
1 parent fad5c49 commit f0e7743
Show file tree
Hide file tree
Showing 4 changed files with 77 additions and 15 deletions.
22 changes: 21 additions & 1 deletion Sources/GoogleAI/GenerateContentResponse.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ public struct CitationMetadata: Decodable {

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

Expand Down Expand Up @@ -297,3 +297,23 @@ extension PromptFeedback: Decodable {
}
}
}

// MARK: - Codable Conformances

@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
extension Citation: Decodable {
enum CodingKeys: CodingKey {
case startIndex
case endIndex
case uri
case license
}

public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
startIndex = try container.decodeIfPresent(Int.self, forKey: .startIndex) ?? 0
endIndex = try container.decode(Int.self, forKey: .endIndex)
uri = try container.decode(String.self, forKey: .uri)
license = try container.decodeIfPresent(String.self, forKey: .license) ?? ""
}
}
6 changes: 3 additions & 3 deletions Tests/GoogleAITests/GenerateContentResponses/streaming-success-citations.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@ data: {"candidates": [{"content": {"parts": [{"text": " More information"}],"rol

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"}]}]}

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

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

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

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": ""}]}}]}
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": ""}]}}]}
21 changes: 18 additions & 3 deletions Tests/GoogleAITests/GenerateContentResponses/unary-success-citations.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,25 @@
"citationMetadata": {
"citationSources": [
{
"startIndex": 574,
"endIndex": 705,
"uri": "https://www.example.com/some-citation",
"endIndex": 128,
"uri": "https://www.example.com/some-citation-1",
},
{
"startIndex": 130,
"endIndex": 265,
"uri": "https://www.example.com/some-citation-2",
},
{
"startIndex": 272,
"endIndex": 431,
"uri": "https://www.example.com/some-citation-3",
"license": ""
},
{
"startIndex": 444,
"endIndex": 630,
"uri": "https://www.example.com/some-citation-4",
"license": "mit"
}
]
}
Expand Down
43 changes: 35 additions & 8 deletions Tests/GoogleAITests/GenerativeModelTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,12 +104,27 @@ final class GenerativeModelTests: XCTestCase {
XCTAssertEqual(candidate.content.parts.count, 1)
XCTAssertEqual(response.text, "Some information cited from an external source")
let citationMetadata = try XCTUnwrap(candidate.citationMetadata)
XCTAssertEqual(citationMetadata.citationSources.count, 1)
let citationSource = try XCTUnwrap(citationMetadata.citationSources.first)
XCTAssertEqual(citationSource.uri, "https://www.example.com/some-citation")
XCTAssertEqual(citationSource.startIndex, 574)
XCTAssertEqual(citationSource.endIndex, 705)
XCTAssertEqual(citationSource.license, "")
XCTAssertEqual(citationMetadata.citationSources.count, 4)
let citationSource1 = try XCTUnwrap(citationMetadata.citationSources[0])
XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
XCTAssertEqual(citationSource1.startIndex, 0)
XCTAssertEqual(citationSource1.endIndex, 128)
XCTAssertEqual(citationSource1.license, "")
let citationSource2 = try XCTUnwrap(citationMetadata.citationSources[1])
XCTAssertEqual(citationSource2.uri, "https://www.example.com/some-citation-2")
XCTAssertEqual(citationSource2.startIndex, 130)
XCTAssertEqual(citationSource2.endIndex, 265)
XCTAssertEqual(citationSource2.license, "")
let citationSource3 = try XCTUnwrap(citationMetadata.citationSources[2])
XCTAssertEqual(citationSource3.uri, "https://www.example.com/some-citation-3")
XCTAssertEqual(citationSource3.startIndex, 272)
XCTAssertEqual(citationSource3.endIndex, 431)
XCTAssertEqual(citationSource3.license, "")
let citationSource4 = try XCTUnwrap(citationMetadata.citationSources[3])
XCTAssertEqual(citationSource4.uri, "https://www.example.com/some-citation-4")
XCTAssertEqual(citationSource4.startIndex, 444)
XCTAssertEqual(citationSource4.endIndex, 630)
XCTAssertEqual(citationSource4.license, "mit")
}

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

XCTAssertEqual(citations.count, 8)
XCTAssertTrue(citations
.contains(where: { $0.startIndex == 574 && $0.endIndex == 705 && !$0.uri.isEmpty }))
.contains(where: {
$0.startIndex == 0 && $0.endIndex == 128 && !$0.uri.isEmpty && $0.license.isEmpty
}))
XCTAssertTrue(citations
.contains(where: { $0.startIndex == 899 && $0.endIndex == 1026 && !$0.uri.isEmpty }))
.contains(where: {
$0.startIndex == 130 && $0.endIndex == 265 && !$0.uri.isEmpty && $0.license.isEmpty
}))
XCTAssertTrue(citations
.contains(where: {
$0.startIndex == 272 && $0.endIndex == 431 && !$0.uri.isEmpty && $0.license.isEmpty
}))
XCTAssertTrue(citations
.contains(where: {
$0.startIndex == 444 && $0.endIndex == 630 && !$0.uri.isEmpty && $0.license == "mit"
}))
}

func testGenerateContentStream_errorMidStream() async throws {
Expand Down

0 comments on commit f0e7743

Please sign in to comment.