{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Intro to Controlled Generation with the Gemini API\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fcontrolled-generation%2Fintro_controlled_generation.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/controlled-generation/intro_controlled_generation.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://goo.gle/3Pyftqr\">\n",
        "      <img width=\"32px\" src=\"https://cdn.qwiklabs.com/assets/gcp_cloud-e3a77215f0b8bfa9b3f611c0d2208c7e8708ed31.svg\" alt=\"Google Cloud logo\"><br> Open in  Cloud Skills Boost\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/controlled-generation/intro_controlled_generation.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| Author |\n",
        "| --- |\n",
        "| [Eric Dong](https://github.com/gericdong) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "The notebook provides a tutorial on using the Gemini API's controlled generation capabilities to generate structured output. This is important for application reliability, as it ensures the model's response is in a predictable format like JSON, which is easier to process than free-form text.\n",
        "\n",
        "Learn more about [structured output](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/control-generated-output).\n",
        "\n",
        "\n",
        "### Objectives\n",
        "\n",
        "In this tutorial, you learn how to use the controlled generation capability in the Gemini API in Vertex AI to generate model responses in a structured data format.\n",
        "\n",
        "You will complete the following tasks:\n",
        "\n",
        "- Generate JSON\n",
        "    - Option 1: Using Pydantic models\n",
        "    - Option 2: Using OpenAPI Schema\n",
        "    - Option 3: Using JSON Schema\n",
        "- Generate Enum values\n",
        "- Apply controlled generation to several use cases"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = \"global\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "09720c707f1c"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e45ea9a28734"
      },
      "outputs": [],
      "source": [
        "from enum import Enum\n",
        "\n",
        "from google import genai\n",
        "from google.genai.types import GenerateContentConfig, Part, SafetySetting\n",
        "from pydantic import BaseModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qsjb0UBvXxMt"
      },
      "source": [
        "### Connect to the generative AI service on Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k3nnPB7QXVsU"
      },
      "outputs": [],
      "source": [
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_5Y_CA5TXrRH"
      },
      "source": [
        "### Supported Models\n",
        "\n",
        "You can use this feature in the [supported models](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/control-generated-output). The following examples use Gemini 2.5 Flash (`gemini-2.5-flash`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81cbb6bd51d8"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.5-flash\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "52aeea15a479"
      },
      "source": [
        "## Generating JSON\n",
        "\n",
        "The Gemini models allow you to define a response schema to specify the structure of a model's output, the field names, and the expected data type for each field. The response schema is specified in the `response_schema` or the `response_json_schema` parameter in `config`, and the model output will strictly follow that schema.\n",
        "\n",
        "You can configure response schema using the following options:\n",
        "\n",
        "- Using Pydantic models with `response_schema`\n",
        "- Using OpenAPI Schema with `response_schema`\n",
        "- Using JSON Schema with `response_json_schema`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IlnXqujOu8v2"
      },
      "source": [
        "### **Option 1**: Using Pydantic models\n",
        "\n",
        "You can provide the schemas as [Pydantic](https://docs.pydantic.dev/) models or a [JSON](https://www.json.org/json-en.html) string and the model will respond as JSON or an [Enum](https://docs.python.org/3/library/enum.html) depending on the value set in `response_mime_type`.\n",
        "\n",
        "When prompting the model to generate the content, pass the schema to the `response_schema` field of the `generation_config`.\n",
        "\n",
        "You also need to specify the model output format in the `response_mime_type` field. Output formats such as `application/json` and `text/x.enum` are supported."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e4402900f4c6"
      },
      "outputs": [],
      "source": [
        "class CountryInfo(BaseModel):\n",
        "    name: str\n",
        "    population: int\n",
        "    capital: str\n",
        "    continent: str\n",
        "    gdp: int\n",
        "    official_language: str\n",
        "    total_area_sq_mi: int\n",
        "\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=\"Fetch key facts for the United States.\",\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_schema=CountryInfo,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fed9413c2c56"
      },
      "source": [
        "You can either parse the response string as JSON, or use the `parsed` field to get the response as an object or dictionary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf75d83da42d"
      },
      "outputs": [],
      "source": [
        "country: CountryInfo = response.parsed\n",
        "print(country.name)\n",
        "print(country.population)\n",
        "print(country.capital)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "766346c046f9"
      },
      "source": [
        "### **Option 2**: Using OpenAPI Schema\n",
        "\n",
        "You also can define a response schema using a subset of the OpenAPI schema. The supported fields are listed below. All other fields are ignored.\n",
        "\n",
        "- `enum`\n",
        "- `items`\n",
        "- `maxItems`\n",
        "- `nullable`\n",
        "- `properties`\n",
        "- `required`\n",
        "\n",
        "By default, fields are optional, meaning the model can populate the fields or skip them. You can set fields as required to force the model to provide a value. If there's insufficient context in the associated input prompt, the model generates responses mainly based on the data it was trained on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "af3fa1fbff4f"
      },
      "outputs": [],
      "source": [
        "response_schema = {\n",
        "    \"type\": \"OBJECT\",\n",
        "    \"properties\": {\n",
        "        \"name\": {\"type\": \"STRING\"},\n",
        "        \"population\": {\"type\": \"INTEGER\"},\n",
        "        \"capital\": {\"type\": \"STRING\"},\n",
        "        \"continent\": {\"type\": \"STRING\"},\n",
        "        \"gdp\": {\"type\": \"INTEGER\"},\n",
        "        \"official_language\": {\"type\": \"STRING\"},\n",
        "        \"total_area_sq_mi\": {\"type\": \"INTEGER\"},\n",
        "    },\n",
        "    \"required\": [\n",
        "        \"name\",\n",
        "        \"population\",\n",
        "        \"capital\",\n",
        "        \"continent\",\n",
        "        \"gdp\",\n",
        "        \"official_language\",\n",
        "        \"total_area_sq_mi\",\n",
        "    ],\n",
        "}\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=\"Fetch key facts for Canada.\",\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_schema=response_schema,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m9HwMC8GvwOa"
      },
      "source": [
        "### **Option 3**: Using JSON Schema\n",
        "\n",
        "[JSON Schema](https://json-schema.org/) is a standard for describing the structure and validation rules for JSON data.\n",
        "\n",
        "Support for the JSON Schema standard is available in preview. You pass the schema to the `response_json_schema` parameter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "smYYAFYzv4s-"
      },
      "outputs": [],
      "source": [
        "# An advanced JSON Schema with conditional logic\n",
        "country_info_schema = {\n",
        "    \"type\": \"object\",\n",
        "    # 1. Define ALL possible properties at the top level.\n",
        "    \"properties\": {\n",
        "        \"country_name\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The official name of the country.\",\n",
        "        },\n",
        "        \"capital\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The capital city of the country.\",\n",
        "        },\n",
        "        \"is_g7_member\": {\n",
        "            \"type\": \"boolean\",\n",
        "            \"description\": \"True if the country is a member of the G7.\",\n",
        "        },\n",
        "        \"government_type\": {\n",
        "            \"type\": \"string\",\n",
        "            \"enum\": [\"Republic\", \"Constitutional Monarchy\"],\n",
        "            \"description\": \"The type of government.\",\n",
        "        },\n",
        "        \"head_of_state\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The name of the current President or head of state (for Republics).\",\n",
        "        },\n",
        "        \"monarch\": {\n",
        "            \"type\": \"string\",\n",
        "            \"description\": \"The name of the current monarch (for Monarchies).\",\n",
        "        },\n",
        "        \"g7_join_year\": {\n",
        "            \"type\": \"integer\",\n",
        "            \"description\": \"The year the country joined the G7 (if applicable).\",\n",
        "        },\n",
        "    },\n",
        "    # 2. List only the properties that are ALWAYS required.\n",
        "    \"required\": [\"country_name\", \"capital\", \"is_g7_member\", \"government_type\"],\n",
        "    # 3. Use allOf to layer on the conditional requirements.\n",
        "    \"allOf\": [\n",
        "        {\n",
        "            \"if\": {\"properties\": {\"government_type\": {\"const\": \"Republic\"}}},\n",
        "            \"then\": {\"required\": [\"head_of_state\"]},\n",
        "        },\n",
        "        {\n",
        "            \"if\": {\n",
        "                \"properties\": {\"government_type\": {\"const\": \"Constitutional Monarchy\"}}\n",
        "            },\n",
        "            \"then\": {\"required\": [\"monarch\"]},\n",
        "        },\n",
        "        {\n",
        "            \"if\": {\"properties\": {\"is_g7_member\": {\"const\": True}}},\n",
        "            \"then\": {\"required\": [\"g7_join_year\"]},\n",
        "        },\n",
        "    ],\n",
        "}\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=\"Fetch key facts for the United Kingdom.\",\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_json_schema=country_info_schema,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q01-752pnNZ1"
      },
      "source": [
        "## Generating Enum Values\n",
        "\n",
        "To force the model to choose from a predefined list of options, you can use an\n",
        "enum. This ensures data consistency by preventing unpredictable responses.\n",
        "\n",
        "To generate an enum, set the `response_mime_type` field to `text/x.enum`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ca2ZIxobnnVY"
      },
      "outputs": [],
      "source": [
        "class InstrumentEnum(Enum):\n",
        "    PERCUSSION = \"Percussion\"\n",
        "    STRING = \"String\"\n",
        "    WOODWIND = \"Woodwind\"\n",
        "    BRASS = \"Brass\"\n",
        "    KEYBOARD = \"Keyboard\"\n",
        "\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=\"What instrument plays multiple notes at once?\",\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"text/x.enum\",\n",
        "        response_schema=InstrumentEnum,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69450c61bc07"
      },
      "source": [
        "## Use Case Examples\n",
        "\n",
        "Controlled generation can be used to ensure that model outputs adhere to a specific structure (e.g., JSON), instruct the model to perform pure multiple choices (e.g., sentiment classification), or follow certain style or guidelines.\n",
        "\n",
        "Let's use controlled generation in the following use cases that require output constraints."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eba9ef4d4b50"
      },
      "source": [
        "### **Example**: Generate game character profile\n",
        "\n",
        "In this example, you instruct the model to create a game character profile with some specific requirements, and constraint the model output to a structured format. This example also demonstrates how to configure the `response_schema` and `response_mime_type` fields in `config` in conjunction with `safety_settings`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1411f729f2f7"
      },
      "outputs": [],
      "source": [
        "response_schema = {\n",
        "    \"type\": \"ARRAY\",\n",
        "    \"items\": {\n",
        "        \"type\": \"OBJECT\",\n",
        "        \"properties\": {\n",
        "            \"name\": {\"type\": \"STRING\"},\n",
        "            \"age\": {\"type\": \"INTEGER\"},\n",
        "            \"occupation\": {\"type\": \"STRING\"},\n",
        "            \"background\": {\"type\": \"STRING\"},\n",
        "            \"playable\": {\"type\": \"BOOLEAN\"},\n",
        "            \"children\": {\n",
        "                \"type\": \"ARRAY\",\n",
        "                \"items\": {\n",
        "                    \"type\": \"OBJECT\",\n",
        "                    \"properties\": {\n",
        "                        \"name\": {\"type\": \"STRING\"},\n",
        "                        \"age\": {\"type\": \"INTEGER\"},\n",
        "                    },\n",
        "                    \"required\": [\"name\", \"age\"],\n",
        "                },\n",
        "            },\n",
        "        },\n",
        "        \"required\": [\"name\", \"age\", \"occupation\", \"children\"],\n",
        "    },\n",
        "}\n",
        "\n",
        "prompt = \"\"\"\n",
        "    Generate a character profile for a video game, including the character's name, age, occupation, background, names of their\n",
        "    three children, and whether they can be controlled by the player.\n",
        "\"\"\"\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=prompt,\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_schema=response_schema,\n",
        "        safety_settings=[\n",
        "            SafetySetting(\n",
        "                category=\"HARM_CATEGORY_DANGEROUS_CONTENT\",\n",
        "                threshold=\"BLOCK_LOW_AND_ABOVE\",\n",
        "            ),\n",
        "            SafetySetting(\n",
        "                category=\"HARM_CATEGORY_HARASSMENT\",\n",
        "                threshold=\"BLOCK_LOW_AND_ABOVE\",\n",
        "            ),\n",
        "            SafetySetting(\n",
        "                category=\"HARM_CATEGORY_HATE_SPEECH\",\n",
        "                threshold=\"BLOCK_LOW_AND_ABOVE\",\n",
        "            ),\n",
        "            SafetySetting(\n",
        "                category=\"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n",
        "                threshold=\"BLOCK_LOW_AND_ABOVE\",\n",
        "            ),\n",
        "        ],\n",
        "    ),\n",
        ")\n",
        "\n",
        "character: dict = response.parsed\n",
        "print(character)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e02769d61054"
      },
      "source": [
        "### **Example**: Extract errors from log data\n",
        "\n",
        "In this example, you use the model to pull out specific error messages from unstructured log data, extract key information, and constraint the model output to a structured format.\n",
        "\n",
        "Some properties are set to nullable so the model can return a null value when it doesn't have enough context to generate a meaningful response.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "007c0394cadc"
      },
      "outputs": [],
      "source": [
        "response_schema = {\n",
        "    \"type\": \"ARRAY\",\n",
        "    \"items\": {\n",
        "        \"type\": \"OBJECT\",\n",
        "        \"properties\": {\n",
        "            \"timestamp\": {\"type\": \"STRING\"},\n",
        "            \"error_code\": {\"type\": \"INTEGER\", \"nullable\": True},\n",
        "            \"error_message\": {\"type\": \"STRING\"},\n",
        "        },\n",
        "        \"required\": [\"timestamp\", \"error_message\", \"error_code\"],\n",
        "    },\n",
        "}\n",
        "\n",
        "prompt = \"\"\"\n",
        "[15:43:28] ERROR: Could not process image upload: Unsupported file format. (Error Code: 308)\n",
        "[15:44:10] INFO: Search index updated successfully.\n",
        "[15:45:02] ERROR: Service dependency unavailable (payment gateway). Retrying... (Error Code: 5522)\n",
        "[15:45:33] ERROR: Application crashed due to out-of-memory exception. (Error Code: 9001)\n",
        "\"\"\"\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=prompt,\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_schema=response_schema,\n",
        "    ),\n",
        ")\n",
        "\n",
        "log_data: dict = response.parsed\n",
        "print(log_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a74594893037"
      },
      "source": [
        "### **Example**: Analyze product review data\n",
        "\n",
        "In this example, you instruct the model to analyze product review data, extract key entities, perform sentiment classification (multiple choices), provide additional explanation, and output the results in JSON format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sG9uSOqnstqS"
      },
      "outputs": [],
      "source": [
        "response_schema = {\n",
        "    \"type\": \"ARRAY\",\n",
        "    \"items\": {\n",
        "        \"type\": \"ARRAY\",\n",
        "        \"items\": {\n",
        "            \"type\": \"OBJECT\",\n",
        "            \"properties\": {\n",
        "                \"rating\": {\"type\": \"INTEGER\"},\n",
        "                \"flavor\": {\"type\": \"STRING\"},\n",
        "                \"sentiment\": {\n",
        "                    \"type\": \"STRING\",\n",
        "                    \"enum\": [\"POSITIVE\", \"NEGATIVE\", \"NEUTRAL\"],\n",
        "                },\n",
        "                \"explanation\": {\"type\": \"STRING\"},\n",
        "            },\n",
        "            \"required\": [\"rating\", \"flavor\", \"sentiment\", \"explanation\"],\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "prompt = \"\"\"\n",
        "  Analyze the following product reviews, output the sentiment classification and give an explanation.\n",
        "\n",
        "  - \"Absolutely loved it! Best ice cream I've ever had.\" Rating: 4, Flavor: Strawberry Cheesecake\n",
        "  - \"Quite good, but a bit too sweet for my taste.\" Rating: 1, Flavor: Mango Tango\n",
        "\"\"\"\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=prompt,\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_schema=response_schema,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10971b23afcf"
      },
      "source": [
        "### **Example**: Detect objects in images\n",
        "\n",
        "You can also use controlled generation in multimodality use cases. In this example, you instruct the model to detect objects in the images and output the results in JSON format. These images are stored in a Google Storage bucket.\n",
        "\n",
        "- [office-desk.jpeg](https://storage.googleapis.com/cloud-samples-data/generative-ai/image/office-desk.jpeg)\n",
        "- [gardening-tools.jpeg](https://storage.googleapis.com/cloud-samples-data/generative-ai/image/gardening-tools.jpeg)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1f3e9935e2da"
      },
      "outputs": [],
      "source": [
        "response_schema = {\n",
        "    \"type\": \"ARRAY\",\n",
        "    \"items\": {\n",
        "        \"type\": \"ARRAY\",\n",
        "        \"items\": {\n",
        "            \"type\": \"OBJECT\",\n",
        "            \"properties\": {\n",
        "                \"object\": {\"type\": \"STRING\"},\n",
        "            },\n",
        "        },\n",
        "    },\n",
        "}\n",
        "\n",
        "prompt = \"Generate a list of objects in the images.\"\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=[\n",
        "        Part.from_uri(\n",
        "            file_uri=\"gs://cloud-samples-data/generative-ai/image/office-desk.jpeg\",\n",
        "            mime_type=\"image/jpeg\",\n",
        "        ),\n",
        "        Part.from_uri(\n",
        "            file_uri=\"gs://cloud-samples-data/generative-ai/image/gardening-tools.jpeg\",\n",
        "            mime_type=\"image/jpeg\",\n",
        "        ),\n",
        "        prompt,\n",
        "    ],\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",\n",
        "        response_schema=response_schema,\n",
        "    ),\n",
        ")\n",
        "\n",
        "object_list = response.parsed\n",
        "print(object_list)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8e47be074e75"
      },
      "source": [
        "### **Example**: Respond with a single plain text enum value\n",
        "\n",
        "This example identifies the genre of a movie based on its description. The output is one plain-text enum value that the model selects from a list of values that are defined in the response schema. Note that in this example, the `response_mime_type` field is set to `text/x.enum`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f4f0a68dda49"
      },
      "outputs": [],
      "source": [
        "response_schema = {\"type\": \"STRING\", \"enum\": [\"drama\", \"comedy\", \"documentary\"]}\n",
        "\n",
        "prompt = (\n",
        "    \"The film aims to educate and inform viewers about real-life subjects, events, or people.\"\n",
        "    \"It offers a factual record of a particular topic by combining interviews, historical footage, \"\n",
        "    \"and narration. The primary purpose of a film is to present information and provide insights \"\n",
        "    \"into various aspects of reality.\"\n",
        ")\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=prompt,\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"text/x.enum\",\n",
        "        response_schema=response_schema,\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08cd5ace6daa"
      },
      "source": [
        "### **Example**: Use a JSON schema with conditional logic\n",
        "\n",
        "This example demonstrates how to use a JSON schema with conditional logic to guide the model to generate highly structured data. The schema defines a format for order status updates:\n",
        "\n",
        "- Every update includes an order_id and a status.\n",
        "- If the status is `SHIPPED`, a tracking_number is required.\n",
        "- If the status is `DELIVERED`, a delivery_date is required.\n",
        "\n",
        "The prompt is a simple request for three different order updates. By using the provided JSON schema in the request, the model is forced to generate a structured JSON array that includes the conditional fields for each order, resulting in a predictable JSON output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "928e77f64ca4"
      },
      "outputs": [],
      "source": [
        "# A JSON schema for order status updates with conditional fields\n",
        "order_schema = {\n",
        "    \"type\": \"object\",\n",
        "    # 1. Define all possible properties that could appear.\n",
        "    \"properties\": {\n",
        "        \"order_id\": {\n",
        "            \"type\": \"string\",\n",
        "        },\n",
        "        \"status\": {\n",
        "            \"type\": \"string\",\n",
        "            \"enum\": [\"PENDING\", \"SHIPPED\", \"DELIVERED\"],\n",
        "        },\n",
        "        \"tracking_number\": {\n",
        "            \"type\": \"string\",\n",
        "        },\n",
        "        \"delivery_date\": {\n",
        "            \"type\": \"string\",\n",
        "            \"format\": \"date\",\n",
        "        },\n",
        "    },\n",
        "    # 2. List only the properties that are ALWAYS required.\n",
        "    \"required\": [\"order_id\", \"status\"],\n",
        "    # 3. Add the conditional logic.\n",
        "    \"allOf\": [\n",
        "        {\n",
        "            \"if\": {\"properties\": {\"status\": {\"const\": \"SHIPPED\"}}},\n",
        "            \"then\": {\"required\": [\"tracking_number\"]},\n",
        "        },\n",
        "        {\n",
        "            \"if\": {\"properties\": {\"status\": {\"const\": \"DELIVERED\"}}},\n",
        "            \"then\": {\"required\": [\"delivery_date\"]},\n",
        "        },\n",
        "    ],\n",
        "}\n",
        "\n",
        "# A prompt asking for updates on three different orders\n",
        "prompt = \"\"\"\n",
        "Provide a status update for the following orders:\n",
        "- Order 98765 is still processing.\n",
        "- Order 54321 was shipped. (Tracking number: TRK54321XYZ)\n",
        "- Order 12345 was delivered on July 31, 2025.\n",
        "\"\"\"\n",
        "\n",
        "# The API call using the JSON schema\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=prompt,\n",
        "    config=GenerateContentConfig(\n",
        "        response_mime_type=\"application/json\",  # required\n",
        "        response_json_schema={\"type\": \"array\", \"items\": order_schema},\n",
        "    ),\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "intro_controlled_generation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
