Make Citation.license optional (google-gemini#141)

SsomsakTH · Jun 21, 2024 · ae96e10 · ae96e10
1 parent 039114b
commit ae96e10
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 9 deletions.
diff --git a/Sources/GoogleAI/GenerateContentResponse.swift b/Sources/GoogleAI/GenerateContentResponse.swift
@@ -184,8 +184,8 @@ public struct Citation {
   /// A link to the cited source.
   public let uri: String
 
-  /// The license the cited source work is distributed under.
-  public let license: String
+  /// The license the cited source work is distributed under, if specified.
+  public let license: String?
 }
 
 /// A value enumerating possible reasons for a model to terminate a content generation request.
@@ -314,6 +314,11 @@ extension Citation: Decodable {
     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) ?? ""
+    if let license = try container.decodeIfPresent(String.self, forKey: .license),
+       !license.isEmpty {
+      self.license = license
+    } else {
+      license = nil
+    }
   }
 }
diff --git a/Tests/GoogleAITests/GenerativeModelTests.swift b/Tests/GoogleAITests/GenerativeModelTests.swift
@@ -109,17 +109,17 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
     XCTAssertEqual(citationSource1.startIndex, 0)
     XCTAssertEqual(citationSource1.endIndex, 128)
-    XCTAssertEqual(citationSource1.license, "")
+    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)
-    XCTAssertEqual(citationSource2.license, "")
+    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, "")
+    XCTAssertNil(citationSource3.license)
     let citationSource4 = try XCTUnwrap(citationMetadata.citationSources[3])
     XCTAssertEqual(citationSource4.uri, "https://www.example.com/some-citation-4")
     XCTAssertEqual(citationSource4.startIndex, 444)
@@ -740,15 +740,15 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(citations.count, 8)
     XCTAssertTrue(citations
       .contains(where: {
-        $0.startIndex == 0 && $0.endIndex == 128 && !$0.uri.isEmpty && $0.license.isEmpty
+        $0.startIndex == 0 && $0.endIndex == 128 && !$0.uri.isEmpty && $0.license == nil
       }))
     XCTAssertTrue(citations
       .contains(where: {
-        $0.startIndex == 130 && $0.endIndex == 265 && !$0.uri.isEmpty && $0.license.isEmpty
+        $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.isEmpty
+        $0.startIndex == 272 && $0.endIndex == 431 && !$0.uri.isEmpty && $0.license == nil
       }))
     XCTAssertTrue(citations
       .contains(where: {