Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading