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..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 @@ -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,11 @@ 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..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 @@ -21,7 +21,7 @@ class Candidate internal constructor( val content: Content, val safetyRatings: List, - val citationMetadata: List, + val citationMetadata: CitationMetadata?, val finishReason: FinishReason? ) @@ -37,15 +37,25 @@ internal constructor( ) /** - * Provides citation metadata for sourcing of content provided by the model between a given + * 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. + */ +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 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 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 CitationMetadata +class Citation internal constructor( val startIndex: Int = 0, val endIndex: Int, 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..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.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 + } } } }