controlled generation(genai): Add new code samples (#13109)

* controlled generation(genai): Add new code samples

* controlled generation(genai): Add test case

* controlled generation(genai): Update test file

* controlled generation(genai): Update region tag prefix

* controlled generation(genai): fix lint issues and add requirements-test.txt

* controlled generation(genai): reviewer suggestions

Author: msampathkumar

genai/controlled_generation/ctrlgen_with_class_schema.py

@@ -0,0 +1,59 @@
+# Copyright 2025 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.
+
+
+def generate_content() -> str:
+    # [START googlegenaisdk_ctrlgen_with_class_schema]
+    from google import genai
+    from pydantic import BaseModel, TypeAdapter
+
+    class Recipe(BaseModel):
+        recipe_name: str
+        ingredients: list[str]
+
+    client = genai.Client()
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001",
+        contents="List a few popular cookie recipes.",
+        config={
+            "response_mime_type": "application/json",
+            "response_schema": list[Recipe],
+        },
+    )
+    # Use the response as a JSON string.
+    print(response.text)
+
+    # Example output:
+    # [Recipe(recipe_name='Chocolate Chip Cookies', ingredients=['2 1/4 cups all-purpose flour'
+    #   {
+    #     "ingredients": [
+    #       "2 1/4 cups all-purpose flour",
+    #       "1 teaspoon baking soda",
+    #       "1 teaspoon salt",
+    #       "1 cup (2 sticks) unsalted butter, softened",
+    #       "3/4 cup granulated sugar",
+    #       "3/4 cup packed brown sugar",
+    #       "1 teaspoon vanilla extract",
+    #       "2 large eggs",
+    #       "2 cups chocolate chips"
+    #     ],
+    #     "recipe_name": "Classic Chocolate Chip Cookies"
+    #   }, ... ]
+
+    # [END googlegenaisdk_ctrlgen_with_class_schema]
+    return response.text
+
+
+if __name__ == "__main__":
+    generate_content()

genai/controlled_generation/ctrlgen_with_enum_schema.py

@@ -0,0 +1,42 @@
+# Copyright 2025 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.
+
+
+def generate_content() -> str:
+    # [START googlegenaisdk_ctrlgen_with_enum_schema]
+    from google import genai
+
+    client = genai.Client()
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001",
+        contents="What type of instrument is an oboe?",
+        config={
+            "response_mime_type": "text/x.enum",
+            "response_schema": {
+                "type": "STRING",
+                "enum": ["Percussion", "String", "Woodwind", "Brass", "Keyboard"],
+            },
+        },
+    )
+
+    print(response.text)
+    # Example output:
+    # Woodwind
+
+    # [END googlegenaisdk_ctrlgen_with_enum_schema]
+    return response.text
+
+
+if __name__ == "__main__":
+    generate_content()

genai/controlled_generation/ctrlgen_with_nested_class_schema.py

@@ -0,0 +1,54 @@
+# Copyright 2025 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.
+
+
+def generate_content() -> str:
+    # [START googlegenaisdk_ctrlgen_with_nested_class_schema]
+    from google import genai
+
+    import enum
+    from pydantic import BaseModel
+
+    class Grade(enum.Enum):
+        A_PLUS = "a+"
+        A = "a"
+        B = "b"
+        C = "c"
+        D = "d"
+        F = "f"
+
+    class Recipe(BaseModel):
+        recipe_name: str
+        rating: Grade
+
+    client = genai.Client()
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001",
+        contents="List about 10 home-baked cookies and give them grades based on tastiness.",
+        config={
+            "response_mime_type": "application/json",
+            "response_schema": list[Recipe],
+        },
+    )
+
+    print(response.text)
+    # Example output:
+    # [{"rating": "a+", "recipe_name": "Classic Chocolate Chip Cookies"}, ...]
+
+    # [END googlegenaisdk_ctrlgen_with_nested_class_schema]
+    return response.text
+
+
+if __name__ == "__main__":
+    generate_content()

genai/controlled_generation/ctrlgen_with_nullable_schema.py

@@ -0,0 +1,76 @@
+# Copyright 2025 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.
+
+
+def generate_content() -> str:
+    # [START googlegenaisdk_ctrlgen_with_nullable_schema]
+    from google import genai
+
+    response_schema = {
+        "type": "OBJECT",
+        "properties": {
+            "forecast": {
+                "type": "ARRAY",
+                "items": {
+                    "type": "OBJECT",
+                    "properties": {
+                        "Day": {"type": "STRING", "nullable": True},
+                        "Forecast": {"type": "STRING", "nullable": True},
+                        "Temperature": {"type": "INTEGER", "nullable": True},
+                        "Humidity": {"type": "STRING", "nullable": True},
+                        "Wind Speed": {"type": "INTEGER", "nullable": True},
+                    },
+                    "required": ["Day", "Temperature", "Forecast", "Wind Speed"],
+                },
+            }
+        },
+    }
+
+    prompt = """
+        The week ahead brings a mix of weather conditions.
+        Sunday is expected to be sunny with a temperature of 77°F and a humidity level of 50%. Winds will be light at around 10 km/h.
+        Monday will see partly cloudy skies with a slightly cooler temperature of 72°F and the winds will pick up slightly to around 15 km/h.
+        Tuesday brings rain showers, with temperatures dropping to 64°F and humidity rising to 70%.
+        Wednesday may see thunderstorms, with a temperature of 68°F.
+        Thursday will be cloudy with a temperature of 66°F and moderate humidity at 60%.
+        Friday returns to partly cloudy conditions, with a temperature of 73°F and the Winds will be light at 12 km/h.
+        Finally, Saturday rounds off the week with sunny skies, a temperature of 80°F, and a humidity level of 40%. Winds will be gentle at 8 km/h.
+    """
+
+    client = genai.Client()
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001",
+        contents=prompt,
+        config={
+            "response_mime_type": "application/json",
+            "response_schema": response_schema,
+        },
+    )
+
+    print(response.text)
+    # Example output:
+    # {"forecast": [{"Day": "Sunday", "Forecast": "sunny", "Temperature": 77, "Wind Speed": 10, "Humidity": "50%"},
+    #   {"Day": "Monday", "Forecast": "partly cloudy", "Temperature": 72, "Wind Speed": 15},
+    #   {"Day": "Tuesday", "Forecast": "rain showers", "Temperature": 64, "Wind Speed": null, "Humidity": "70%"},
+    #   {"Day": "Wednesday", "Forecast": "thunderstorms", "Temperature": 68, "Wind Speed": null},
+    #   {"Day": "Thursday", "Forecast": "cloudy", "Temperature": 66, "Wind Speed": null, "Humidity": "60%"},
+    #   {"Day": "Friday", "Forecast": "partly cloudy", "Temperature": 73, "Wind Speed": 12},
+    #   {"Day": "Saturday", "Forecast": "sunny", "Temperature": 80, "Wind Speed": 8, "Humidity": "40%"}]}
+
+    # [END googlegenaisdk_ctrlgen_with_nullable_schema]
+    return response.text
+
+
+if __name__ == "__main__":
+    generate_content()

genai/controlled_generation/noxfile_config.py

@@ -0,0 +1,42 @@
+# Copyright 2021 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
+#
+#      http://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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be imported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.12"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+    # If you need to use a specific version of pip,
+    # change pip_version_override to the string representation
+    # of the version number, for example, "20.2.4"
+    "pip_version_override": None,
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}

genai/controlled_generation/requirements-test.txt

@@ -0,0 +1,4 @@
+backoff==2.2.1
+google-api-core==2.19.0
+pytest==8.2.0
+pytest-asyncio==0.23.6

genai/controlled_generation/test_controlled_generation_samples.py

@@ -0,0 +1,45 @@
+# Copyright 2025 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.
+
+#
+# Using Google Cloud Vertex AI to test the code samples.
+#
+
+import os
+
+import ctrlgen_with_class_schema
+import ctrlgen_with_enum_schema
+import ctrlgen_with_nested_class_schema
+import ctrlgen_with_nullable_schema
+
+os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
+os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
+# The project name is included in the CICD pipeline
+# os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
+
+
+def test_ctrlgen_with_class_schema() -> None:
+    assert ctrlgen_with_class_schema.generate_content()
+
+
+def test_ctrlgen_with_enum_schema() -> None:
+    assert ctrlgen_with_enum_schema.generate_content()
+
+
+def test_ctrlgen_with_nested_class_schema() -> None:
+    assert ctrlgen_with_nested_class_schema.generate_content()
+
+
+def test_ctrlgen_with_nullable_schema() -> None:
+    assert ctrlgen_with_nullable_schema.generate_content()