-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: [vertexai] Add FunctionDeclarationMaker (#10592)
PiperOrigin-RevId: 617858145 Co-authored-by: Zhenyi Qi <[email protected]>
- Loading branch information
1 parent
7078fb7
commit a8aa591
Showing
2 changed files
with
240 additions
and
0 deletions.
There are no files selected for viewing
63 changes: 63 additions & 0 deletions
63
...rtexai/src/main/java/com/google/cloud/vertexai/generativeai/FunctionDeclarationMaker.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,63 @@ | ||
/* | ||
* 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 | ||
* | ||
* https://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. | ||
*/ | ||
package com.google.cloud.vertexai.generativeai; | ||
|
||
import static com.google.common.base.Preconditions.checkArgument; | ||
import static com.google.common.base.Preconditions.checkNotNull; | ||
|
||
import com.google.cloud.vertexai.api.FunctionDeclaration; | ||
import com.google.common.base.Strings; | ||
import com.google.gson.JsonObject; | ||
import com.google.protobuf.InvalidProtocolBufferException; | ||
import com.google.protobuf.util.JsonFormat; | ||
|
||
/** Helper class to create {@link com.google.cloud.vertexai.api.FunctionDeclaration} */ | ||
public final class FunctionDeclarationMaker { | ||
|
||
/** | ||
* Creates a FunctionDeclaration from a JsonString | ||
* | ||
* @param jsonString A valid Json String that can be parsed to a FunctionDeclaration. | ||
* @return a {@link FunctionDeclaration} by parsing the input json String. | ||
* @throws InvalidProtocolBufferException if the String can't be parsed into a FunctionDeclaration | ||
* proto. | ||
*/ | ||
public static FunctionDeclaration fromJsonString(String jsonString) | ||
throws InvalidProtocolBufferException { | ||
checkArgument(!Strings.isNullOrEmpty(jsonString), "Input String can't be null or empty."); | ||
FunctionDeclaration.Builder builder = FunctionDeclaration.newBuilder(); | ||
JsonFormat.parser().merge(jsonString, builder); | ||
FunctionDeclaration declaration = builder.build(); | ||
if (declaration.getName().isEmpty()) { | ||
throw new IllegalArgumentException("name field must be present."); | ||
} | ||
return declaration; | ||
} | ||
|
||
/** | ||
* Creates a FunctionDeclaration from a JsonObject | ||
* | ||
* @param jsonObject A valid Json Object that can be parsed to a FunctionDeclaration. | ||
* @return a {@link FunctionDeclaration} by parsing the input json Object. | ||
* @throws InvalidProtocolBufferException if the jsonObject can't be parsed into a | ||
* FunctionDeclaration proto. | ||
*/ | ||
public static FunctionDeclaration fromJsonObject(JsonObject jsonObject) | ||
throws InvalidProtocolBufferException { | ||
checkNotNull(jsonObject, "JsonObject can't be null."); | ||
return fromJsonString(jsonObject.toString()); | ||
} | ||
} |
177 changes: 177 additions & 0 deletions
177
...ai/src/test/java/com/google/cloud/vertexai/generativeai/FunctionDeclarationMakerTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,177 @@ | ||
/* | ||
* 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 | ||
* | ||
* https://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. | ||
*/ | ||
package com.google.cloud.vertexai.generativeai; | ||
|
||
import static com.google.common.truth.Truth.assertThat; | ||
import static org.junit.Assert.assertThrows; | ||
|
||
import com.google.cloud.vertexai.api.FunctionDeclaration; | ||
import com.google.cloud.vertexai.api.Schema; | ||
import com.google.cloud.vertexai.api.Type; | ||
import com.google.gson.JsonObject; | ||
import com.google.protobuf.InvalidProtocolBufferException; | ||
import org.junit.Test; | ||
import org.junit.runner.RunWith; | ||
import org.junit.runners.JUnit4; | ||
|
||
@RunWith(JUnit4.class) | ||
public final class FunctionDeclarationMakerTest { | ||
|
||
@Test | ||
public void fromValidJsonStringTested_returnsFunctionDeclaration() | ||
throws InvalidProtocolBufferException { | ||
String jsonString = | ||
"{\n" | ||
+ " \"name\": \"functionName\",\n" | ||
+ " \"description\": \"functionDescription\",\n" | ||
+ " \"parameters\": {\n" | ||
+ " \"type\": \"OBJECT\", \n" | ||
+ " \"properties\": {\n" | ||
+ " \"param1\": {\n" | ||
+ " \"type\": \"STRING\",\n" | ||
+ " \"description\": \"param1Description\"\n" | ||
+ " }\n" | ||
+ " }\n" | ||
+ " }\n" | ||
+ "}"; | ||
|
||
FunctionDeclaration functionDeclaration = FunctionDeclarationMaker.fromJsonString(jsonString); | ||
|
||
FunctionDeclaration expectedFunctionDeclaration = | ||
FunctionDeclaration.newBuilder() | ||
.setName("functionName") | ||
.setDescription("functionDescription") | ||
.setParameters( | ||
Schema.newBuilder() | ||
.setType(Type.OBJECT) | ||
.putProperties( | ||
"param1", | ||
Schema.newBuilder() | ||
.setType(Type.STRING) | ||
.setDescription("param1Description") | ||
.build())) | ||
.build(); | ||
assertThat(functionDeclaration).isEqualTo(expectedFunctionDeclaration); | ||
} | ||
|
||
@Test | ||
public void fromJsonStringWithInvalidType_throwsInvalidProtocolBufferException() | ||
throws InvalidProtocolBufferException { | ||
// Here we use "parameter" (singular) instead of "parameters" | ||
String jsonString = | ||
"{\n" | ||
+ " \"name\": \"functionName\",\n" | ||
+ " \"parameter\": {\n" | ||
+ " \"type\": \"OBJECT\", \n" | ||
+ " \"properties\": {\n" | ||
+ " \"param1\": {\n" | ||
+ " \"type\": \"STRING\",\n" | ||
+ " \"description\": \"param1Description\"\n" | ||
+ " }\n" | ||
+ " }\n" | ||
+ " }\n" | ||
+ "}"; | ||
assertThrows( | ||
InvalidProtocolBufferException.class, | ||
() -> FunctionDeclarationMaker.fromJsonString(jsonString)); | ||
} | ||
|
||
@Test | ||
public void fromJsonStringNameMissing_throwsIllegalArgumentException() | ||
throws InvalidProtocolBufferException { | ||
String jsonString = | ||
"{\n" | ||
+ " \"description\": \"functionDescription\",\n" | ||
+ " \"parameters\": {\n" | ||
+ " \"type\": \"OBJECT\", \n" | ||
+ " \"properties\": {\n" | ||
+ " \"param1\": {\n" | ||
+ " \"type\": \"STRING\",\n" | ||
+ " \"description\": \"param1Description\"\n" | ||
+ " }\n" | ||
+ " }\n" | ||
+ " }\n" | ||
+ "}"; | ||
|
||
IllegalArgumentException thrown = | ||
assertThrows( | ||
IllegalArgumentException.class, | ||
() -> FunctionDeclarationMaker.fromJsonString(jsonString)); | ||
assertThat(thrown).hasMessageThat().isEqualTo("name field must be present."); | ||
} | ||
|
||
@Test | ||
public void fromEmptyString_throwsIllegalArgumentException() | ||
throws InvalidProtocolBufferException { | ||
String jsonString = ""; | ||
|
||
IllegalArgumentException thrown = | ||
assertThrows( | ||
IllegalArgumentException.class, | ||
() -> FunctionDeclarationMaker.fromJsonString(jsonString)); | ||
assertThat(thrown).hasMessageThat().isEqualTo("Input String can't be null or empty."); | ||
} | ||
|
||
@Test | ||
public void fromJsonStringStringIsNull_throwsIllegalArgumentException() | ||
throws InvalidProtocolBufferException { | ||
String jsonString = null; | ||
|
||
IllegalArgumentException thrown = | ||
assertThrows( | ||
IllegalArgumentException.class, | ||
() -> FunctionDeclarationMaker.fromJsonString(jsonString)); | ||
assertThat(thrown).hasMessageThat().isEqualTo("Input String can't be null or empty."); | ||
} | ||
|
||
@Test | ||
public void fromValidJsonObject_returnsFunctionDeclaration() | ||
throws InvalidProtocolBufferException { | ||
JsonObject param1JsonObject = new JsonObject(); | ||
param1JsonObject.addProperty("type", "STRING"); | ||
param1JsonObject.addProperty("description", "param1Description"); | ||
|
||
JsonObject propertiesJsonObject = new JsonObject(); | ||
propertiesJsonObject.add("param1", param1JsonObject); | ||
|
||
JsonObject parametersJsonObject = new JsonObject(); | ||
parametersJsonObject.addProperty("type", "OBJECT"); | ||
parametersJsonObject.add("properties", propertiesJsonObject); | ||
|
||
JsonObject jsonObject = new JsonObject(); | ||
jsonObject.addProperty("name", "functionName"); | ||
jsonObject.addProperty("description", "functionDescription"); | ||
jsonObject.add("parameters", parametersJsonObject); | ||
|
||
FunctionDeclaration functionDeclaration = FunctionDeclarationMaker.fromJsonObject(jsonObject); | ||
|
||
FunctionDeclaration expectedFunctionDeclaration = | ||
FunctionDeclaration.newBuilder() | ||
.setName("functionName") | ||
.setDescription("functionDescription") | ||
.setParameters( | ||
Schema.newBuilder() | ||
.setType(Type.OBJECT) | ||
.putProperties( | ||
"param1", | ||
Schema.newBuilder() | ||
.setType(Type.STRING) | ||
.setDescription("param1Description") | ||
.build())) | ||
.build(); | ||
assertThat(functionDeclaration).isEqualTo(expectedFunctionDeclaration); | ||
} | ||
} |