From 2b138b8e6d0a3ab5d8d0fffddde9850e0b035eac Mon Sep 17 00:00:00 2001 From: Rodrigo Lazo Paz Date: Thu, 19 Sep 2024 11:39:43 -0400 Subject: [PATCH 1/2] Wrap citations inside CitationMetadata We were previously unwrapping Citations within CitationMetadata, but we've decided to better align with the proto. b/368310789 --- .../vertexai/internal/util/conversions.kt | 12 ++++++---- .../firebase/vertexai/type/Candidate.kt | 23 +++++++++++++------ .../vertexai/StreamingSnapshotTests.kt | 2 +- .../firebase/vertexai/UnarySnapshotTests.kt | 13 +++++++---- 4 files changed, 33 insertions(+), 17 deletions(-) diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt index 5a0cb4cadbe..fd92bb55c80 100644 --- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt +++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt @@ -20,7 +20,6 @@ import android.graphics.Bitmap import android.graphics.BitmapFactory import android.util.Base64 import com.google.firebase.vertexai.common.client.Schema -import com.google.firebase.vertexai.common.server.CitationSources import com.google.firebase.vertexai.common.shared.Blob import com.google.firebase.vertexai.common.shared.FileData import com.google.firebase.vertexai.common.shared.FunctionCall @@ -32,6 +31,7 @@ import com.google.firebase.vertexai.type.BlobPart import com.google.firebase.vertexai.type.BlockReason import com.google.firebase.vertexai.type.BlockThreshold import com.google.firebase.vertexai.type.Candidate +import com.google.firebase.vertexai.type.Citation import com.google.firebase.vertexai.type.CitationMetadata import com.google.firebase.vertexai.type.Content import com.google.firebase.vertexai.type.CountTokensResponse @@ -181,7 +181,7 @@ internal fun JSONObject.toInternal() = Json.decodeFromString(toStrin internal fun com.google.firebase.vertexai.common.server.Candidate.toPublic(): Candidate { val safetyRatings = safetyRatings?.map { it.toPublic() }.orEmpty() - val citations = citationMetadata?.citationSources?.map { it.toPublic() }.orEmpty() + val citations = citationMetadata?.toPublic() val finishReason = finishReason.toPublic() return Candidate( @@ -228,8 +228,12 @@ internal fun com.google.firebase.vertexai.common.shared.Part.toPublic(): Part { } } -internal fun CitationSources.toPublic() = - CitationMetadata(startIndex = startIndex, endIndex = endIndex, uri = uri ?: "", license = license) +internal fun com.google.firebase.vertexai.common.server.CitationSources.toPublic() = + Citation(startIndex = startIndex, endIndex = endIndex, uri = uri, license = license) + + +internal fun com.google.firebase.vertexai.common.server.CitationMetadata.toPublic() = + CitationMetadata(citationSources.map { it.toPublic() }) internal fun com.google.firebase.vertexai.common.server.SafetyRating.toPublic() = SafetyRating( diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt index 18663908ff5..1d4e2bf6ca5 100644 --- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt +++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt @@ -21,7 +21,7 @@ class Candidate internal constructor( val content: Content, val safetyRatings: List, - val citationMetadata: List, + val citationMetadata: CitationMetadata?, val finishReason: FinishReason? ) @@ -37,16 +37,25 @@ internal constructor( ) /** - * Provides citation metadata for sourcing of content provided by the model between a given - * [startIndex] and [endIndex]. + * A collection of source attributions for a piece of content. * - * @property startIndex The beginning of the citation. - * @property endIndex The end of the citation. - * @property uri The URI of the cited work. - * @property license The license under which the cited work is distributed. + * @property citations A list of individual cited sources and the parts of the content to which they apply. */ class CitationMetadata internal constructor( + val citations: List +) + +/** + * Provides citation information for sourcing of content provided by the model between a given + * [startIndex] and [endIndex]. + * + * @property startIndex The inclusive beginning of a sequence in a model response that derives from a cited source. + * @property endIndex The exclusive end of a sequence in a model response that derives from a cited source. + * @property uri A link to the cited source, if available. + * @property license The license the cited source work is distributed under, if specified. + */ +class Citation internal constructor( val startIndex: Int = 0, val endIndex: Int, val uri: String? = null, diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt index 941f7647edd..9c54536a50f 100644 --- a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt +++ b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt @@ -169,7 +169,7 @@ internal class StreamingSnapshotTests { withTimeout(testTimeout) { val responseList = responses.toList() - responseList.any { it.candidates.any { it.citationMetadata.isNotEmpty() } } shouldBe true + responseList.any { it.candidates.any { it.citationMetadata?.citations?.isNotEmpty() ?: false} } shouldBe true } } diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt index 04f47a0949b..a2a12f632d6 100644 --- a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt +++ b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/UnarySnapshotTests.kt @@ -229,7 +229,7 @@ internal class UnarySnapshotTests { val response = model.generateContent("prompt") response.candidates.isEmpty() shouldBe false - response.candidates.first().citationMetadata.size shouldBe 3 + response.candidates.first().citationMetadata?.citations?.size shouldBe 3 } } @@ -240,11 +240,14 @@ internal class UnarySnapshotTests { val response = model.generateContent("prompt") response.candidates.isEmpty() shouldBe false - response.candidates.first().citationMetadata.isEmpty() shouldBe false + response.candidates.first().citationMetadata?.citations?.isEmpty() shouldBe false // Verify the values in the citation source - with(response.candidates.first().citationMetadata.first()) { - license shouldBe null - startIndex shouldBe 0 + val firstCitation = response.candidates.first().citationMetadata?.citations?.first() + if (firstCitation != null) { + with(firstCitation) { + license shouldBe null + startIndex shouldBe 0 + } } } } From 414da2d18e7d318de62cc81108938ed01a56e230 Mon Sep 17 00:00:00 2001 From: Rodrigo Lazo Paz Date: Thu, 19 Sep 2024 12:13:32 -0400 Subject: [PATCH 2/2] Fix formatting --- .../vertexai/internal/util/conversions.kt | 1 - .../google/firebase/vertexai/type/Candidate.kt | 17 +++++++++-------- .../firebase/vertexai/StreamingSnapshotTests.kt | 4 +++- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt index fd92bb55c80..28f77a00324 100644 --- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt +++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/internal/util/conversions.kt @@ -231,7 +231,6 @@ internal fun com.google.firebase.vertexai.common.shared.Part.toPublic(): Part { internal fun com.google.firebase.vertexai.common.server.CitationSources.toPublic() = Citation(startIndex = startIndex, endIndex = endIndex, uri = uri, license = license) - internal fun com.google.firebase.vertexai.common.server.CitationMetadata.toPublic() = CitationMetadata(citationSources.map { it.toPublic() }) diff --git a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt index 1d4e2bf6ca5..1a23885cbe0 100644 --- a/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt +++ b/firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Candidate.kt @@ -39,23 +39,24 @@ internal constructor( /** * A collection of source attributions for a piece of content. * - * @property citations A list of individual cited sources and the parts of the content to which they apply. + * @property citations A list of individual cited sources and the parts of the content to which they + * apply. */ -class CitationMetadata -internal constructor( - val citations: List -) +class CitationMetadata internal constructor(val citations: List) /** * Provides citation information for sourcing of content provided by the model between a given * [startIndex] and [endIndex]. * - * @property startIndex The inclusive beginning of a sequence in a model response that derives from a cited source. - * @property endIndex The exclusive end of a sequence in a model response that derives from a cited source. + * @property startIndex The inclusive beginning of a sequence in a model response that derives from + * a cited source. + * @property endIndex The exclusive end of a sequence in a model response that derives from a cited + * source. * @property uri A link to the cited source, if available. * @property license The license the cited source work is distributed under, if specified. */ -class Citation internal constructor( +class Citation +internal constructor( val startIndex: Int = 0, val endIndex: Int, val uri: String? = null, diff --git a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt index 9c54536a50f..9a6beb057a1 100644 --- a/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt +++ b/firebase-vertexai/src/test/java/com/google/firebase/vertexai/StreamingSnapshotTests.kt @@ -169,7 +169,9 @@ internal class StreamingSnapshotTests { withTimeout(testTimeout) { val responseList = responses.toList() - responseList.any { it.candidates.any { it.citationMetadata?.citations?.isNotEmpty() ?: false} } shouldBe true + responseList.any { + it.candidates.any { it.citationMetadata?.citations?.isNotEmpty() ?: false } + } shouldBe true } }