From ddcece5dc6a62b46896967f8fc0380599755c3e1 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Fri, 19 Apr 2024 16:05:41 -0400
Subject: [PATCH] Make `Citation.license` optional (#141)

---
 Sources/GoogleAI/GenerateContentResponse.swift | 11 ++++++++---
 Tests/GoogleAITests/GenerativeModelTests.swift | 12 ++++++------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/Sources/GoogleAI/GenerateContentResponse.swift b/Sources/GoogleAI/GenerateContentResponse.swift
index 6ac4576..e623d74 100644
--- 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
index d5efcd8..1ce8816 100644
--- 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: {