gemini/grounding/intro-grounding-gemini.ipynb

{
  "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": [
        "# Getting Started with Grounding with Gemini in Vertex AI\n",
        "\n",
        "> **NOTE:** This notebook uses the Vertex AI SDK, which does not support Gemini 2.0; refer to [Intro to Gemini 2.0 Flash](https://github.com/GoogleCloudPlatform/generative-ai/blob/sdk-updates/gemini/getting-started/intro_gemini_2_0_flash.ipynb) which uses the Google Gen AI SDK.\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/grounding/intro-grounding-gemini.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Run 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%2Fgrounding%2Fintro-grounding-gemini.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/grounding/intro-grounding-gemini.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" 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/gemini/grounding/intro-grounding-gemini.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" 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/gemini/grounding/intro-grounding-gemini.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/grounding/intro-grounding-gemini.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/grounding/intro-grounding-gemini.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/gemini/grounding/intro-grounding-gemini.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/grounding/intro-grounding-gemini.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": "49e1e41cea0d"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Holt Skinner](https://github.com/holtskinner), [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/generative-ai/docs/multimodal/ground-gemini) lets you use generative text models 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",
        "2. [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",
        "### Allowlisting\n",
        "\n",
        "Some of the features in this sample notebook require access to certain features via an allowlist. [Grounding with Vertex AI Search](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/ground-gemini) is available in Public Preview, whereas Grounding with Google Search results is generally available.\n",
        "\n",
        "If you use this service in a production application, you will also need to [use a Google Search entry point](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/grounding-search-entry-points)."
      ]
    },
    {
      "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",
        "\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 and Vertex AI Agent Builder APIs](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,discoveryengine.googleapis.com).\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": null,
      "metadata": {
        "id": "2b4ef9b72d43"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "Restart the kernel after installing packages:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb05997a457d"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "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": null,
      "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": "WReHDGG5g0XY"
      },
      "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).\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)\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": 2,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "from IPython.display import Markdown, display\n",
        "from vertexai.generative_models import (\n",
        "    GenerationResponse,\n",
        "    GenerativeModel,\n",
        "    Tool,\n",
        "    grounding,\n",
        ")\n",
        "from vertexai.preview.generative_models import grounding as preview_grounding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "55cf2dd17690"
      },
      "source": [
        "Initialize the Gemini model from Vertex AI:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "652a8969dd5a"
      },
      "outputs": [],
      "source": [
        "model = GenerativeModel(\"gemini-1.5-flash\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff453ac772d4"
      },
      "source": [
        "## Helper functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "3dce100cab74"
      },
      "outputs": [],
      "source": [
        "def print_grounding_response(response: GenerationResponse):\n",
        "    \"\"\"Prints Gemini response with grounding citations.\"\"\"\n",
        "    grounding_metadata = response.candidates[0].grounding_metadata\n",
        "\n",
        "    # Citation indices are in byte units\n",
        "    ENCODING = \"utf-8\"\n",
        "    text_bytes = response.text.encode(ENCODING)\n",
        "\n",
        "    prev_index = 0\n",
        "    markdown_text = \"\"\n",
        "\n",
        "    for grounding_support in grounding_metadata.grounding_supports:\n",
        "        text_segment = text_bytes[\n",
        "            prev_index : grounding_support.segment.end_index\n",
        "        ].decode(ENCODING)\n",
        "\n",
        "        footnotes_text = \"\"\n",
        "        for grounding_chunk_index in grounding_support.grounding_chunk_indices:\n",
        "            footnotes_text += f\"[{grounding_chunk_index + 1}]\"\n",
        "\n",
        "        markdown_text += f\"{text_segment} {footnotes_text}\\n\"\n",
        "        prev_index = grounding_support.segment.end_index\n",
        "\n",
        "    if prev_index < len(text_bytes):\n",
        "        markdown_text += str(text_bytes[prev_index:], encoding=ENCODING)\n",
        "\n",
        "    markdown_text += \"\\n----\\n## Grounding Sources\\n\"\n",
        "\n",
        "    if grounding_metadata.web_search_queries:\n",
        "        markdown_text += (\n",
        "            f\"\\n**Web Search Queries:** {grounding_metadata.web_search_queries}\\n\"\n",
        "        )\n",
        "        if grounding_metadata.search_entry_point:\n",
        "            markdown_text += f\"\\n**Search Entry Point:**\\n {grounding_metadata.search_entry_point.rendered_content}\\n\"\n",
        "    elif grounding_metadata.retrieval_queries:\n",
        "        markdown_text += (\n",
        "            f\"\\n**Retrieval Queries:** {grounding_metadata.retrieval_queries}\\n\"\n",
        "        )\n",
        "\n",
        "    markdown_text += \"### Grounding Chunks\\n\"\n",
        "\n",
        "    for index, grounding_chunk in enumerate(\n",
        "        grounding_metadata.grounding_chunks, start=1\n",
        "    ):\n",
        "        context = grounding_chunk.web or grounding_chunk.retrieved_context\n",
        "        if not context:\n",
        "            print(f\"Skipping Grounding Chunk {grounding_chunk}\")\n",
        "            continue\n",
        "\n",
        "        markdown_text += f\"{index}. [{context.title}]({context.uri})\\n\"\n",
        "\n",
        "    display(Markdown(markdown_text))"
      ]
    },
    {
      "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": null,
      "metadata": {
        "id": "6a28ca4abb52"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"You are an expert in astronomy. When is the next solar eclipse in the US?\""
      ]
    },
    {
      "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": null,
      "metadata": {
        "id": "a2e348ff93e6"
      },
      "outputs": [],
      "source": [
        "response = model.generate_content(PROMPT)\n",
        "\n",
        "display(Markdown(response.text))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5d7cb7cceb99"
      },
      "source": [
        "### Text generation grounded in Google Search results\n",
        "\n",
        "Now you can add the `tools` keyword arg with a `grounding` tool of `grounding.GoogleSearchRetrieval()` to instruct the LLM to first perform a Google Search with the prompt, then construct an answer based on the web search results.\n",
        "\n",
        "The search queries and [Search Entry Point](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/grounding-search-entry-points) are available for each `Candidate` in the response. The helper function `print_grounding_response()` prints the response text with citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1d9fb83b0ab9"
      },
      "outputs": [],
      "source": [
        "tool = Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval())\n",
        "\n",
        "response = model.generate_content(PROMPT, tools=[tool])\n",
        "\n",
        "print_grounding_response(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6d3920bb2ac0"
      },
      "source": [
        "Note that the response without grounding only has limited information from the LLM about solar eclipses. 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).\n",
        "\n",
        "The data store will contain internal documents from a fictional bank, Cymbal Bank. These documents aren't available on the public internet, so the Gemini model won't have any information about them by default."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1b308548c68b"
      },
      "source": [
        "### Creating a data store in Vertex AI Search\n",
        "\n",
        "In this example, you'll use a Google Cloud Storage bucket with a few sample internal documents for our bank. There's some docs about booking business travel, strategic plan for this Fiscal Year and HR docs describing the different jobs available in the company.\n",
        "\n",
        "Follow the tutorial steps in the Vertex AI Search documentation to:\n",
        "\n",
        "1. [Create a data store with unstructured data](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#unstructured-data) that loads in documents from the GCS folder `gs://cloud-samples-data/gen-app-builder/search/cymbal-bank-employee`.\n",
        "2. [Create a search app](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_a_search_app) that is attached to that data store. You should also enable the **Enterprise edition features** so that you can search indexed records within the data store.\n",
        "\n",
        "Note: The data store must be in the same project that you are using for Gemini.\n",
        "\n",
        "Once you've created a data store, obtain the Data Store ID and input it below.\n",
        "\n",
        "Note: You will need to wait for data ingestion to finish before using a data store with grounding. For more information, see [create a data store](https://cloud.google.com/generative-ai-app-builder/docs/create-data-store-es)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "fcd767476241"
      },
      "outputs": [],
      "source": [
        "DATA_STORE_PROJECT_ID = PROJECT_ID  # @param {type:\"string\"}\n",
        "DATA_STORE_REGION = \"global\"  # @param {type:\"string\"}\n",
        "# Replace this with your data store ID from Vertex AI Search\n",
        "DATA_STORE_ID = \"your-data-store-id_1234567890123\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ccc156676e0a"
      },
      "source": [
        "Now you can ask a question about the company culture:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "9c1e1b1743bd"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"What is the company culture like?\""
      ]
    },
    {
      "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": null,
      "metadata": {
        "id": "299818ae71e9"
      },
      "outputs": [],
      "source": [
        "response = model.generate_content(PROMPT)\n",
        "\n",
        "display(Markdown(response.text))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "073f2ec42ff6"
      },
      "source": [
        "### Text generation grounded in Vertex AI Search results\n",
        "\n",
        "Now we can add the `tools` keyword arg with a grounding tool of `grounding.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": null,
      "metadata": {
        "id": "d4c5d53a37b4"
      },
      "outputs": [],
      "source": [
        "tool = Tool.from_retrieval(\n",
        "    preview_grounding.Retrieval(\n",
        "        preview_grounding.VertexAISearch(\n",
        "            datastore=DATA_STORE_ID,\n",
        "            project=DATA_STORE_PROJECT_ID,\n",
        "            location=DATA_STORE_REGION,\n",
        "        )\n",
        "    )\n",
        ")\n",
        "\n",
        "response = model.generate_content(PROMPT, tools=[tool])\n",
        "print_grounding_response(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e3f985c704cd"
      },
      "source": [
        "Note that the response without grounding doesn't have any context about what company we are asking about. Whereas the response that was grounded in Vertex AI Search results contains information from the documents provided, along with citations of the information.\n",
        "\n",
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ Important notes:</b><br>\n",
        "<br>\n",
        "<b>If you get an error when running the previous cell:</b><br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;In order for this sample notebook to work with data store in Vertex AI Search,<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;you'll need to create a <a href=\"https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_a_data_store\">data store</a> <b>and</b> a <a href=\"https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#create_a_search_app\">search app</a> associated with it in Vertex AI Search.<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;If you only create a data store, the previous request will return errors when making queries against the data store.\n",
        "<br><br>\n",
        "<b>If you get an empty response when running the previous cell:</b><br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;You will need to wait for data ingestion to finish before using a data store with grounding.<br>\n",
        "&nbsp;&nbsp;&nbsp;&nbsp;For more information, see <a href=\"https://cloud.google.com/generative-ai-app-builder/docs/create-data-store-es\">create a data store</a>.\n",
        "</div>\n",
        "</div>"
      ]
    },
    {
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "490cf1ed3399"
      },
      "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": "b59783e4f1ce"
      },
      "source": [
        "### Chat session grounded in Google Search results\n",
        "\n",
        "Now you can add the `tools` keyword arg with a grounding tool of `grounding.GoogleSearchRetrieval()` 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": null,
      "metadata": {
        "id": "58edb2bd860f"
      },
      "outputs": [],
      "source": [
        "chat = model.start_chat()\n",
        "tool = Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval())\n",
        "\n",
        "response = chat.send_message(PROMPT, tools=[tool])\n",
        "print_grounding_response(response)\n",
        "\n",
        "response = chat.send_message(PROMPT_FOLLOWUP, tools=[tool])\n",
        "print_grounding_response(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87be7f661f14"
      },
      "source": [
        "### Chat session grounded in Vertex AI Search results\n",
        "\n",
        "Now we can add the `tools` keyword arg with a grounding tool of `grounding.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": 18,
      "metadata": {
        "id": "8fdad0c3f1f3"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"How do I book business travel?\"\n",
        "PROMPT_FOLLOWUP = \"Give me more details.\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1a824202a8f0"
      },
      "outputs": [],
      "source": [
        "chat = model.start_chat()\n",
        "tool = Tool.from_retrieval(\n",
        "    preview_grounding.Retrieval(\n",
        "        preview_grounding.VertexAISearch(\n",
        "            datastore=DATA_STORE_ID,\n",
        "            project=DATA_STORE_PROJECT_ID,\n",
        "            location=DATA_STORE_REGION,\n",
        "        )\n",
        "    )\n",
        ")\n",
        "response = chat.send_message(PROMPT, tools=[tool])\n",
        "print_grounding_response(response)\n",
        "\n",
        "response = chat.send_message(PROMPT_FOLLOWUP, tools=[tool])\n",
        "print_grounding_response(response)"
      ]
    },
    {
      "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-gemini.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}