{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tuOe1ymfHZPu"
      },
      "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": {
        "id": "3f21741c44b3"
      },
      "source": [
        "# Grounding with Vertex AI Search\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/grounding_with_vais.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fgrounding%2Fgrounding_with_vais.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/grounding/grounding_with_vais.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",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/grounding/grounding_with_vais.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</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/grounding_with_vais.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/grounding_with_vais.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/grounding_with_vais.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/grounding/grounding_with_vais.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/grounding_with_vais.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": "2a457a9fb5a8"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) |  [Diem Vu](https://github.com/diemtvu/) [Zhen Hu](https://github.com/undertwig/) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a10d34f145f2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use Vertex AI Search for grounding Vertex LLMs.  For more general information on grounding, see [Getting Started with Grounding with Gemini in Vertex AI](./intro-grounding-gemini.ipynb).\n",
        "\n",
        "In this tutorial, we will cover:\n",
        "\n",
        "* How to create a Vertex AI Search datastore with your data.\n",
        "* Provide an example LLM request that uses this data for grounding."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cab0c8509c9"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uQwqEaFLHGlL"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6029CKnWG75v"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet google-cloud-aiplatform google-cloud-discoveryengine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f42d12d15616"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f3d98395d9a4"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff050a2aa492"
      },
      "source": [
        "[link text](https://)<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. In Colab or Colab Enterprise, you might see an error message that says \"Your session crashed for an unknown reason.\" This is expected. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xu2xiyI56I17"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b36981f9f93b"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "# fmt: on\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "# You also need to choose a region for your Vertex AI Search datastore.\n",
        "# Vertex AI Search is a multi-region service, supporting us, eu, and global regions.\n",
        "VAIS_LOCATION = \"global\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0139cefe4f5"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8cb01510d7e4"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8dbdd4e7885d"
      },
      "source": [
        "## Create Vertex AI Search engine\n",
        "\n",
        "We will create a Vertex AI Search engine index by ingesting files from Google Cloud Storage (GCS) using the SDK. To explore additional features and UI-based instructions, see the [Vertex AI Search documentation](https://cloud.google.com/enterprise-search).\n",
        "\n",
        "This involves 2 steps:\n",
        "\n",
        "1. Create a datastore (corpus) and import documents: For this demo, we will use the PDF files available at `gs://cloud-samples-data/gen-app-builder/search/alphabet-sec-filings`. You can try with your own data.\n",
        "\n",
        "2. Create `SEARCH_TIER_ENTERPRISE` engine with `SEARCH_ADD_ON_LLM` feature on top of the data store. Enterprise tier is required to get extractive answers from a search query and advanced LLM features to ensure the grounding quality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P719DMtK8t-p"
      },
      "outputs": [],
      "source": [
        "# @title Initialize the VAIS clients\n",
        "\n",
        "import time\n",
        "\n",
        "import google.cloud.discoveryengine_v1 as vais\n",
        "from google.api_core.client_options import ClientOptions\n",
        "\n",
        "client_options = (\n",
        "    ClientOptions(api_endpoint=f\"{VAIS_LOCATION}-discoveryengine.googleapis.com\")\n",
        "    if VAIS_LOCATION != \"global\"\n",
        "    else None\n",
        ")\n",
        "data_store_service_client = vais.DataStoreServiceClient(client_options=client_options)\n",
        "document_service_client = vais.DocumentServiceClient(client_options=client_options)\n",
        "engine_client = vais.EngineServiceClient(client_options=client_options)\n",
        "\n",
        "\n",
        "def wait_for_operation_finish(operation):\n",
        "    while not operation.done:\n",
        "        time.sleep(2)  # sleep 2 seconds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AaPZiXVwEDHZ"
      },
      "outputs": [],
      "source": [
        "# @title Create data store\n",
        "# fmt: off\n",
        "DATA_STORE_ID = \"cymbal\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "# Create a generic search data store with unstructured data content.\n",
        "data_store = vais.DataStore(\n",
        "    display_name=\"Data Store for Vertex LLM Grounding demo\",\n",
        "    industry_vertical=\"GENERIC\",\n",
        "    solution_types=[\"SOLUTION_TYPE_SEARCH\"],\n",
        "    content_config=\"CONTENT_REQUIRED\",\n",
        ")\n",
        "\n",
        "create_data_store_request = vais.CreateDataStoreRequest(\n",
        "    parent=f\"projects/{PROJECT_ID}/locations/{VAIS_LOCATION}/collections/default_collection\",\n",
        "    data_store=data_store,\n",
        "    data_store_id=DATA_STORE_ID,\n",
        ")\n",
        "\n",
        "# The api returns long running operation as response.\n",
        "create_data_store_operation = data_store_service_client.create_data_store(\n",
        "    create_data_store_request\n",
        ")\n",
        "print(\n",
        "    f\"Waiting for create datastore operation to complete: {create_data_store_operation.operation.name}\"\n",
        ")\n",
        "created_data_store = create_data_store_operation.result()\n",
        "data_store_name = created_data_store.name\n",
        "print(f\"Data store {data_store_name} is created.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sLxqk7cEdsbn"
      },
      "outputs": [],
      "source": [
        "# @title Ingest documents\n",
        "\n",
        "# fmt: off\n",
        "GCS_SOURCE = \"gs://cloud-samples-data/gen-app-builder/search/cymbal-bank-employee\"  # @param {type:\"string\"}\n",
        "# fmt: on\n",
        "\n",
        "branch_path = document_service_client.branch_path(\n",
        "    project=PROJECT_ID,\n",
        "    location=VAIS_LOCATION,\n",
        "    data_store=DATA_STORE_ID,\n",
        "    branch=\"default_branch\",\n",
        ")\n",
        "\n",
        "document_service_client.import_documents(\n",
        "    request=vais.ImportDocumentsRequest(\n",
        "        parent=branch_path,\n",
        "        gcs_source=vais.GcsSource(\n",
        "            input_uris=[f\"{GCS_SOURCE}/*\"], data_schema=\"content\"\n",
        "        ),\n",
        "        reconciliation_mode=vais.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4277e5df42ad"
      },
      "outputs": [],
      "source": [
        "# @title Wait for documents imported\n",
        "\n",
        "document_imported = False\n",
        "while not document_imported:\n",
        "    for doc in document_service_client.list_documents(parent=branch_path):\n",
        "        print(doc.content.uri)\n",
        "        document_imported = True\n",
        "    time.sleep(10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a634cbdf1c81"
      },
      "outputs": [],
      "source": [
        "# @title Create engine\n",
        "\n",
        "engine_id = f\"{DATA_STORE_ID}_engine\"\n",
        "create_engine_request = vais.CreateEngineRequest(\n",
        "    parent=engine_client.collection_path(\n",
        "        project=PROJECT_ID, location=VAIS_LOCATION, collection=\"default_collection\"\n",
        "    ),\n",
        "    engine=vais.Engine(\n",
        "        display_name=\"Engine for Vertex LLM Grounding demo\",\n",
        "        solution_type=vais.SolutionType.SOLUTION_TYPE_SEARCH,\n",
        "        search_engine_config=vais.Engine.SearchEngineConfig(\n",
        "            search_tier=vais.SearchTier.SEARCH_TIER_ENTERPRISE,\n",
        "            search_add_ons=[vais.SearchAddOn.SEARCH_ADD_ON_LLM],\n",
        "        ),\n",
        "        common_config=vais.Engine.CommonConfig(company_name=\"Cymbal\"),\n",
        "        industry_vertical=vais.IndustryVertical.GENERIC,\n",
        "        data_store_ids=[DATA_STORE_ID],\n",
        "    ),\n",
        "    engine_id=engine_id,\n",
        ")\n",
        "create_engine_operation = engine_client.create_engine(create_engine_request)\n",
        "wait_for_operation_finish(create_engine_operation)\n",
        "print(\"Successfully create engine\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5dcbed43322"
      },
      "source": [
        "Give it a few minutes to build the index. To verify the engine is ready to use, you may try a search request against in. For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "16cc2dd8a944"
      },
      "outputs": [],
      "source": [
        "search_client = vais.SearchServiceClient(client_options=client_options)\n",
        "\n",
        "while True:\n",
        "    try:\n",
        "        search_client.search(\n",
        "            vais.SearchRequest(\n",
        "                serving_config=(\n",
        "                    f\"projects/{PROJECT_ID}/locations/{VAIS_LOCATION}/collections/default_collection/engines/{engine_id}/servingConfigs/default_search\"\n",
        "                ),\n",
        "                query=\"Who is the CEO?\",\n",
        "                page_size=10,\n",
        "            )\n",
        "        )\n",
        "    except Exception:\n",
        "        time.sleep(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "502482711fb8"
      },
      "source": [
        "## Calling Vertex AI LLM grounded on your Vertex AI Search engine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df114061bfc8"
      },
      "source": [
        "Once your engine is ready, you can use it as a grounding source in a Vertex AI LLM call, as shown below:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2354eb2ca244"
      },
      "outputs": [],
      "source": [
        "# @title Initialize Gen AI client\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "from google import genai\n",
        "from google.genai.types import GenerateContentConfig, Retrieval, Tool, VertexAISearch\n",
        "\n",
        "# fmt: off\n",
        "MODEL_ID = \"gemini-2.0-flash\"  # @param {type: \"string\"}\n",
        "# fmt: on\n",
        "\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "vais_tool = Tool(\n",
        "    retrieval=Retrieval(\n",
        "        vertex_ai_search=VertexAISearch(\n",
        "            engine=f\"projects/{PROJECT_ID}/locations/global/collections/default_collection/engines/{engine_id}\",\n",
        "        )\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l6vmUHS0cUd-"
      },
      "outputs": [],
      "source": [
        "# fmt: off\n",
        "PROMPT = \"Instruct me how to book the flight for my business trip\"  # @param {type: \"string\"}\n",
        "# fmt: on\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=PROMPT,\n",
        "    config=GenerateContentConfig(tools=[vais_tool]),\n",
        ")\n",
        "\n",
        "if response.candidates:\n",
        "    display(Markdown(response.text))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7862dbae55a9"
      },
      "source": [
        "Behind the scenes, the model will send a search query to the Vertex AI engine to retrieve information that helps it generate a response. If the generated text makes a claim strongly based on the retrieved result, the model will also output that correlation under grounding metadata. For example, you can see the supporting documents for specific claims in the generated output:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cca93ea35d73"
      },
      "outputs": [],
      "source": [
        "if response.candidates[0].grounding_metadata.grounding_supports:\n",
        "    for s in response.candidates[0].grounding_metadata.grounding_supports:\n",
        "        display(Markdown(f\"{s.segment.text} {s.grounding_chunk_indices}\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72d87754bd54"
      },
      "source": [
        "And to see all retrieved results:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "60cf74052062"
      },
      "outputs": [],
      "source": [
        "for i, chunk in enumerate(response.candidates[0].grounding_metadata.grounding_chunks):\n",
        "    print(f\"{i}.\\n\")\n",
        "    display(Markdown(chunk.retrieved_context.text))\n",
        "    print(chunk.retrieved_context.uri)\n",
        "    print(\"\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ebfaaf2852e9"
      },
      "source": [
        "Note: The model may not always output grounding support, even with successful retrieval. This occurs when the model doesn't find a strong enough corroboration in the retrieved information."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n7jUXKkZVOjn"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "Delete `engine` and `dataStore` created in this tutorial"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3caf56cb3422"
      },
      "outputs": [],
      "source": [
        "# Delete engine\n",
        "engine_full_resource_path = engine_client.engine_path(\n",
        "    project=PROJECT_ID,\n",
        "    location=VAIS_LOCATION,\n",
        "    collection=\"default_collection\",\n",
        "    engine=engine_id,\n",
        ")\n",
        "\n",
        "delete_engine_operation = engine_client.delete_engine(engine_full_resource_path)\n",
        "print(\n",
        "    f\"Waiting for delete engine operation to complete: {delete_engine_operation.operation.name}\"\n",
        ")\n",
        "wait_for_operation_finish(delete_engine_operation)\n",
        "print(f\"Successfully deleted engine {engine_full_resource_path}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LaAt7gOB3nId"
      },
      "outputs": [],
      "source": [
        "# Delete the data store\n",
        "data_store_full_resource_path = data_store_service_client.data_store_path(\n",
        "    project=PROJECT_ID, location=VAIS_LOCATION, data_store=DATA_STORE_ID\n",
        ")\n",
        "delete_data_store_operation = data_store_service_client.delete_data_store(\n",
        "    name=data_store_full_resource_path\n",
        ")\n",
        "print(\n",
        "    f\"Waiting for delete data store operation to complete: {delete_data_store_operation.operation.name}\"\n",
        ")\n",
        "wait_for_operation_finish(delete_data_store_operation)\n",
        "print(f\"Successfully deleted data store {data_store_full_resource_path}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "grounding_with_vais.ipynb",
      "toc_visible": true
    },
    "google": {
      "image_path": "/site-assets/images/share.png",
      "keywords": [
        "examples",
        "samplecode",
        "python",
        "embed"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
