gemini/rag-engine/rag_engine_vector_search.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": [
        "# Vertex AI RAG Engine with Vertex AI Vector 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/rag-engine/rag_engine_vector_search.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%2Frag-engine%2Frag_engine_vector_search.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/rag-engine/rag_engine_vector_search.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/rag-engine/rag_engine_vector_search.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",
        "</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/rag-engine/rag_engine_vector_search.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/rag-engine/rag_engine_vector_search.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/rag-engine/rag_engine_vector_search.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/rag-engine/rag_engine_vector_search.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/rag-engine/rag_engine_vector_search.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) | [Holt Skinner](https://github.com/holtskinner) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook illustrates how to use [Vertex AI RAG Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-overview) with [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview) as a vector database.\n",
        "\n",
        "For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/use-vertexai-vector-search).\n",
        "\n",
        "For more details on RAG corpus/file management and detailed support please visit [Vertex AI RAG Engine API](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/rag-api)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[33mWARNING: Skipping /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_aiplatform-1.50.0.dist-info due to invalid metadata entry 'name'\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Skipping /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_storage-2.16.0.dist-info due to invalid metadata entry 'name'\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mDEPRECATION: Loading egg at /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/fsspec-2024.3.1-py3.11.egg is deprecated. pip 24.3 will enforce this behaviour change. A possible replacement is to use pip for package installation.. Discussion can be found at https://github.com/pypa/pip/issues/12330\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mDEPRECATION: Loading egg at /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_documentai_toolbox-0.12.2a0-py3.11.egg is deprecated. pip 24.3 will enforce this behaviour change. A possible replacement is to use pip for package installation.. Discussion can be found at https://github.com/pypa/pip/issues/12330\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mDEPRECATION: Loading egg at /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_documentai_toolbox-0.11.1a0-py3.11.egg is deprecated. pip 24.3 will enforce this behaviour change. A possible replacement is to use pip for package installation.. Discussion can be found at https://github.com/pypa/pip/issues/12330\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Skipping /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_aiplatform-1.50.0.dist-info due to invalid metadata entry 'name'\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Skipping /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_storage-2.16.0.dist-info due to invalid metadata entry 'name'\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Skipping /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_aiplatform-1.50.0.dist-info due to invalid metadata entry 'name'\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Skipping /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/google_cloud_aiplatform-1.50.0.dist-info due to invalid metadata entry 'name'\u001b[0m\u001b[33m\n",
            "\u001b[0m\n",
            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n",
            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip3.11 install --upgrade pip\u001b[0m\n",
            "Note: you may need to restart the kernel to use updated packages.\n"
          ]
        }
      ],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform google-genai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "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": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "from google import genai\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "adbe5c6b3549"
      },
      "source": [
        "## (Optional) Setup Vertex AI Vector Search index and index endpoint\n",
        "\n",
        "In this section, we have some helper methods to help you setup your Vector Search index.\n",
        "\n",
        "This section is not required if you already have a Vector Search index ready to use.\n",
        "\n",
        "The index has to meet the following criteria:\n",
        "\n",
        "1. `IndexUpdateMethod` must be `STREAM_UPDATE`, see [Create stream index]({{docs_path}}vector-search/create-manage-index#create-stream-index).\n",
        "\n",
        "2. Distance measure type must be explicitly set to one of the following:\n",
        "\n",
        "   * `DOT_PRODUCT_DISTANCE`\n",
        "   * `COSINE_DISTANCE`\n",
        "\n",
        "3. Dimension of the vector must be consistent with the embedding model you plan\n",
        "   to use in the RAG corpus. Other parameters can be tuned based on\n",
        "   your choices, which determine whether the additional parameters can be\n",
        "   tuned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ee177c9bc175"
      },
      "outputs": [],
      "source": [
        "# create the index\n",
        "my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
        "    display_name=\"your_display_name\",\n",
        "    description=\"your_description\",\n",
        "    dimensions=768,\n",
        "    approximate_neighbors_count=10,\n",
        "    leaf_node_embedding_count=500,\n",
        "    leaf_nodes_to_search_percent=7,\n",
        "    distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n",
        "    feature_norm_type=\"UNIT_L2_NORM\",\n",
        "    index_update_method=\"STREAM_UPDATE\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02e52d8dcda6"
      },
      "source": [
        "RAG Engine supports [public endpoints](https://cloud.google.com/vertex-ai/docs/vector-search/deploy-index-public)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce6e0e85adf1"
      },
      "outputs": [],
      "source": [
        "# create IndexEndpoint\n",
        "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
        "    display_name=\"your_display_name\", public_endpoint_enabled=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c4c3f91ab95f"
      },
      "source": [
        "Deploy the index to the index endpoint.\n",
        "\n",
        "If it's the first time that you're deploying an index to an index endpoint, it\n",
        "takes approximately 30 minutes to automatically build and initiate the backend\n",
        "before the index can be stored. After the first deployment, the index is ready\n",
        "in seconds. To see the status of the index deployment, open the\n",
        "[**Vector Search Console**](https://console.cloud.google.com/vertex-ai/matching-engine/index-endpoints),\n",
        "select the **Index endpoints** tab, and choose your index endpoint.\n",
        "\n",
        "Identify the resource name of your index and index endpoint, which have the\n",
        "following the formats:\n",
        "\n",
        "* `projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexes/${INDEX_ID}`\n",
        "* `projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexEndpoints/${INDEX_ENDPOINT_ID}`.\n",
        "\n",
        "If you aren't sure about the resource name, you can use the following command to\n",
        "check:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "382010f08560"
      },
      "outputs": [],
      "source": [
        "print(my_index_endpoint.resource_name)\n",
        "print(my_index.resource_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6dab214cd107"
      },
      "outputs": [],
      "source": [
        "# Deploy Index\n",
        "my_index_endpoint.deploy_index(\n",
        "    index=my_index, deployed_index_id=\"your_deployed_index_id\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Use Vertex AI Vector Search in RAG Engine\n",
        "\n",
        "After the Vector Search instance is set up, follow the steps in this section to set the Vector Search instance as the vector database for the RAG application.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cd05469c3e71"
      },
      "source": [
        "### Set the vector database to create a RAG corpus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b9ad5442bd4e"
      },
      "outputs": [],
      "source": [
        "from google.genai.types import (\n",
        "    GenerateContentConfig,\n",
        "    Retrieval,\n",
        "    Tool,\n",
        "    VertexRagStore,\n",
        "    VertexRagStoreRagResource,\n",
        ")\n",
        "from vertexai import rag"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53865b3ea33e"
      },
      "outputs": [],
      "source": [
        "vector_db = rag.VertexVectorSearch(\n",
        "    index=my_index.resource_name, index_endpoint=my_index_endpoint.resource_name\n",
        ")\n",
        "\n",
        "# Name your corpus\n",
        "DISPLAY_NAME = \"\"  # @param  {type:\"string\"}\n",
        "\n",
        "# Create RAG Corpus\n",
        "rag_corpus = rag.create_corpus(\n",
        "    display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db)\n",
        ")\n",
        "print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93a3296647a2"
      },
      "source": [
        "## Upload a file to the corpus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7f31cc83fb04"
      },
      "outputs": [],
      "source": [
        "%%writefile test.txt\n",
        "\n",
        "Here's a demo for Vertex AI Vector Search RAG."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7bab0e824c3d"
      },
      "outputs": [],
      "source": [
        "rag_file = rag.upload_file(\n",
        "    corpus_name=rag_corpus.name,\n",
        "    path=\"test.txt\",\n",
        "    display_name=\"test.txt\",\n",
        "    description=\"my test\",\n",
        ")\n",
        "print(f\"Uploaded file to resource: {rag_file.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e51b5bcd1739"
      },
      "source": [
        "## Import files from Google Cloud Storage\n",
        "\n",
        "Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Google Cloud Storage bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e0e53a05445e"
      },
      "outputs": [],
      "source": [
        "GCS_BUCKET = \"\"  # @param {type:\"string\", \"placeholder\": \"your-gs-bucket\"}\n",
        "\n",
        "response = rag.import_files(  # noqa: F704\n",
        "    corpus_name=rag_corpus.name,\n",
        "    paths=[GCS_BUCKET],\n",
        "    transformation_config=rag.TransformationConfig(\n",
        "        chunking_config=rag.ChunkingConfig(\n",
        "            chunk_size=512,\n",
        "            chunk_overlap=50,\n",
        "        )\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "48313a38ef52"
      },
      "outputs": [],
      "source": [
        "# Check the files just imported. It may take a few seconds to process the imported files.\n",
        "rag.list_files(corpus_name=rag_corpus.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ceab91983444"
      },
      "source": [
        "## Import files from Google Drive\n",
        "\n",
        "Eligible paths can be:\n",
        "\n",
        "- `https://drive.google.com/drive/folders/{folder_id}`\n",
        "- `https://drive.google.com/file/d/{file_id}`\n",
        "\n",
        "Remember to grant \"Viewer\" access to the \"Vertex RAG Data Service Agent\" (with the format of `service-{project_number}@gcp-sa-vertex-rag.iam.gserviceaccount.com`) for your Drive folder/files.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ea8a5c97ad80"
      },
      "outputs": [],
      "source": [
        "FILE_ID = \"\"  # @param {type:\"string\", \"placeholder\": \"your-file-id\"}\n",
        "FILE_PATH = f\"https://drive.google.com/file/d/{FILE_ID}\"\n",
        "\n",
        "rag.import_files(\n",
        "    corpus_name=rag_corpus.name,\n",
        "    paths=[FILE_PATH],\n",
        "    transformation_config=rag.TransformationConfig(\n",
        "        chunking_config=rag.ChunkingConfig(\n",
        "            chunk_size=1024,\n",
        "            chunk_overlap=100,\n",
        "        )\n",
        "    ),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e71887752baa"
      },
      "outputs": [],
      "source": [
        "# Check the files just imported. It may take a few seconds to process the imported files.\n",
        "rag.list_files(corpus_name=rag_corpus.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "346ceb446e7c"
      },
      "source": [
        "## Use your RAG Corpus to add context to your Gemini queries\n",
        "\n",
        "When retrieved contexts similarity distance < `vector_distance_threshold`, the contexts (from `RagStore`) will be used for content generation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fec72ac982c3"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.0-flash-001\"\n",
        "\n",
        "rag_retrieval_tool = Tool(\n",
        "    retrieval=Retrieval(\n",
        "        vertex_rag_store=VertexRagStore(\n",
        "            rag_resources=[\n",
        "                VertexRagStoreRagResource(\n",
        "                    rag_corpus=rag_corpus.name  # Currently only 1 corpus is allowed.\n",
        "                )\n",
        "            ],\n",
        "            similarity_top_k=10,\n",
        "            vector_distance_threshold=0.4,\n",
        "        )\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cc0ee39e50f6"
      },
      "outputs": [],
      "source": [
        "GENERATE_CONTENT_PROMPT = \"What is RAG and why it is helpful?\"  # @param {type:\"string\"}\n",
        "\n",
        "response = client.models.generate_content(\n",
        "    model=MODEL_ID,\n",
        "    contents=GENERATE_CONTENT_PROMPT,\n",
        "    config=GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
        ")\n",
        "\n",
        "display(Markdown(response.text))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2899daa12fac"
      },
      "source": [
        "## Using other generation API with Rag Retrieval Tool\n",
        "\n",
        "The retrieved contexts can be passed to any SDK or model generation API to generate final results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d549fb12733f"
      },
      "outputs": [],
      "source": [
        "RETRIEVAL_QUERY = \"What is RAG and why it is helpful?\"  # @param {type:\"string\"}\n",
        "\n",
        "rag_resource = rag.RagResource(\n",
        "    rag_corpus=rag_corpus.name,\n",
        "    # Need to manually get the ids from rag.list_files.\n",
        "    # rag_file_ids=[],\n",
        ")\n",
        "\n",
        "response = rag.retrieval_query(\n",
        "    rag_resources=[rag_resource],  # Currently only 1 corpus is allowed.\n",
        "    text=RETRIEVAL_QUERY,\n",
        "    rag_retrieval_config=rag.RagRetrievalConfig(\n",
        "        top_k=10,  # Optional\n",
        "        filter=rag.Filter(\n",
        "            vector_distance_threshold=0.5,  # Optional\n",
        "        ),\n",
        "    ),\n",
        ")\n",
        "\n",
        "# The retrieved context can be passed to any SDK or model generation API to generate final results.\n",
        "retrieved_context = \" \".join(\n",
        "    [context.text for context in response.contexts.contexts]\n",
        ").replace(\"\\n\", \"\")\n",
        "\n",
        "retrieved_context"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "Clean up resources created in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ea74a96756a3"
      },
      "outputs": [],
      "source": [
        "delete_rag_corpus = False  # @param {type:\"boolean\"}\n",
        "\n",
        "if delete_rag_corpus:\n",
        "    rag.delete_corpus(name=rag_corpus.name)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "rag_engine_vector_search.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}