language/grounding/intro-grounding.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": [
        "# Getting Started with Grounding in Vertex AI\n",
        "\n",
        "> **NOTE:** This notebook uses the PaLM generative model, which will reach its [discontinuation date in October 2024](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text#model_versions). Please refer to [this updated notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/grounding/intro-grounding-gemini.ipynb) for a version which uses the latest Gemini model.\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/language/grounding/intro-grounding.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/language/grounding/intro-grounding.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\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/language/grounding/intro-grounding.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\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/language/grounding/intro-grounding.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/language/grounding/intro-grounding.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/language/grounding/intro-grounding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/53/X_logo_2023_original.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/language/grounding/intro-grounding.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/language/grounding/intro-grounding.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": "84e7e432e6ff"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Kristopher Overholt](https://github.com/koverholt) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24743cf4a1e1"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.11"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "[Grounding in Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/grounding/ground-language-models) lets you use language models (e.g., [`text-bison` and `chat-bison`](https://cloud.google.com/vertex-ai/docs/generative-ai/language-model-overview)) to generate content grounded in your own documents and data. This capability lets the model access information at runtime that goes beyond its training data. By grounding model responses in Google Search results or data stores within [Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/enterprise-search-introduction), LLMs that are grounded in data can produce more accurate, up-to-date, and relevant responses.\n",
        "\n",
        "Grounding provides the following benefits:\n",
        "\n",
        "- Reduces model hallucinations (instances where the model generates content that isn't factual)\n",
        "- Anchors model responses to specific information, documents, and data sources\n",
        "- Enhances the trustworthiness, accuracy, and applicability of the generated content\n",
        "\n",
        "In the context of grounding in Vertex AI, you can configure two different sources of grounding:\n",
        "\n",
        "1. Google Search results for data that is publicly available and indexed\n",
        "1. [Data stores in Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/create-datastore-ingest), which can include your own data in the form of website data, unstructured data, or structured data\n",
        "\n",
        "**NOTE:** Some of the features in this sample notebook require early access to certain features via an allowlist. [Grounding with Vertex AI Search](https://cloud.google.com/vertex-ai/docs/generative-ai/grounding/ground-language-models) is available in Public Preview, whereas Grounding with Google Web Search results is available in Private Preview. To request early access to features in Private Preview, contact your account representative or [Google Cloud Support](https://cloud.google.com/contact)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to:\n",
        "\n",
        "- Generate LLM text and chat model responses grounded in Google Search results\n",
        "- Compare the results of ungrounded LLM responses with grounded LLM responses\n",
        "- Create and use a data store in Vertex AI Search to ground responses in custom documents and data\n",
        "- Generate LLM text and chat model responses grounded in Vertex AI Search results\n",
        "- Use the asynchronous text and chat models APIs with grounding\n",
        "\n",
        "This tutorial uses the following Google Cloud AI services and resources:\n",
        "\n",
        "- Vertex AI\n",
        "- Vertex AI Search and Conversation\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Configuring the LLM and prompt for various examples\n",
        "- Sending example prompts to generative text and chat models in Vertex AI\n",
        "- Setting up a data store in Vertex AI Search with your own data\n",
        "- Sending example prompts with various levels of grounding (no grounding, web grounding, data store grounding)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### 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 free credit towards your compute/storage costs.\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "1. Enable the [Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com) and [Vertex AI Search and Conversation API](https://console.cloud.google.com/flows/enableapi?apiid=discoveryengine.googleapis.com).\n",
        "1. If you want to use Grounding with Google Web Search results, your project must also be allowlisted for this feature while it is in the Private Preview stage.\n",
        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "### Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform==1.38.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "Restart the kernel after installing packages:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "### Configure your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Updated property [core/project].\n"
          ]
        }
      ],
      "source": [
        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project ID\n",
        "!gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "### Configure your region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "If you are running this notebook on Google Colab, you will need to authenticate your environment. To do this, run the new cell below. This step is not required if you are using Vertex AI Workbench."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    # Authenticate user to Google Cloud\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "from vertexai.language_models import ChatModel, GroundingSource, TextGenerationModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "outputs": [],
      "source": [
        "vertexai.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4e76133e3d70"
      },
      "source": [
        "Initialize the generative text and chat models from Vertex AI:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "274baabc5db6"
      },
      "outputs": [],
      "source": [
        "text_model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
        "chat_model = ChatModel.from_pretrained(\"chat-bison\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e336da7161af"
      },
      "source": [
        "## Example: Grounding with Google Search results\n",
        "\n",
        "In this example, you'll compare LLM responses with no grounding with responses that are grounded in the results of a Google Search. You'll ask a question about a recent hardware release from the Google Store."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "3944de52a061"
      },
      "outputs": [],
      "source": [
        "PROMPT = (\n",
        "    \"What are the price, available colors, and storage size options of a Pixel Tablet?\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "25955ce5d263"
      },
      "source": [
        "### Text generation without grounding\n",
        "\n",
        "Make a prediction request to the LLM with no grounding:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "e8f9084d463c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              " **Price:**\n",
              "\n",
              "* Starting at $399 for the Wi-Fi-only model with 128GB of storage\n",
              "* $499 for the Wi-Fi + 5G model with 128GB of storage\n",
              "* $599 for the Wi-Fi + 5G model with 256GB of storage\n",
              "\n",
              "**Available Colors:**\n",
              "\n",
              "* Chalk (white)\n",
              "* Charcoal (black)\n",
              "* Sage (green)\n",
              "\n",
              "**Storage Size Options:**\n",
              "\n",
              "* 128GB\n",
              "* 256GB"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "response = text_model.predict(PROMPT)\n",
        "response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96af98bfdb56"
      },
      "source": [
        "### Text generation grounded in Google Search results\n",
        "\n",
        "Now you can add the `grounding_source` keyword arg with a grounding source of `GroundingSource.WebSearch()` to instruct the LLM to first perform a Google Search with the prompt, then construct an answer based on the web search results:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "6be1d992a399"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( The Pixel Tablet starts at $499 in the US, £599 in the UK, €679 throughout selected European regions, and CAD $699 in Canada. It comes in three colors: Porcelain (white), Rose (pink), and Hazel (green). The storage size options are 128GB and 256GB.,\n",
              " GroundingMetadata(citations=[GroundingCitation(start_index=1, end_index=129, url='https://www.androidauthority.com/google-pixel-tablet-3163922/', title=None, license=None, publication_date=None), GroundingCitation(start_index=130, end_index=206, url='https://www.androidpolice.com/google-pixel-tablet/', title=None, license=None, publication_date=None)], search_queries=['Pixel Tablet price, colors, and storage size options?']))"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "grounding_source = GroundingSource.WebSearch()\n",
        "\n",
        "response = text_model.predict(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b9a1e2b0a864"
      },
      "source": [
        "Note that the response without grounding only has limited information from the LLM about the Pixel tablet. Whereas the response that was grounded in web search results contains the most up to date information from web search results that are returned as part of the LLM with grounding request."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77f0800f8762"
      },
      "source": [
        "## Example: Grounding with custom documents and data\n",
        "\n",
        "In this example, you'll compare LLM responses with no grounding with responses that are grounded in the [results of a data store in Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/create-datastore-ingest). You'll ask a question about a GoogleSQL query to create an [object table in BigQuery](https://cloud.google.com/bigquery/docs/object-table-introduction)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5c2fc348bed"
      },
      "source": [
        "### Creating a data store in Vertex AI Search\n",
        "\n",
        "Follow the steps in the [Vertex AI Search getting started documentation](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_a_search_app_for_website_data) to create a data store in Vertex AI Search with sample data. In this example, you'll use a website-based data store that contains content from the Google Cloud website, including documentation.\n",
        "\n",
        "Once you've created a data store, obtain the Data Store ID and input it below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "5d09a93b4292"
      },
      "outputs": [],
      "source": [
        "DATA_STORE_ID = \"your-data-store-id_1234567890123\"  # Replace this with your data store ID from Vertex AI Search\n",
        "DATA_STORE_REGION = \"global\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ccc156676e0a"
      },
      "source": [
        "Now you can ask a question about object tables in BigQuery and when to use them:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "9c1e1b1743bd"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"When should I use an object table in BigQuery? And how does it store data?\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f365681544bb"
      },
      "source": [
        "### Text generation without grounding\n",
        "\n",
        "Make a prediction request to the LLM with no grounding:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "29a264acddda"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( **When to use an object table in BigQuery**\n",
              "\n",
              "Object tables are a specialized type of table in BigQuery that is designed for storing and querying semi-structured data. Semi-structured data is data that does not conform to a fixed schema, such as JSON, XML, or Avro.\n",
              "\n",
              "Object tables are useful for storing data that is:\n",
              "\n",
              "* **Complex and hierarchical:** Object tables can store data that is nested or has a complex structure. For example, you could store a JSON object that represents a customer record, which includes the customer's name, address, and order history.\n",
              "* **Changing frequently:** Object,\n",
              " GroundingMetadata(citations=[], search_queries=[]))"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "response = text_model.predict(PROMPT)\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "26a239cfcc37"
      },
      "source": [
        "### Text generation grounded in Vertex AI Search results\n",
        "\n",
        "Now we can add the `grounding_source` keyword arg with a grounding source of `GroundingSource.VertexAISearch()` to instruct the LLM to first perform a search within your custom data store, then construct an answer based on the relevant documents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "3bb350a79e85"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( **When to use an object table in BigQuery**\n",
              "\n",
              "Object tables are useful for storing and analyzing unstructured data, such as images, videos, and audio files. They can also be used to store semi-structured data, such as JSON or XML files.\n",
              "\n",
              "Object tables are particularly useful when you need to:\n",
              "\n",
              "* Store large amounts of unstructured data\n",
              "* Perform complex analysis on unstructured data\n",
              "* Share unstructured data with others\n",
              "* Access unstructured data from multiple locations\n",
              "\n",
              "**How object tables store data**\n",
              "\n",
              "Object tables store data in a columnar format, which makes it efficient to query and analyze large amounts of data. Each column,\n",
              " GroundingMetadata(citations=[], search_queries=['When should I use an object table in BigQuery?']))"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "grounding_source = GroundingSource.VertexAISearch(\n",
        "    data_store_id=DATA_STORE_ID, location=DATA_STORE_REGION\n",
        ")\n",
        "\n",
        "response = text_model.predict(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ee76e67a1b48"
      },
      "source": [
        "Note that the response without grounding only has limited information from the LLM about object tables in BigQuery that might not be accurate. Whereas the response that was grounded in Vertex AI Search results contains the most up to date information from the Google Cloud documentation about BigQuery."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "54562717e2a4"
      },
      "source": [
        "## Example: Grounded chat responses\n",
        "\n",
        "You can also use grounding when working with chat models in Vertex AI. In this example, you'll compare LLM responses with no grounding with responses that are grounded in the results of a Google Search and a data store in Vertex AI Search.\n",
        "\n",
        "You'll ask a question about Vertex AI and a follow up question about managed datasets in Vertex AI:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "8d20f362cd19"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"What are managed datasets in Vertex AI?\"\n",
        "PROMPT_FOLLOWUP = \"What types of data can I use\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b4e8fa97a38e"
      },
      "source": [
        "### Chat session without grounding\n",
        "\n",
        "Start a chat session and send messages to the LLM with no grounding:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "833e4ff0f540"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " Managed datasets are a feature of Vertex AI that allows you to easily create, manage, and version your datasets. With managed datasets, you can:\n",
            "\n",
            "* **Easily create datasets:** You can create datasets from a variety of sources, including Cloud Storage, BigQuery, and CSV files.\n",
            "* **Manage datasets:** You can view, edit, and delete datasets. You can also add and remove columns from datasets.\n",
            "* **Version datasets:** You can create new versions of datasets. This allows you to track changes to your datasets over time.\n",
            "* **Share datasets:** You can share datasets with other users in your organization.\n",
            "* **Use datasets in Vertex AI models:** You can use managed datasets to train and evaluate Vertex AI models.\n",
            "\n",
            "Managed datasets are a powerful tool that can help you to improve the performance of your Vertex AI models.\n",
            " You can use a variety of data types with managed datasets, including:\n",
            "\n",
            "* **Structured data:** Structured data is data that is organized in a tabular format. Examples of structured data include CSV files, JSON files, and BigQuery tables.\n",
            "* **Unstructured data:** Unstructured data is data that is not organized in a tabular format. Examples of unstructured data include images, videos, and text files.\n",
            "* **Semi-structured data:** Semi-structured data is data that has some structure, but not as much as structured data. Examples of semi-structured data include XML files and HTML files.\n",
            "\n",
            "You can also use managed datasets to combine different types of data. For example, you could create a dataset that includes both structured data and unstructured data.\n"
          ]
        }
      ],
      "source": [
        "chat = chat_model.start_chat()\n",
        "\n",
        "response = chat.send_message(PROMPT)\n",
        "print(response.text)\n",
        "\n",
        "response = chat.send_message(PROMPT_FOLLOWUP)\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a98d8da6eee8"
      },
      "source": [
        "### Chat session grounded in Google Search results\n",
        "\n",
        "Now you can add the `grounding_source` keyword arg with a grounding source of `GroundingSource.WebSearch()` to instruct the chat model to first perform a Google Search with the prompt, then construct an answer based on the web search results:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "02c92f2d55ef"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " Managed datasets in Vertex AI are a way to store and manage your data for use in machine learning models. They provide a number of benefits, including:\n",
            "\n",
            "- **Centralized storage:** Managed datasets are stored in a central location, making them easy to access and manage.\n",
            "- **Data versioning:** Managed datasets support data versioning, so you can easily track changes to your data over time.\n",
            "- **Data security:** Managed datasets are encrypted at rest and in transit, so you can be sure that your data is safe.\n",
            "- **Data processing:** Managed datasets can be processed using a variety of tools, including Vertex AI's built-in data processing tools.\n",
            "GroundingMetadata(citations=[], search_queries=['Vertex AI managed datasets?'])\n",
            " You can use a variety of data types in managed datasets, including:\n",
            "\n",
            "- **Structured data:** Structured data is data that is organized in a tabular format, such as CSV files or SQL tables.\n",
            "- **Unstructured data:** Unstructured data is data that is not organized in a tabular format, such as images, videos, and text files.\n",
            "- **Semi-structured data:** Semi-structured data is data that is partially structured, such as JSON files or XML files.\n",
            "GroundingMetadata(citations=[], search_queries=['What types of data can I use in a managed dataset in Vertex AI?'])\n"
          ]
        }
      ],
      "source": [
        "chat = chat_model.start_chat()\n",
        "grounding_source = GroundingSource.WebSearch()\n",
        "\n",
        "response = chat.send_message(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "print(response.text)\n",
        "print(response.grounding_metadata)\n",
        "\n",
        "response = chat.send_message(\n",
        "    PROMPT_FOLLOWUP,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "print(response.text)\n",
        "print(response.grounding_metadata)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69580e133023"
      },
      "source": [
        "### Chat session grounded in Vertex AI Search results\n",
        "\n",
        "Now you can add the `grounding_source` keyword arg with a grounding source of `GroundingSource.VertexAISearch()` to instruct the chat model to first perform a search within your custom data store, then construct an answer based on the relevant documents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "1449a93cda4f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " Managed datasets in Vertex AI are used to provide the source data for training AutoML and custom models.\n",
            "GroundingMetadata(citations=[GroundingCitation(start_index=1, end_index=105, url='https://cloud.google.com/vertex-ai/docs/datasets/overview', title=None, license=None, publication_date=None)], search_queries=['Vertex AI managed datasets?'])\n",
            " Managed datasets in Vertex AI are used to provide the source data for training AutoML and custom models.\n",
            "GroundingMetadata(citations=[GroundingCitation(start_index=1, end_index=105, url='https://cloud.google.com/vertex-ai/docs/datasets/overview', title=None, license=None, publication_date=None)], search_queries=['Vertex AI managed datasets'])\n"
          ]
        }
      ],
      "source": [
        "chat = chat_model.start_chat()\n",
        "grounding_source = GroundingSource.VertexAISearch(\n",
        "    data_store_id=DATA_STORE_ID, location=DATA_STORE_REGION\n",
        ")\n",
        "\n",
        "response = chat.send_message(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "print(response.text)\n",
        "print(response.grounding_metadata)\n",
        "\n",
        "response = chat.send_message(\n",
        "    PROMPT_FOLLOWUP,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "print(response.text)\n",
        "print(response.grounding_metadata)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "56f5acaa7ce7"
      },
      "source": [
        "## Example: Grounded async text and chat responses\n",
        "\n",
        "You can also use grounding in Vertex AI when working with the asynchronous APIs for the text and chat models. In this example, you'll compare LLM responses with no grounding with responses that are grounded in the results of a data store in Vertex AI Search.\n",
        "\n",
        "You'll ask a question about different services available in Google Cloud."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "91b723359992"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"What are the different types of databases available in Google Cloud?\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61ed9081abbd"
      },
      "source": [
        "### Async text generation grounded in Google Search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "dc8fa27b0c3e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( The different types of databases available in Google Cloud are:\n",
              "\n",
              "1. Cloud Spanner: It provides all the relational database capabilities of Cloud SQL along with horizontal scalability which usually comes with NoSQL databases.\n",
              "\n",
              "2. Cloud Bigtable: Users can store different types of data, including time-series, marketing, financial, IoT, and graph data. Cloud Bigtable also integrates with popular big data.\n",
              "\n",
              "3. Cloud SQL: Provides managed MySQL, PostgreSQL, and SQL Server databases on Google Cloud.\n",
              "\n",
              "4. AlloyDB: It is a fully managed PostgreSQL-compatible database service that offers high performance and scalability for PostgreSQL workloads.,\n",
              " GroundingMetadata(citations=[GroundingCitation(start_index=69, end_index=226, url='https://medium.com/google-cloud/choose-the-right-database-service-in-gcp-8e3803245e1d', title=None, license=None, publication_date=None), GroundingCitation(start_index=230, end_index=352, url='https://www.techtarget.com/searchcloudcomputing/feature/7-Google-Cloud-database-options-to-free-up-your-IT-team', title=None, license=None, publication_date=None), GroundingCitation(start_index=353, end_index=407, url='https://www.techtarget.com/searchcloudcomputing/feature/7-Google-Cloud-database-options-to-free-up-your-IT-team', title=None, license=None, publication_date=None), GroundingCitation(start_index=411, end_index=500, url='https://cloud.google.com/blog/topics/developers-practitioners/your-google-cloud-database-options-explained', title=None, license=None, publication_date=None)], search_queries=['What are the different types of databases available in Google Cloud?']))"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "grounding_source = GroundingSource.WebSearch()\n",
        "\n",
        "response = await text_model.predict_async(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dba56fdc56ff"
      },
      "source": [
        "### Async text generation grounded in Vertex AI Search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "4c6d62c0dab6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( The different types of databases available in Google Cloud are:\n",
              "\n",
              "1. **Cloud SQL**: A fully-managed database service that supports MySQL, PostgreSQL, and SQL Server.\n",
              "2. **BigQuery**: A serverless, highly scalable data warehouse that can handle petabytes of data.\n",
              "3. **Spanner**: A globally distributed, highly scalable relational database that supports ACID transactions.\n",
              "4. **Firestore**: A NoSQL document database that is ideal for real-time applications.\n",
              "5. **Memorystore**: An in-memory data store that is ideal for applications that require fast access to data.,\n",
              " GroundingMetadata(citations=[GroundingCitation(start_index=1, end_index=65, url='https://cloud.google.com/learn/what-is-a-cloud-database', title=None, license=None, publication_date=None)], search_queries=['What are the different types of databases available in Google Cloud?']))"
            ]
          },
          "execution_count": 23,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "grounding_source = GroundingSource.VertexAISearch(\n",
        "    data_store_id=DATA_STORE_ID, location=DATA_STORE_REGION\n",
        ")\n",
        "\n",
        "response = await text_model.predict_async(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8eba08f7d762"
      },
      "source": [
        "### Async chat session grounded in Google Search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "17f740222bbb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( - Cloud Spanner\n",
              "- Cloud Bigtable\n",
              "- Cloud SQL\n",
              "- AlloyDB,\n",
              " GroundingMetadata(citations=[], search_queries=['What are the different types of databases available in Google Cloud?']))"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "chat = chat_model.start_chat()\n",
        "\n",
        "grounding_source = GroundingSource.WebSearch()\n",
        "response = await chat.send_message_async(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6e7e4abd0ee2"
      },
      "source": [
        "### Async chat session grounded in Vertex AI Search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "1994ac2210fc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "( The different types of databases available in Google Cloud are:\n",
              "\n",
              "1. [2] **Cloud SQL**: A fully-managed database service that supports MySQL, PostgreSQL, and SQL Server.\n",
              "2. [2] **BigQuery**: A serverless, highly scalable data warehouse that can handle large amounts of data.\n",
              "3. [2] **Spanner**: A globally-distributed, highly available relational database.\n",
              "4. [2] **Firestore**: A NoSQL document database that is ideal for real-time applications.\n",
              "5. [2] **Memorystore**: An in-memory data store that is optimized for high-performance applications.,\n",
              " GroundingMetadata(citations=[GroundingCitation(start_index=1, end_index=65, url='https://cloud.google.com/learn/what-is-a-cloud-database', title=None, license=None, publication_date=None)], search_queries=['What are the different types of databases available in Google Cloud?']))"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "chat = chat_model.start_chat()\n",
        "\n",
        "grounding_source = GroundingSource.VertexAISearch(\n",
        "    data_store_id=DATA_STORE_ID, location=DATA_STORE_REGION\n",
        ")\n",
        "response = await chat.send_message_async(\n",
        "    PROMPT,\n",
        "    grounding_source=grounding_source,\n",
        ")\n",
        "\n",
        "response, response.grounding_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid incurring charges to your Google Cloud account for the resources used in this notebook, follow these steps:\n",
        "\n",
        "1. To avoid unnecessary Google Cloud charges, use the [Google Cloud console](https://console.cloud.google.com/) to delete your project if you do not need it. Learn more in the Google Cloud documentation for [managing and deleting your project](https://cloud.google.com/resource-manager/docs/creating-managing-projects).\n",
        "1. If you used an existing Google Cloud project, delete the resources you created to avoid incurring charges to your account. For more information, refer to the documentation to [Delete data from a data store in Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/delete-datastores), then delete your data store.\n",
        "1. Disable the [Vertex AI Search and Conversation API](https://console.cloud.google.com/apis/api/discoveryengine.googleapis.com) and [Vertex AI API](https://console.cloud.google.com/apis/api/aiplatform.googleapis.com) in the Google Cloud Console."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "intro-grounding.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}