From 679e7e6f48b0875146c15241218f2dc2803ba370 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Mon, 22 Apr 2024 14:34:12 -0400 Subject: [PATCH] Update Citation decoding to handle optional startIndex --- .../Sources/GenerateContentResponse.swift | 24 +++++++++++- .../streaming-success-citations.txt | 10 ++--- .../unary-success-citations.json | 38 ++++++++++++------- .../Tests/Unit/GenerativeModelTests.swift | 32 +++++++++++----- 4 files changed, 72 insertions(+), 32 deletions(-) diff --git a/FirebaseVertexAI/Sources/GenerateContentResponse.swift b/FirebaseVertexAI/Sources/GenerateContentResponse.swift index 29c5b14f552..48f4ff22413 100644 --- a/FirebaseVertexAI/Sources/GenerateContentResponse.swift +++ b/FirebaseVertexAI/Sources/GenerateContentResponse.swift @@ -197,7 +197,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 @@ -207,7 +207,7 @@ public struct Citation: Decodable { /// A link to the cited source. public let uri: String - /// The license the cited source work is distributed under. + /// The license the cited source work is distributed under, if specified. public let license: String? } @@ -337,3 +337,23 @@ extension GenerateContentResponse.UsageMetadata: Decodable { totalTokenCount = try container.decodeIfPresent(Int.self, forKey: .totalTokenCount) ?? 0 } } + +// 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) + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt index a530736b308..fdc1f20d8b6 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt @@ -1,9 +1,7 @@ -data: {"candidates": [{"content": {"role": "model","parts": [{"text": "Some information"}]}}]} +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}]}]} -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}]}]} +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"}]}}]} -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"}]}}]} +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"}]}}]} -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"}]}}]} - -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}} +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}} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json index a0dacc3e883..68fc9648773 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json @@ -14,46 +14,56 @@ { "category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE", - "probabilityScore": 0.16013464, + "probabilityScore": 0.16926852, "severity": "HARM_SEVERITY_NEGLIGIBLE", - "severityScore": 0.074500255 + "severityScore": 0.08181271 }, { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "probability": "NEGLIGIBLE", - "probabilityScore": 0.09687653, + "probabilityScore": 0.15636235, "severity": "HARM_SEVERITY_NEGLIGIBLE", - "severityScore": 0.049313594 + "severityScore": 0.02981654 }, { "category": "HARM_CATEGORY_HARASSMENT", "probability": "NEGLIGIBLE", - "probabilityScore": 0.16817278, + "probabilityScore": 0.33133608, "severity": "HARM_SEVERITY_NEGLIGIBLE", - "severityScore": 0.09451043 + "severityScore": 0.10875559 }, { "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "probability": "NEGLIGIBLE", - "probabilityScore": 0.05023736, + "probabilityScore": 0.17766814, "severity": "HARM_SEVERITY_NEGLIGIBLE", - "severityScore": 0.034553625 + "severityScore": 0.18787657 } ], "citationMetadata": { "citations": [ { - "startIndex": 179, - "endIndex": 366, - "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": "mit" } ] } } ], "usageMetadata": { - "promptTokenCount": 11, - "candidatesTokenCount": 135, - "totalTokenCount": 146 + "promptTokenCount": 15, + "candidatesTokenCount": 253, + "totalTokenCount": 268 } } diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift index e7ea54f5dd9..c9442e05ba7 100644 --- a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -108,12 +108,22 @@ 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, 179) - XCTAssertEqual(citationSource.endIndex, 366) - XCTAssertNil(citationSource.license) + XCTAssertEqual(citationMetadata.citationSources.count, 3) + 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) + XCTAssertNil(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) + XCTAssertNil(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, "mit") } func testGenerateContent_success_quoteReply() async throws { @@ -778,13 +788,15 @@ final class GenerativeModelTests: XCTestCase { XCTAssertEqual(citations.count, 3) XCTAssertTrue(citations .contains(where: { - $0.startIndex == 31 && $0.endIndex == 187 && $0 - .uri == "https://www.example.com/citation-1" && $0.license == nil + $0.startIndex == 0 && $0.endIndex == 128 && !$0.uri.isEmpty && $0.license == nil })) XCTAssertTrue(citations .contains(where: { - $0.startIndex == 133 && $0.endIndex == 272 && $0 - .uri == "https://www.example.com/citation-3" && $0.license == "mit" + $0.startIndex == 130 && $0.endIndex == 265 && !$0.uri.isEmpty && $0.license == nil + })) + XCTAssertTrue(citations + .contains(where: { + $0.startIndex == 272 && $0.endIndex == 431 && !$0.uri.isEmpty && $0.license == "mit" })) }