Skip to content

Commit aecf49e

Browse files
cfloresscfloress
authored
feat(genai): Add samples for controlled generation (#10206)
* genai: added controlled generation samples * genai: PR comments * genai: PR refactor * genai: lint fixed
1 parent 900f7fe commit aecf49e

4 files changed

Lines changed: 362 additions & 0 deletions

File tree

‎genai/snippets/src/main/java/genai/controlledgeneration/ControlledGenerationWithEnumClassSchema.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
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 genai.controlledgeneration;
18+
19+
// [START googlegenaisdk_ctrlgen_with_enum_class_schema]
20+
21+
import com.google.genai.Client;
22+
import com.google.genai.types.GenerateContentConfig;
23+
import com.google.genai.types.GenerateContentResponse;
24+
import com.google.genai.types.HttpOptions;
25+
import com.google.genai.types.Schema;
26+
import com.google.genai.types.Type;
27+
import java.util.List;
28+
29+
public class ControlledGenerationWithEnumClassSchema {
30+
31+
public static void main(String[] args) {
32+
// TODO(developer): Replace these variables before running the sample.
33+
String modelId = "gemini-2.5-flash";
34+
String prompt = "What type of instrument is a guitar?";
35+
generateContent(modelId, prompt);
36+
}
37+
38+
// Generates content with an enum class response schema
39+
public static String generateContent(String modelId, String contents) {
40+
// Initialize client that will be used to send requests. This client only needs to be created
41+
// once, and can be reused for multiple requests.
42+
try (Client client =
43+
Client.builder()
44+
.location("global")
45+
.vertexAI(true)
46+
.httpOptions(HttpOptions.builder().apiVersion("v1").build())
47+
.build()) {
48+
49+
// Build schema using enum values
50+
Schema responseSchema =
51+
Schema.builder()
52+
.type(Type.Known.STRING)
53+
.enum_(
54+
List.of(
55+
InstrumentClass.PERCUSSION.getValue(),
56+
InstrumentClass.STRING.getValue(),
57+
InstrumentClass.WOODWIND.getValue(),
58+
InstrumentClass.BRASS.getValue(),
59+
InstrumentClass.KEYBOARD.getValue()))
60+
.build();
61+
62+
GenerateContentConfig config =
63+
GenerateContentConfig.builder()
64+
.responseMimeType("text/x.enum")
65+
.responseSchema(responseSchema)
66+
.build();
67+
68+
GenerateContentResponse response = client.models.generateContent(modelId, contents, config);
69+
70+
System.out.println(response.text());
71+
// Example response:
72+
// String
73+
return response.text();
74+
}
75+
}
76+
77+
// Enum mirroring the Python sample
78+
public enum InstrumentClass {
79+
PERCUSSION("Percussion"),
80+
STRING("String"),
81+
WOODWIND("Woodwind"),
82+
BRASS("Brass"),
83+
KEYBOARD("Keyboard");
84+
85+
private final String value;
86+
87+
InstrumentClass(String value) {
88+
this.value = value;
89+
}
90+
91+
public String getValue() {
92+
return value;
93+
}
94+
}
95+
}
96+
// [END googlegenaisdk_ctrlgen_with_enum_class_schema]

‎genai/snippets/src/main/java/genai/controlledgeneration/ControlledGenerationWithNullableSchema.java‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
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 genai.controlledgeneration;
18+
19+
// [START googlegenaisdk_ctrlgen_with_nullable_schema]
20+
21+
import com.google.genai.Client;
22+
import com.google.genai.types.GenerateContentConfig;
23+
import com.google.genai.types.GenerateContentResponse;
24+
import com.google.genai.types.HttpOptions;
25+
import com.google.genai.types.Schema;
26+
import com.google.genai.types.Type;
27+
import java.util.List;
28+
import java.util.Map;
29+
30+
public class ControlledGenerationWithNullableSchema {
31+
32+
public static void main(String[] args) {
33+
// TODO(developer): Replace these variables before running the sample.
34+
String modelId = "gemini-2.5-flash";
35+
36+
String prompt =
37+
"The week ahead brings a mix of weather conditions.\n"
38+
+ "Sunday is expected to be sunny with a temperature "
39+
+ "of 77°F and a humidity level of 50%. "
40+
+ "Winds will be light at around 10 km/h.\n"
41+
+ "Monday will see partly cloudy skies with "
42+
+ "a slightly cooler temperature of 72°F and the winds "
43+
+ "will pick up slightly to around 15 km/h.\n"
44+
+ "Tuesday brings rain showers, with temperatures dropping "
45+
+ "to 64°F and humidity rising to 70%.\n"
46+
+ "Wednesday may see thunderstorms, with a temperature of 68°F.\n"
47+
+ "Thursday will be cloudy with a temperature of 66°F and moderate humidity at 60%.\n"
48+
+ "Friday returns to partly cloudy conditions, with "
49+
+ "a temperature of 73°F and the Winds will be "
50+
+ "light at 12 km/h.\n"
51+
+ "Finally, Saturday rounds off the week with sunny skies, a "
52+
+ "temperature of 80°F, and a humidity "
53+
+ "level of 40%. Winds will be gentle at 8 km/h.\n";
54+
55+
generateContent(modelId, prompt);
56+
}
57+
58+
// Generates content with a nullable response schema
59+
public static String generateContent(String modelId, String contents) {
60+
// Initialize client that will be used to send requests. This client only needs to be created
61+
// once, and can be reused for multiple requests.
62+
try (Client client =
63+
Client.builder()
64+
.location("global")
65+
.vertexAI(true)
66+
.httpOptions(HttpOptions.builder().apiVersion("v1").build())
67+
.build()) {
68+
69+
// Define schema for array items (each weather entry object)
70+
Schema dayForecastSchema =
71+
Schema.builder()
72+
.type(Type.Known.OBJECT)
73+
.properties(
74+
Map.of(
75+
"Day", Schema.builder().type(Type.Known.STRING).nullable(true).build(),
76+
"Forecast", Schema.builder().type(Type.Known.STRING).nullable(true).build(),
77+
"Temperature",
78+
Schema.builder().type(Type.Known.INTEGER).nullable(true).build(),
79+
"Humidity", Schema.builder().type(Type.Known.STRING).nullable(true).build(),
80+
"Wind Speed",
81+
Schema.builder().type(Type.Known.INTEGER).nullable(true).build()))
82+
.required(List.of("Day", "Temperature", "Forecast", "Wind Speed"))
83+
.build();
84+
85+
// Full response schema
86+
Schema responseSchema =
87+
Schema.builder()
88+
.type(Type.Known.OBJECT)
89+
.properties(
90+
Map.of(
91+
"forecast",
92+
Schema.builder().type(Type.Known.ARRAY).items(dayForecastSchema).build()))
93+
.build();
94+
95+
GenerateContentConfig config =
96+
GenerateContentConfig.builder()
97+
.responseMimeType("application/json")
98+
.responseSchema(responseSchema)
99+
.build();
100+
101+
GenerateContentResponse response = client.models.generateContent(modelId, contents, config);
102+
103+
System.out.println(response.text());
104+
// Example response:
105+
// {"forecast": [{"Day": "Sunday", "Forecast": "sunny", "Temperature": 77, "Wind Speed": 10,
106+
// "Humidity": "50%"},
107+
// {"Day": "Monday", "Forecast": "partly cloudy", "Temperature": 72, "Wind Speed": 15},
108+
// {"Day": "Tuesday", "Forecast": "rain showers", "Temperature": 64, "Wind Speed": null,
109+
// "Humidity": "70%"},
110+
// {"Day": "Wednesday", "Forecast": "thunderstorms", "Temperature": 68, "Wind Speed": null},
111+
// {"Day": "Thursday", "Forecast": "cloudy", "Temperature": 66, "Wind Speed": null,
112+
// "Humidity": "60%"},
113+
// {"Day": "Friday", "Forecast": "partly cloudy", "Temperature": 73, "Wind Speed": 12},
114+
// {"Day": "Saturday", "Forecast": "sunny", "Temperature": 80, "Wind Speed": 8, "Humidity":
115+
// "40%"}]}
116+
return response.text();
117+
}
118+
}
119+
}
120+
// [END googlegenaisdk_ctrlgen_with_nullable_schema]

‎genai/snippets/src/main/java/genai/controlledgeneration/ControlledGenerationWithResponseSchema.java‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
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 genai.controlledgeneration;
18+
19+
// [START googlegenaisdk_ctrlgen_with_resp_schema]
20+
21+
import com.google.genai.Client;
22+
import com.google.genai.types.GenerateContentConfig;
23+
import com.google.genai.types.GenerateContentResponse;
24+
import com.google.genai.types.HttpOptions;
25+
import com.google.genai.types.Schema;
26+
import com.google.genai.types.Type;
27+
import java.util.List;
28+
import java.util.Map;
29+
30+
public class ControlledGenerationWithResponseSchema {
31+
32+
public static void main(String[] args) {
33+
// TODO(developer): Replace these variables before running the sample.
34+
String modelId = "gemini-2.5-flash";
35+
36+
String contents = "List a few popular cookie recipes.";
37+
38+
generateContent(modelId, contents);
39+
}
40+
41+
// Generates content with a response schema
42+
public static String generateContent(String modelId, String contents) {
43+
// Initialize client that will be used to send requests. This client only needs to be created
44+
// once, and can be reused for multiple requests.
45+
try (Client client =
46+
Client.builder()
47+
.location("global")
48+
.vertexAI(true)
49+
.httpOptions(HttpOptions.builder().apiVersion("v1").build())
50+
.build()) {
51+
52+
// Schema for each item in array
53+
Schema recipeSchema =
54+
Schema.builder()
55+
.type(Type.Known.OBJECT)
56+
.properties(
57+
Map.of(
58+
"recipe_name", Schema.builder().type(Type.Known.STRING).build(),
59+
"ingredients",
60+
Schema.builder()
61+
.type(Type.Known.ARRAY)
62+
.items(Schema.builder().type(Type.Known.STRING).build())
63+
.build()))
64+
.required(List.of("recipe_name", "ingredients"))
65+
.build();
66+
67+
// Full root schema (array)
68+
Schema responseSchema = Schema.builder().type(Type.Known.ARRAY).items(recipeSchema).build();
69+
70+
GenerateContentConfig config =
71+
GenerateContentConfig.builder()
72+
.responseMimeType("application/json")
73+
.responseSchema(responseSchema)
74+
.build();
75+
76+
GenerateContentResponse response = client.models.generateContent(modelId, contents, config);
77+
78+
System.out.println(response.text());
79+
// Example response:
80+
// [
81+
// {
82+
// "ingredients": [
83+
// "2 1/4 cups all-purpose flour",
84+
// "1 teaspoon baking soda",
85+
// "1 teaspoon salt",
86+
// "1 cup (2 sticks) unsalted butter, softened",
87+
// "3/4 cup granulated sugar",
88+
// "3/4 cup packed brown sugar",
89+
// "1 teaspoon vanilla extract",
90+
// "2 large eggs",
91+
// "2 cups chocolate chips",
92+
// ],
93+
// "recipe_name": "Chocolate Chip Cookies",
94+
// }
95+
// ]
96+
return response.text();
97+
}
98+
}
99+
}
100+
// [END googlegenaisdk_ctrlgen_with_resp_schema]

‎genai/snippets/src/test/java/genai/controlledgeneration/ControlledGenerationIT.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,4 +65,50 @@ public void testControlledGenerationWithEnumSchema() {
6565
String response = ControlledGenerationWithEnumSchema.generateContent(GEMINI_FLASH, prompt);
6666
assertThat(response).isNotEmpty();
6767
}
68+
69+
@Test
70+
public void testControlledGenerationWithEnumClassSchema() {
71+
String prompt = "What type of instrument is a guitar?";
72+
String response = ControlledGenerationWithEnumClassSchema.generateContent(GEMINI_FLASH, prompt);
73+
74+
assertThat(response).isNotEmpty();
75+
assertThat(response).isEqualTo("String");
76+
}
77+
78+
@Test
79+
public void testControlledGenerationWithNullableSchema() {
80+
String prompt =
81+
"The week ahead brings a mix of weather conditions.\n"
82+
+ "Sunday is expected to be sunny with a temperature "
83+
+ "of 77°F and a humidity level of 50%. "
84+
+ "Winds will be light at around 10 km/h.\n"
85+
+ "Monday will see partly cloudy skies with "
86+
+ "a slightly cooler temperature of 72°F and the winds "
87+
+ "will pick up slightly to around 15 km/h.\n"
88+
+ "Tuesday brings rain showers, with temperatures dropping "
89+
+ "to 64°F and humidity rising to 70%.\n"
90+
+ "Wednesday may see thunderstorms, with a temperature of 68°F.\n"
91+
+ "Thursday will be cloudy with a temperature of 66°F and moderate humidity at 60%.\n"
92+
+ "Friday returns to partly cloudy conditions, with "
93+
+ "a temperature of 73°F and the Winds will be "
94+
+ "light at 12 km/h.\n"
95+
+ "Finally, Saturday rounds off the week with sunny skies, a "
96+
+ "temperature of 80°F, and a humidity "
97+
+ "level of 40%. Winds will be gentle at 8 km/h.\n";
98+
99+
String response = ControlledGenerationWithNullableSchema.generateContent(GEMINI_FLASH, prompt);
100+
101+
assertThat(response).isNotEmpty();
102+
assertThat(response).contains("forecast");
103+
}
104+
105+
@Test
106+
public void testControlledGenerationWithResponseSchema() {
107+
String prompt = "List a few popular cookie recipes.";
108+
109+
String response = ControlledGenerationWithResponseSchema.generateContent(GEMINI_FLASH, prompt);
110+
111+
assertThat(response).isNotEmpty();
112+
assertThat(response).contains("recipe_name");
113+
}
68114
}

0 commit comments

Comments
 (0)