Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Vertex AI] Add Citation.publicationDate #13893

Merged
merged 8 commits into from
Oct 15, 2024
2 changes: 2 additions & 0 deletions FirebaseVertexAI/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
26 changes: 26 additions & 0 deletions FirebaseVertexAI/Sources/GenerateContentResponse.swift
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,9 @@ 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.
public let publicationDate: Date?
}

/// A value enumerating possible reasons for a model to terminate a content generation request.
Expand Down Expand Up @@ -363,28 +366,51 @@ 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
) {
do {
publicationDate = try publicationProtoDate.asDate()
} catch let error as ProtoDate.DateConversionError {
throw DecodingError.dataCorrupted(
DecodingError.Context(
codingPath: [CodingKeys.publicationDate],
debugDescription: "Invalid citation publicationDate.",
underlyingError: error
)
)
}
} else {
publicationDate = nil
}
}
}

Expand Down
137 changes: 137 additions & 0 deletions FirebaseVertexAI/Sources/Types/Internal/ProtoDate.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
// 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
)
}

/// Returns a `Date` representation of the `ProtoDate`.
///
/// - Throws: An error of type `DateConversionError` if the `ProtoDate` cannot be represented as
/// a `Date`.
func asDate() throws -> Date {
guard year != nil else {
throw DateConversionError(message: "Missing a year: \(self)")
}
guard month != nil else {
throw DateConversionError(message: "Missing a month: \(self)")
}
guard day != nil else {
throw DateConversionError(message: "Missing a day: \(self)")
}
guard dateComponents.isValidDate, let date = dateComponents.date else {
throw DateConversionError(message: "Invalid date: \(self)")
}
return date
}

struct DateConversionError: Error {
let localizedDescription: String

init(message: String) {
localizedDescription = message
}
}
}

// 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 {
guard year >= 1 && year <= 9999 else {
throw DecodingError.dataCorrupted(
.init(codingPath: [CodingKeys.year], debugDescription: "Invalid year: \(year)")
andrewheard marked this conversation as resolved.
Show resolved Hide resolved
)
}
self.year = year
} else {
year = nil
}

if let month = try container.decodeIfPresent(Int.self, forKey: .month), month != 0 {
guard month >= 1 && month <= 12 else {
throw DecodingError.dataCorrupted(
.init(codingPath: [CodingKeys.month], debugDescription: "Invalid month: \(month)")
)
}
self.month = month
} else {
month = nil
}

if let day = try container.decodeIfPresent(Int.self, forKey: .day), day != 0 {
guard day >= 1 && day <= 31 else {
andrewheard marked this conversation as resolved.
Show resolved Hide resolved
throw DecodingError.dataCorrupted(
.init(codingPath: [CodingKeys.day], debugDescription: "Invalid day: \(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"
))
}
}
}
15 changes: 13 additions & 2 deletions FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,11 @@ final class GenerativeModelTests: XCTestCase {
].sorted()
let testModelResourceName =
"projects/test-project-id/locations/test-location/publishers/google/models/test-model"
let dateFormatter = {
let dateFormatter = DateFormatter()
dateFormatter.dateFormat = "yyyy-MM-dd"
return dateFormatter
}()

var urlSession: URLSession!
var model: GenerativeModel!
Expand Down Expand Up @@ -134,6 +139,7 @@ final class GenerativeModelTests: XCTestCase {
forResource: "unary-success-citations",
withExtension: "json"
)
let expectedPublicationDate = try XCTUnwrap(dateFormatter.date(from: "2019-05-10"))

let response = try await model.generateContent(testPrompt)

Expand All @@ -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)
Expand All @@ -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 {
Expand Down Expand Up @@ -1052,6 +1061,7 @@ final class GenerativeModelTests: XCTestCase {
forResource: "streaming-success-citations",
withExtension: "txt"
)
let expectedPublicationDate = try XCTUnwrap(dateFormatter.date(from: "2014-03-30"))

let stream = try model.generateContentStream("Hi")
var citations = [Citation]()
Expand All @@ -1072,18 +1082,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 })
Expand Down
Loading
Loading