Skip to content

Commit

Permalink
[Vertex AI] Add Citation.publicationDate (#13893)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard authored Oct 15, 2024
1 parent b99e3d7 commit 4ca1413
Show file tree
Hide file tree
Showing 6 changed files with 401 additions and 2 deletions.
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
24 changes: 24 additions & 0 deletions FirebaseVertexAI/Sources/GenerateContentResponse.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
}
}

Expand Down
117 changes: 117 additions & 0 deletions FirebaseVertexAI/Sources/Types/Internal/ProtoDate.swift
Original file line number Diff line number Diff line change
@@ -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"
))
}
}
}
4 changes: 4 additions & 0 deletions FirebaseVertexAI/Sources/VertexLog.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 18 additions & 2 deletions FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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)

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,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]()
Expand All @@ -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 })
Expand Down
Loading

0 comments on commit 4ca1413

Please sign in to comment.