[Vertex AI] Add Citation.publicationDate (#13893)

Author: andrewheard

FirebaseVertexAI/CHANGELOG.md

@@ -72,6 +72,8 @@
   that may be reported when a prompt is blocked. (#13861)
 - [added] Added the `PromptFeedback` property `blockReasonMessage` that *may* be
   provided alongside the `blockReason`. (#13891)
+- [added] Added an optional `publicationDate` property that *may* be provided in
+  `Citation`. (#13893)
 
 # 11.3.0
 - [added] Added `Decodable` conformance for `FunctionResponse`. (#13606)

FirebaseVertexAI/Sources/GenerateContentResponse.swift

@@ -141,6 +141,11 @@ public struct Citation: Sendable {
 
   /// The license the cited source work is distributed under, if specified.
   public let license: String?
+
+  /// The publication date of the cited source, if available.
+  ///
+  /// > Tip: `DateComponents` can be converted to a `Date` using the `date` computed property.
+  public let publicationDate: DateComponents?
 }
 
 /// A value enumerating possible reasons for a model to terminate a content generation request.
@@ -363,28 +368,47 @@ extension Citation: Decodable {
     case uri
     case title
     case license
+    case publicationDate
   }
 
   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)
+
     if let uri = try container.decodeIfPresent(String.self, forKey: .uri), !uri.isEmpty {
       self.uri = uri
     } else {
       uri = nil
     }
+
     if let title = try container.decodeIfPresent(String.self, forKey: .title), !title.isEmpty {
       self.title = title
     } else {
       title = nil
     }
+
     if let license = try container.decodeIfPresent(String.self, forKey: .license),
        !license.isEmpty {
       self.license = license
     } else {
       license = nil
     }
+
+    if let publicationProtoDate = try container.decodeIfPresent(
+      ProtoDate.self,
+      forKey: .publicationDate
+    ) {
+      publicationDate = publicationProtoDate.dateComponents
+      if let publicationDate, !publicationDate.isValidDate {
+        VertexLog.warning(
+          code: .decodedInvalidCitationPublicationDate,
+          "Decoded an invalid citation publication date: \(publicationDate)"
+        )
+      }
+    } else {
+      publicationDate = nil
+    }
   }
 }
 

FirebaseVertexAI/Sources/Types/Internal/ProtoDate.swift

@@ -0,0 +1,117 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import Foundation
+
+/// Represents a whole or partial calendar date, such as a birthday.
+///
+/// The time of day and time zone are either specified elsewhere or are insignificant. The date is
+/// relative to the Gregorian Calendar. This can represent one of the following:
+/// - A full date, with non-zero year, month, and day values
+/// - A month and day value, with a zero year, such as an anniversary
+/// - A year on its own, with zero month and day values
+/// - A year and month value, with a zero day, such as a credit card expiration date
+///
+/// This represents a
+/// [`google.type.Date`](https://cloud.google.com/vertex-ai/docs/reference/rest/Shared.Types/Date).
+struct ProtoDate {
+  /// Year of the date.
+  ///
+  /// Must be from 1 to 9999, or 0 to specify a date without a year.
+  let year: Int?
+
+  /// Month of a year.
+  ///
+  /// Must be from 1 to 12, or 0 to specify a year without a month and day.
+  let month: Int?
+
+  /// Day of a month.
+  ///
+  /// Must be from 1 to 31 and valid for the year and month, or 0 to specify a year by itself or a
+  /// year and month where the day isn't significant.
+  let day: Int?
+
+  /// Returns the a `DateComponents` representation of the `ProtoDate`.
+  ///
+  /// > Note: This uses the Gregorian `Calendar` to match the `google.type.Date` definition.
+  var dateComponents: DateComponents {
+    DateComponents(
+      calendar: Calendar(identifier: .gregorian),
+      year: year,
+      month: month,
+      day: day
+    )
+  }
+}
+
+// MARK: - Codable Conformance
+
+extension ProtoDate: Decodable {
+  enum CodingKeys: CodingKey {
+    case year
+    case month
+    case day
+  }
+
+  init(from decoder: any Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    if let year = try container.decodeIfPresent(Int.self, forKey: .year), year != 0 {
+      if year < 0 || year > 9999 {
+        VertexLog.warning(
+          code: .decodedInvalidProtoDateYear,
+          """
+          Invalid year: \(year); must be from 1 to 9999, or 0 for a date without a specified year.
+          """
+        )
+      }
+      self.year = year
+    } else {
+      year = nil
+    }
+
+    if let month = try container.decodeIfPresent(Int.self, forKey: .month), month != 0 {
+      if month < 0 || month > 12 {
+        VertexLog.warning(
+          code: .decodedInvalidProtoDateMonth,
+          """
+          Invalid month: \(month); must be from 1 to 12, or 0 for a year date without a specified \
+          month and day.
+          """
+        )
+      }
+      self.month = month
+    } else {
+      month = nil
+    }
+
+    if let day = try container.decodeIfPresent(Int.self, forKey: .day), day != 0 {
+      if day < 0 || day > 31 {
+        VertexLog.warning(
+          code: .decodedInvalidProtoDateDay,
+          "Invalid day: \(day); must be from 1 to 31, or 0 for a date without a specified day."
+        )
+      }
+      self.day = day
+    } else {
+      day = nil
+    }
+
+    guard year != nil || month != nil || day != nil else {
+      throw DecodingError.dataCorrupted(.init(
+        codingPath: [CodingKeys.year, CodingKeys.month, CodingKeys.day],
+        debugDescription: "Invalid date: missing year, month and day"
+      ))
+    }
+  }
+}

FirebaseVertexAI/Sources/VertexLog.swift

@@ -50,6 +50,10 @@ enum VertexLog {
     case generateContentResponseUnrecognizedHarmProbability = 3005
     case generateContentResponseUnrecognizedHarmCategory = 3006
     case generateContentResponseUnrecognizedHarmSeverity = 3007
+    case decodedInvalidProtoDateYear = 3008
+    case decodedInvalidProtoDateMonth = 3009
+    case decodedInvalidProtoDateDay = 3010
+    case decodedInvalidCitationPublicationDate = 3011
 
     // SDK State Errors
     case generateContentResponseNoCandidates = 4000

FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

@@ -134,6 +134,12 @@ final class GenerativeModelTests: XCTestCase {
         forResource: "unary-success-citations",
         withExtension: "json"
       )
+    let expectedPublicationDate = DateComponents(
+      calendar: Calendar(identifier: .gregorian),
+      year: 2019,
+      month: 5,
+      day: 10
+    )
 
     let response = try await model.generateContent(testPrompt)
 
@@ -149,8 +155,10 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(citationSource1.endIndex, 128)
     XCTAssertNil(citationSource1.title)
     XCTAssertNil(citationSource1.license)
+    XCTAssertNil(citationSource1.publicationDate)
     let citationSource2 = try XCTUnwrap(citationMetadata.citations[1])
     XCTAssertEqual(citationSource2.title, "some-citation-2")
+    XCTAssertEqual(citationSource2.publicationDate, expectedPublicationDate)
     XCTAssertEqual(citationSource2.startIndex, 130)
     XCTAssertEqual(citationSource2.endIndex, 265)
     XCTAssertNil(citationSource2.uri)
@@ -161,6 +169,7 @@ final class GenerativeModelTests: XCTestCase {
     XCTAssertEqual(citationSource3.endIndex, 431)
     XCTAssertEqual(citationSource3.license, "mit")
     XCTAssertNil(citationSource3.title)
+    XCTAssertNil(citationSource3.publicationDate)
   }
 
   func testGenerateContent_success_quoteReply() async throws {
@@ -1052,6 +1061,12 @@ final class GenerativeModelTests: XCTestCase {
         forResource: "streaming-success-citations",
         withExtension: "txt"
       )
+    let expectedPublicationDate = DateComponents(
+      calendar: Calendar(identifier: .gregorian),
+      year: 2014,
+      month: 3,
+      day: 30
+    )
 
     let stream = try model.generateContentStream("Hi")
     var citations = [Citation]()
@@ -1072,18 +1087,19 @@ final class GenerativeModelTests: XCTestCase {
       .contains {
         $0.startIndex == 0 && $0.endIndex == 128
           && $0.uri == "https://www.example.com/some-citation-1" && $0.title == nil
-          && $0.license == nil
+          && $0.license == nil && $0.publicationDate == nil
       })
     XCTAssertTrue(citations
       .contains {
         $0.startIndex == 130 && $0.endIndex == 265 && $0.uri == nil
           && $0.title == "some-citation-2" && $0.license == nil
+          && $0.publicationDate == expectedPublicationDate
       })
     XCTAssertTrue(citations
       .contains {
         $0.startIndex == 272 && $0.endIndex == 431
           && $0.uri == "https://www.example.com/some-citation-3" && $0.title == nil
-          && $0.license == "mit"
+          && $0.license == "mit" && $0.publicationDate == nil
       })
     XCTAssertFalse(citations.contains { $0.uri?.isEmpty ?? false })
     XCTAssertFalse(citations.contains { $0.title?.isEmpty ?? false })