{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "# Copyright 2025 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": {}, "source": [ "# Format LLM output using an output schema\n", "\n", "\n", "\n", " \n", " \n", " \n", "
\n", " \n", " \"Colab Run in Colab\n", " \n", " \n", " \n", " \"GitHub\n", " View on GitHub\n", " \n", " \n", " \n", " \"BQ\n", " Open in BigQuery Studio\n", " \n", "
\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This notebook shows you how to create structured LLM output by specifying an output schema when generating predictions with a Gemini model." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Costs\n", "\n", "This tutorial uses billable components of Google Cloud:\n", "\n", "* BigQuery (compute)\n", "* BigQuery ML\n", "* Generative AI support on Vertex AI\n", "\n", "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models), [Generative AI support on Vertex AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing),\n", "and [BigQuery ML pricing](https://cloud.google.com/bigquery/pricing#section-11),\n", "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n", "to generate a cost estimate based on your projected usage." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Before you begin\n", "\n", "Complete the tasks in this section to set up your environment." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Set up your Google Cloud project\n", "\n", "**The following steps are required, regardless of your notebook environment.**\n", "\n", "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 credit towards your compute/storage costs.\n", "\n", "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", "\n", "3. [Click here](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com,bigqueryconnection.googleapis.com,aiplatform.googleapis.com) to enable the following APIs:\n", "\n", " * BigQuery API\n", " * BigQuery Connection API\n", " * Vertex AI API\n", "\n", "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Authenticate your Google Cloud account\n", "\n", "Depending on your Jupyter environment, you might have to manually authenticate. Follow the relevant instructions below.\n", "\n", "**BigQuery Studio** or **Vertex AI Workbench**\n", "\n", "Do nothing, you are already authenticated.\n", "\n", "**Local JupyterLab instance**\n", "\n", "Uncomment and run the following cell:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ! gcloud auth login" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Colab**\n", "\n", "Uncomment and run the following cell:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# from google.colab import auth\n", "# auth.authenticate_user()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Set up your project" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Set your project and import necessary modules. If you don't know your project ID, see [Locate the project ID](https://support.google.com/googleapi/answer/7014113)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "PROJECT = \"\" # replace with your project\n", "import bigframes\n", "bigframes.options.bigquery.project = PROJECT\n", "bigframes.options.display.progress_bar = None\n", "\n", "import bigframes.pandas as bpd\n", "from bigframes.ml import llm" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Create a DataFrame and a Gemini model\n", "Create a simple [DataFrame](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.dataframe.DataFrame) of several cities:" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/global_session.py:103: DefaultLocationWarning: No explicit location is set, so using location US for the session.\n", " _global_session = bigframes.session.connect(\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
city
0Seattle
1New York
2Shanghai
\n", "

3 rows × 1 columns

\n", "
[3 rows x 1 columns in total]" ], "text/plain": [ " city\n", "0 Seattle\n", "1 New York\n", "2 Shanghai\n", "\n", "[3 rows x 1 columns]" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "df = bpd.DataFrame({\"city\": [\"Seattle\", \"New York\", \"Shanghai\"]})\n", "df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Connect to a Gemini model using the [`GeminiTextGenerator` class](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator):" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/log_adapter.py:175: FutureWarning: Since upgrading the default model can cause unintended breakages, the\n", "default model will be removed in BigFrames 3.0. Please supply an\n", "explicit model to avoid this message.\n", " return method(*args, **kwargs)\n" ] } ], "source": [ "gemini = llm.GeminiTextGenerator()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Generate structured output data\n", "Previously, LLMs could only generate text output. For example, you could generate output that identifies whether a given city is a US city:" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/array_value.py:109: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n", "`db_dtypes` is a preview feature and subject to change.\n", " warnings.warn(msg, bfe.PreviewWarning)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
cityml_generate_text_llm_result
0SeattleYes, Seattle is a city in the United States. I...
1New YorkYes, New York City is a city in the United Sta...
2ShanghaiNo, Shanghai is not a US city. It is a major c...
\n", "

3 rows × 2 columns

\n", "
[3 rows x 2 columns in total]" ], "text/plain": [ " city ml_generate_text_llm_result\n", "0 Seattle Yes, Seattle is a city in the United States. I...\n", "1 New York Yes, New York City is a city in the United Sta...\n", "2 Shanghai No, Shanghai is not a US city. It is a major c...\n", "\n", "[3 rows x 2 columns]" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = gemini.predict(df, prompt=[df[\"city\"], \"is a US city?\"])\n", "result[[\"city\", \"ml_generate_text_llm_result\"]]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The output is text that a human can read. However, if you want the output to be more useful for analysis, it is better to format the output as structured data. This is especially true when you want to have Boolean, integer, or float values to work with instead of string values. Previously, formatting the output in this way wasn't easy.\n", "\n", "Now, you can get structured output out-of-the-box by specifying the `output_schema` parameter when calling the Gemini model's [`predict` method](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.llm.GeminiTextGenerator#bigframes_ml_llm_GeminiTextGenerator_predict). In the following example, the model output is formatted as Boolean values:" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/array_value.py:109: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n", "`db_dtypes` is a preview feature and subject to change.\n", " warnings.warn(msg, bfe.PreviewWarning)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
cityis_us_city
0SeattleTrue
1New YorkTrue
2ShanghaiFalse
\n", "

3 rows × 2 columns

\n", "
[3 rows x 2 columns in total]" ], "text/plain": [ " city is_us_city\n", "0 Seattle True\n", "1 New York True\n", "2 Shanghai False\n", "\n", "[3 rows x 2 columns]" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = gemini.predict(df, prompt=[df[\"city\"], \"is a US city?\"], output_schema={\"is_us_city\": \"bool\"})\n", "result[[\"city\", \"is_us_city\"]]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can also format model output as float or integer values. In the following example, the model output is formatted as float values to show the city's population in millions:" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/array_value.py:109: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n", "`db_dtypes` is a preview feature and subject to change.\n", " warnings.warn(msg, bfe.PreviewWarning)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
citypopulation_in_millions
0Seattle0.75
1New York19.68
2Shanghai26.32
\n", "

3 rows × 2 columns

\n", "
[3 rows x 2 columns in total]" ], "text/plain": [ " city population_in_millions\n", "0 Seattle 0.75\n", "1 New York 19.68\n", "2 Shanghai 26.32\n", "\n", "[3 rows x 2 columns]" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = gemini.predict(df, prompt=[\"what is the population in millions of\", df[\"city\"]], output_schema={\"population_in_millions\": \"float64\"})\n", "result[[\"city\", \"population_in_millions\"]]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In the following example, the model output is formatted as integer values to show the count of the city's rainy days:" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/array_value.py:109: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n", "`db_dtypes` is a preview feature and subject to change.\n", " warnings.warn(msg, bfe.PreviewWarning)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
cityrainy_days
0Seattle152
1New York123
2Shanghai123
\n", "

3 rows × 2 columns

\n", "
[3 rows x 2 columns in total]" ], "text/plain": [ " city rainy_days\n", "0 Seattle 152\n", "1 New York 123\n", "2 Shanghai 123\n", "\n", "[3 rows x 2 columns]" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = gemini.predict(df, prompt=[\"how many rainy days per year in\", df[\"city\"]], output_schema={\"rainy_days\": \"int64\"})\n", "result[[\"city\", \"rainy_days\"]]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Format output as multiple data types in one prediction\n", "Within a single prediction, you can generate multiple columns of output that use different data types. \n", "\n", "The input doesn't have to be dedicated prompts as long as the output column names are informative to the model." ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/array_value.py:109: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n", "`db_dtypes` is a preview feature and subject to change.\n", " warnings.warn(msg, bfe.PreviewWarning)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
cityis_US_citypopulation_in_millionsrainy_days_per_year
0SeattleTrue0.75152
1New YorkTrue8.8121
2ShanghaiFalse26.32115
\n", "

3 rows × 4 columns

\n", "
[3 rows x 4 columns in total]" ], "text/plain": [ " city is_US_city population_in_millions rainy_days_per_year\n", "0 Seattle True 0.75 152\n", "1 New York True 8.8 121\n", "2 Shanghai False 26.32 115\n", "\n", "[3 rows x 4 columns]" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = gemini.predict(df, prompt=[df[\"city\"]], output_schema={\"is_US_city\": \"bool\", \"population_in_millions\": \"float64\", \"rainy_days_per_year\": \"int64\"})\n", "result[[\"city\", \"is_US_city\", \"population_in_millions\", \"rainy_days_per_year\"]]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Format output as a composite data type" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can generate composite data types like arrays and structs. The following example generates a `places_to_visit` column as an array of strings and a `gps_coordinates` column as a struct of floats:" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/usr/local/google/home/garrettwu/src/bigframes/bigframes/core/array_value.py:109: PreviewWarning: JSON column interpretation as a custom PyArrow extention in\n", "`db_dtypes` is a preview feature and subject to change.\n", " warnings.warn(msg, bfe.PreviewWarning)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
cityis_US_citypopulation_in_millionsrainy_days_per_yearplaces_to_visitgps_coordinates
0SeattleTrue0.74150['Space Needle' 'Pike Place Market' 'Museum of...{'latitude': 47.6062, 'longitude': -122.3321}
1New YorkTrue8.4121['Times Square' 'Central Park' 'Statue of Libe...{'latitude': 40.7128, 'longitude': -74.006}
2ShanghaiFalse26.32115['The Bund' 'Yu Garden' 'Shanghai Museum' 'Ori...{'latitude': 31.2304, 'longitude': 121.4737}
\n", "

3 rows × 6 columns

\n", "
[3 rows x 6 columns in total]" ], "text/plain": [ " city is_US_city population_in_millions rainy_days_per_year \\\n", "0 Seattle True 0.74 150 \n", "1 New York True 8.4 121 \n", "2 Shanghai False 26.32 115 \n", "\n", " places_to_visit \\\n", "0 ['Space Needle' 'Pike Place Market' 'Museum of... \n", "1 ['Times Square' 'Central Park' 'Statue of Libe... \n", "2 ['The Bund' 'Yu Garden' 'Shanghai Museum' 'Ori... \n", "\n", " gps_coordinates \n", "0 {'latitude': 47.6062, 'longitude': -122.3321} \n", "1 {'latitude': 40.7128, 'longitude': -74.006} \n", "2 {'latitude': 31.2304, 'longitude': 121.4737} \n", "\n", "[3 rows x 6 columns]" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "result = gemini.predict(df, prompt=[df[\"city\"]], output_schema={\"is_US_city\": \"bool\", \"population_in_millions\": \"float64\", \"rainy_days_per_year\": \"int64\", \"places_to_visit\": \"array\", \"gps_coordinates\": \"struct\"})\n", "result[[\"city\", \"is_US_city\", \"population_in_millions\", \"rainy_days_per_year\", \"places_to_visit\", \"gps_coordinates\"]]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Clean up\n", "\n", "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", "\n", "Otherwise, run the following cell to delete the temporary cloud artifacts created during the BigFrames session:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "bpd.close_session()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Next steps\n", "\n", "Learn more about BigQuery DataFrames in the [documentation](https://cloud.google.com/python/docs/reference/bigframes/latest) and find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks)." ] } ], "metadata": { "kernelspec": { "display_name": "venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.14" } }, "nbformat": 4, "nbformat_minor": 2 }