Skip to content

Commit d32c474

Browse files
rlazodaymxn
authored and
tejasd
committed
[VertexAI] Add initial support to export covered API (#6749)
The server produces a discovery document with the details of the API surface. https://aiplatform.googleapis.com/$discovery/rest?version=v1beta1 This change introduces code that can generate similar a description of the API covered by the SDK. This will enable us to track difference between both. In a follow up PR we can implement the logic to fully export the surface. --------- Co-authored-by: Daymon <[email protected]>
1 parent d5801dc commit d32c474

File tree

3 files changed

+389
-1
lines changed

3 files changed

+389
-1
lines changed

‎firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/common/util/serialization.kt

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import kotlinx.serialization.KSerializer
2323
import kotlinx.serialization.SerialName
2424
import kotlinx.serialization.descriptors.SerialDescriptor
2525
import kotlinx.serialization.descriptors.buildClassSerialDescriptor
26+
import kotlinx.serialization.descriptors.element
2627
import kotlinx.serialization.encoding.Decoder
2728
import kotlinx.serialization.encoding.Encoder
2829

@@ -36,7 +37,12 @@ import kotlinx.serialization.encoding.Encoder
3637
*/
3738
internal class FirstOrdinalSerializer<T : Enum<T>>(private val enumClass: KClass<T>) :
3839
KSerializer<T> {
39-
override val descriptor: SerialDescriptor = buildClassSerialDescriptor("FirstOrdinalSerializer")
40+
override val descriptor: SerialDescriptor =
41+
buildClassSerialDescriptor("FirstOrdinalSerializer") {
42+
for (enumValue in enumClass.enumValues()) {
43+
element<String>(enumValue.toString())
44+
}
45+
}
4046

4147
override fun deserialize(decoder: Decoder): T {
4248
val name = decoder.decodeString()

‎firebase-vertexai/src/test/java/com/google/firebase/vertexai/SerializationTests.kt

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.firebase.vertexai
18+
19+
import com.google.firebase.vertexai.common.util.descriptorToJson
20+
import com.google.firebase.vertexai.type.Candidate
21+
import com.google.firebase.vertexai.type.CountTokensResponse
22+
import com.google.firebase.vertexai.type.GenerateContentResponse
23+
import com.google.firebase.vertexai.type.ModalityTokenCount
24+
import com.google.firebase.vertexai.type.Schema
25+
import io.kotest.assertions.json.shouldEqualJson
26+
import org.junit.Test
27+
28+
internal class SerializationTests {
29+
@Test
30+
fun `test countTokensResponse serialization as Json`() {
31+
val expectedJsonAsString =
32+
"""
33+
{
34+
"id": "CountTokensResponse",
35+
"type": "object",
36+
"properties": {
37+
"totalTokens": {
38+
"type": "integer"
39+
},
40+
"totalBillableCharacters": {
41+
"type": "integer"
42+
},
43+
"promptTokensDetails": {
44+
"type": "array",
45+
"items": {
46+
"${'$'}ref": "ModalityTokenCount"
47+
}
48+
}
49+
}
50+
}
51+
"""
52+
.trimIndent()
53+
val actualJson = descriptorToJson(CountTokensResponse.Internal.serializer().descriptor)
54+
expectedJsonAsString shouldEqualJson actualJson.toString()
55+
}
56+
57+
@Test
58+
fun `test modalityTokenCount serialization as Json`() {
59+
val expectedJsonAsString =
60+
"""
61+
{
62+
"id": "ModalityTokenCount",
63+
"type": "object",
64+
"properties": {
65+
"modality": {
66+
"type": "string",
67+
"enum": [
68+
"UNSPECIFIED",
69+
"TEXT",
70+
"IMAGE",
71+
"VIDEO",
72+
"AUDIO",
73+
"DOCUMENT"
74+
]
75+
},
76+
"tokenCount": {
77+
"type": "integer"
78+
}
79+
}
80+
}
81+
"""
82+
.trimIndent()
83+
val actualJson = descriptorToJson(ModalityTokenCount.Internal.serializer().descriptor)
84+
expectedJsonAsString shouldEqualJson actualJson.toString()
85+
}
86+
87+
@Test
88+
fun `test GenerateContentResponse serialization as Json`() {
89+
val expectedJsonAsString =
90+
"""
91+
{
92+
"id": "GenerateContentResponse",
93+
"type": "object",
94+
"properties": {
95+
"candidates": {
96+
"type": "array",
97+
"items": {
98+
"${'$'}ref": "Candidate"
99+
}
100+
},
101+
"promptFeedback": {
102+
"${'$'}ref": "PromptFeedback"
103+
},
104+
"usageMetadata": {
105+
"${'$'}ref": "UsageMetadata"
106+
}
107+
}
108+
}
109+
"""
110+
.trimIndent()
111+
val actualJson = descriptorToJson(GenerateContentResponse.Internal.serializer().descriptor)
112+
expectedJsonAsString shouldEqualJson actualJson.toString()
113+
}
114+
115+
@Test
116+
fun `test Candidate serialization as Json`() {
117+
val expectedJsonAsString =
118+
"""
119+
{
120+
"id": "Candidate",
121+
"type": "object",
122+
"properties": {
123+
"content": {
124+
"${'$'}ref": "Content"
125+
},
126+
"finishReason": {
127+
"type": "string",
128+
"enum": [
129+
"UNKNOWN",
130+
"UNSPECIFIED",
131+
"STOP",
132+
"MAX_TOKENS",
133+
"SAFETY",
134+
"RECITATION",
135+
"OTHER",
136+
"BLOCKLIST",
137+
"PROHIBITED_CONTENT",
138+
"SPII",
139+
"MALFORMED_FUNCTION_CALL"
140+
]
141+
},
142+
"safetyRatings": {
143+
"type": "array",
144+
"items": {
145+
"${'$'}ref": "SafetyRating"
146+
}
147+
},
148+
"citationMetadata": {
149+
"${'$'}ref": "CitationMetadata"
150+
},
151+
"groundingMetadata": {
152+
"${'$'}ref": "GroundingMetadata"
153+
}
154+
}
155+
}
156+
"""
157+
.trimIndent()
158+
val actualJson = descriptorToJson(Candidate.Internal.serializer().descriptor)
159+
expectedJsonAsString shouldEqualJson actualJson.toString()
160+
}
161+
162+
@Test
163+
fun `test Schema serialization as Json`() {
164+
/**
165+
* Unlike the actual schema in the background, we don't represent "type" as an enum, but rather
166+
* as a string. This is because we restrict what values can be used (using helper methods,
167+
* rather than type).
168+
*/
169+
val expectedJsonAsString =
170+
"""
171+
{
172+
"id": "Schema",
173+
"type": "object",
174+
"properties": {
175+
"type": {
176+
"type": "string"
177+
},
178+
"format": {
179+
"type": "string"
180+
},
181+
"description": {
182+
"type": "string"
183+
},
184+
"nullable": {
185+
"type": "boolean"
186+
},
187+
"items": {
188+
"${'$'}ref": "Schema"
189+
},
190+
"enum": {
191+
"type": "array",
192+
"items": {
193+
"type": "string"
194+
}
195+
},
196+
"properties": {
197+
"type": "object",
198+
"additionalProperties": {
199+
"${'$'}ref": "Schema"
200+
}
201+
},
202+
"required": {
203+
"type": "array",
204+
"items": {
205+
"type": "string"
206+
}
207+
}
208+
}
209+
}
210+
"""
211+
.trimIndent()
212+
val actualJson = descriptorToJson(Schema.Internal.serializer().descriptor)
213+
expectedJsonAsString shouldEqualJson actualJson.toString()
214+
}
215+
}

0 commit comments

Comments
 (0)