Skip to content

Commit 22f9fe3

Browse files
copybara-service[bot]jaycee-li
copybara-service[bot]
and
jaycee-li
authored
feat: [vertexai] add util class SchemaMaker to create Schema from JsonString or JsonObject (#11118)
PiperOrigin-RevId: 670742083 Co-authored-by: Jaycee Li <jayceeli@google.com>
1 parent aa3c6f6 commit 22f9fe3

File tree

2 files changed

+234
-0
lines changed

2 files changed

+234
-0
lines changed

‎java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/generativeai/SchemaMaker.java

+58
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* Copyright 2024 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+
* https://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+
package com.google.cloud.vertexai.generativeai;
17+
18+
import static com.google.common.base.Preconditions.checkArgument;
19+
import static com.google.common.base.Preconditions.checkNotNull;
20+
21+
import com.google.cloud.vertexai.api.Schema;
22+
import com.google.common.base.Strings;
23+
import com.google.gson.JsonObject;
24+
import com.google.protobuf.InvalidProtocolBufferException;
25+
import com.google.protobuf.util.JsonFormat;
26+
27+
/** Helper class to create {@link com.google.cloud.vertexai.api.Schema} */
28+
public final class SchemaMaker {
29+
30+
private SchemaMaker() {}
31+
32+
/**
33+
* Creates a Schema from a JsonString
34+
*
35+
* @param jsonString A valid Json String that can be parsed to a Schema object.
36+
* @return a {@link com.google.cloud.vertexai.api.Schema} by parsing the input JSON string.
37+
* @throws InvalidProtocolBufferException if the String can't be parsed into a Schema proto.
38+
*/
39+
public static Schema fromJsonString(String jsonString) throws InvalidProtocolBufferException {
40+
checkArgument(!Strings.isNullOrEmpty(jsonString), "Input String can't be null or empty.");
41+
Schema.Builder builder = Schema.newBuilder();
42+
JsonFormat.parser().merge(jsonString, builder);
43+
44+
return builder.build();
45+
}
46+
47+
/**
48+
* Creates a Schema from a JsonObject
49+
*
50+
* @param jsonObject A valid JSON object that can be parsed to a Schema object.
51+
* @return a {@link com.google.cloud.vertexai.api.Schema} by parsing the input JSON object.
52+
* @throws InvalidProtocolBufferException if the jsonObject can't be parsed into a Schema proto.
53+
*/
54+
public static Schema fromJsonObject(JsonObject jsonObject) throws InvalidProtocolBufferException {
55+
checkNotNull(jsonObject, "JsonObject can't be null.");
56+
return fromJsonString(jsonObject.toString());
57+
}
58+
}

‎java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/generativeai/SchemaMakerTest.java

+176
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
/*
2+
* Copyright 2024 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+
* https://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+
package com.google.cloud.vertexai.generativeai;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
import static org.junit.Assert.assertThrows;
20+
21+
import com.google.cloud.vertexai.api.Schema;
22+
import com.google.cloud.vertexai.api.Type;
23+
import com.google.gson.JsonObject;
24+
import com.google.protobuf.InvalidProtocolBufferException;
25+
import org.junit.Test;
26+
import org.junit.runner.RunWith;
27+
import org.junit.runners.JUnit4;
28+
29+
@RunWith(JUnit4.class)
30+
public final class SchemaMakerTest {
31+
private static final String STRING_PARAM_NAME = "stringParam";
32+
private static final String INTEGER_PARAM_NAME = "integerParam";
33+
private static final String DOUBLE_PARAM_NAME = "doubleParam";
34+
private static final String FLOAT_PARAM_NAME = "floatParam";
35+
private static final String BOOLEAN_PARAM_NAME = "booleanParam";
36+
37+
private static final Schema EXPECTED_SCHEMA =
38+
Schema.newBuilder()
39+
.setType(Type.OBJECT)
40+
.putProperties(
41+
STRING_PARAM_NAME,
42+
Schema.newBuilder().setType(Type.STRING).setDescription(STRING_PARAM_NAME).build())
43+
.putProperties(
44+
INTEGER_PARAM_NAME,
45+
Schema.newBuilder().setType(Type.INTEGER).setDescription(INTEGER_PARAM_NAME).build())
46+
.putProperties(
47+
DOUBLE_PARAM_NAME,
48+
Schema.newBuilder().setType(Type.NUMBER).setDescription(DOUBLE_PARAM_NAME).build())
49+
.putProperties(
50+
FLOAT_PARAM_NAME,
51+
Schema.newBuilder().setType(Type.NUMBER).setDescription(FLOAT_PARAM_NAME).build())
52+
.putProperties(
53+
BOOLEAN_PARAM_NAME,
54+
Schema.newBuilder().setType(Type.BOOLEAN).setDescription(BOOLEAN_PARAM_NAME).build())
55+
.build();
56+
57+
@Test
58+
public void fromJsonString_withValidJsonString_returnsSchema()
59+
throws InvalidProtocolBufferException {
60+
String jsonString =
61+
"{\n"
62+
+ " \"type\": \"OBJECT\",\n"
63+
+ " \"properties\": {\n"
64+
+ " \"stringParam\": {\n"
65+
+ " \"type\": \"STRING\",\n"
66+
+ " \"description\": \"stringParam\"\n"
67+
+ " },\n"
68+
+ " \"integerParam\": {\n"
69+
+ " \"type\": \"INTEGER\",\n"
70+
+ " \"description\": \"integerParam\"\n"
71+
+ " },\n"
72+
+ " \"doubleParam\": {\n"
73+
+ " \"type\": \"NUMBER\",\n"
74+
+ " \"description\": \"doubleParam\"\n"
75+
+ " },\n"
76+
+ " \"floatParam\": {\n"
77+
+ " \"type\": \"NUMBER\",\n"
78+
+ " \"description\": \"floatParam\"\n"
79+
+ " },\n"
80+
+ " \"booleanParam\": {\n"
81+
+ " \"type\": \"BOOLEAN\",\n"
82+
+ " \"description\": \"booleanParam\"\n"
83+
+ " }\n"
84+
+ " }\n"
85+
+ "}";
86+
87+
Schema schema = SchemaMaker.fromJsonString(jsonString);
88+
89+
assertThat(schema).isEqualTo(EXPECTED_SCHEMA);
90+
}
91+
92+
@Test
93+
public void fromJsonString_withInvalidJsonString_throwsInvalidProtocolBufferException()
94+
throws InvalidProtocolBufferException {
95+
// Here we use "property" (singular) instead of "properties"
96+
String jsonString =
97+
"{\n"
98+
+ " \"type\": \"OBJECT\",\n"
99+
+ " \"property\": {\n"
100+
+ " \"stringParam\": {\n"
101+
+ " \"type\": \"STRING\",\n"
102+
+ " \"description\": \"stringParam\"\n"
103+
+ " }\n"
104+
+ " }\n"
105+
+ "}";
106+
assertThrows(
107+
InvalidProtocolBufferException.class, () -> SchemaMaker.fromJsonString(jsonString));
108+
}
109+
110+
@Test
111+
public void fromJsonString_withEmptyString_throwsIllegalArgumentException()
112+
throws InvalidProtocolBufferException {
113+
String jsonString = "";
114+
IllegalArgumentException thrown =
115+
assertThrows(IllegalArgumentException.class, () -> SchemaMaker.fromJsonString(jsonString));
116+
assertThat(thrown).hasMessageThat().isEqualTo("Input String can't be null or empty.");
117+
}
118+
119+
@Test
120+
public void fromJsonObject_withValidJsonObject_returnsSchema()
121+
throws InvalidProtocolBufferException {
122+
JsonObject stringParamJsonObject = new JsonObject();
123+
stringParamJsonObject.addProperty("type", "STRING");
124+
stringParamJsonObject.addProperty("description", STRING_PARAM_NAME);
125+
126+
JsonObject integerParamJsonObject = new JsonObject();
127+
integerParamJsonObject.addProperty("type", "INTEGER");
128+
integerParamJsonObject.addProperty("description", INTEGER_PARAM_NAME);
129+
130+
JsonObject doubleParamJsonObject = new JsonObject();
131+
doubleParamJsonObject.addProperty("type", "NUMBER");
132+
doubleParamJsonObject.addProperty("description", DOUBLE_PARAM_NAME);
133+
134+
JsonObject floatParamJsonObject = new JsonObject();
135+
floatParamJsonObject.addProperty("type", "NUMBER");
136+
floatParamJsonObject.addProperty("description", FLOAT_PARAM_NAME);
137+
138+
JsonObject booleanParamJsonObject = new JsonObject();
139+
booleanParamJsonObject.addProperty("type", "BOOLEAN");
140+
booleanParamJsonObject.addProperty("description", BOOLEAN_PARAM_NAME);
141+
142+
JsonObject propertiesJsonObject = new JsonObject();
143+
propertiesJsonObject.add(STRING_PARAM_NAME, stringParamJsonObject);
144+
propertiesJsonObject.add(INTEGER_PARAM_NAME, integerParamJsonObject);
145+
propertiesJsonObject.add(DOUBLE_PARAM_NAME, doubleParamJsonObject);
146+
propertiesJsonObject.add(FLOAT_PARAM_NAME, floatParamJsonObject);
147+
propertiesJsonObject.add(BOOLEAN_PARAM_NAME, booleanParamJsonObject);
148+
149+
JsonObject jsonObject = new JsonObject();
150+
jsonObject.addProperty("type", "OBJECT");
151+
jsonObject.add("properties", propertiesJsonObject);
152+
153+
Schema schema = SchemaMaker.fromJsonObject(jsonObject);
154+
155+
assertThat(schema).isEqualTo(EXPECTED_SCHEMA);
156+
}
157+
158+
@Test
159+
public void fromJsonObject_withInvalidJsonObject_throwsInvalidProtocolBufferException()
160+
throws InvalidProtocolBufferException {
161+
// Here we use "property" (singular) instead of "properties"
162+
JsonObject jsonObject = new JsonObject();
163+
jsonObject.addProperty("type", "OBJECT");
164+
jsonObject.add("property", new JsonObject());
165+
166+
assertThrows(
167+
InvalidProtocolBufferException.class, () -> SchemaMaker.fromJsonObject(jsonObject));
168+
}
169+
170+
@Test
171+
public void fromJsonObject_withNullJsonObject_throwsNullPointerException()
172+
throws InvalidProtocolBufferException {
173+
JsonObject jsonObject = null;
174+
assertThrows(NullPointerException.class, () -> SchemaMaker.fromJsonObject(jsonObject));
175+
}
176+
}

0 commit comments

Comments
 (0)