Add new notebook for Knowledge Engine API with Pinecone.

PiperOrigin-RevId: 684666941

notebooks/community/model_garden/knowledge_engine_pinecone.ipynb

@@ -0,0 +1,891 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "OraQZI7DEmGi"
+      },
+      "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": "YVZwr1OhE03N"
+      },
+      "source": [
+        "# Using Vertex AI Knowledge Engine with Pinecone\n",
+        "\n",
+        "\u003ctable align=\"left\"\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/knowledge_engine_pinecone.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"\u003e\u003cbr\u003e Open in Colab\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fknowledge_engine_pinecone.ipynb\"\"\u003e\n",
+        "      \u003cimg width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"\u003e\u003cbr\u003e Open in Colab Enterprise\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e    \n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/knowledge_engine_pinecone.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"\u003e\u003cbr\u003e Open in Workbench\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "  \u003ctd style=\"text-align: center\"\u003e\n",
+        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/knowledge_engine_pinecone.ipynb\"\u003e\n",
+        "      \u003cimg src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"\u003e\u003cbr\u003e View on GitHub\n",
+        "    \u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
+        "\u003c/table\u003e"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "30QXShFkFmRF"
+      },
+      "source": [
+        "## Overview\n",
+        "This notebook demonstrates how to bring your Pinecone instance for use with the Vertex AI Knowledge Engine to serve as the vector store of choice."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_GTk4dqDE7rO"
+      },
+      "source": [
+        "## Get Started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "PAogwmMGFKaH"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": true,
+        "id": "etbFDw4pDysS"
+      },
+      "outputs": [],
+      "source": [
+        "!pip3 install --force-reinstall google-cloud-aiplatform"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": true,
+        "id": "GetEr1yEVs6S"
+      },
+      "outputs": [],
+      "source": [
+        "!pip3 install --upgrade google-cloud\n",
+        "!pip3 install --upgrade google-cloud-secret-manager\n",
+        "!pip install --upgrade \"pinecone-client[grpc]\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "szsH0vSXFAIs"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YuBw1UeGD1D1"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5UegSBo-FBvv"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "Authenticate your environment on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "eeaUnT2vD3n7"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "r9J7aM6oFFJa"
+      },
+      "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). 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": "czcyjyznFRKY"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "PROJECT_NUMBER = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "LOCATION = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1YAOACezFToI"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "C7wL1hpBD5q9"
+      },
+      "outputs": [],
+      "source": [
+        "import vertexai\n",
+        "from vertexai.preview import rag\n",
+        "from vertexai.preview.generative_models import GenerativeModel, Tool\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nkybwP6qDpZd"
+      },
+      "source": [
+        "### Setup Pinecone Instance\n",
+        "\n",
+        "[**OPTIONAL**]\n",
+        "\n",
+        "In this section, we have some helper methods to help you setup your Pinecone instance. This section is not required if you already have a Pinecone instance ready to use."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "qdLGvk9vGgNS"
+      },
+      "source": [
+        "#### Initialize the Pinecone Python client"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "C4Z_6ZheEWQv"
+      },
+      "outputs": [],
+      "source": [
+        "# Set API Key\n",
+        "# Copy this value from your pinecone.io console\n",
+        "PINECONE_API_KEY = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "\n",
+        "from pinecone.grpc import PineconeGRPC as Pinecone\n",
+        "from pinecone import ServerlessSpec, PodSpec\n",
+        "\n",
+        "pc = Pinecone(api_key=PINECONE_API_KEY)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xYftoymKEu42"
+      },
+      "source": [
+        "#### Create a Pinecone instance\n",
+        "\n",
+        "Use the below section to create Pinecone indexes with a spec of your choice (serverless or pod). Read more about Pinecone indexes [here](https://docs.pinecone.io/guides/indexes/create-an-index)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "O6_68wc3E8qb"
+      },
+      "source": [
+        "##### Serverless Index"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Vcn5ZehiEptZ"
+      },
+      "outputs": [],
+      "source": [
+        "# Index Configs\n",
+        "INDEX_NAME = ''  # @param {type:\"string\"}\n",
+        "\n",
+        "# Choose a distance metric\n",
+        "DISTANCE_METRIC = 'cosine'  # @param [\"cosine\", \"euclidean\", \"dotproduct\"] {allow-input: true}\n",
+        "\n",
+        "# This number should match the dimension size of the embedding model you choose\n",
+        "# for your RAG Corpus.\n",
+        "EMBEDDING_DIMENSION_SIZE = 768  # @param {\"type\":\"number\",\"placeholder\":\"768\"}\n",
+        "\n",
+        "CLOUD_PROVIDER = 'gcp'  # @param [\"gcp\", \"aws\", \"azure\"] {allow-input: true}\n",
+        "# Choose the right region for your cloud provider of choice\n",
+        "# Refer https://docs.pinecone.io/guides/indexes/understanding-indexes#cloud-regions\n",
+        "CLOUD_REGION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "\n",
+        "\n",
+        "\n",
+        "# Create the index\n",
+        "pc.create_index(\n",
+        "  name=INDEX_NAME,\n",
+        "  dimension=EMBEDDING_DIMENSION_SIZE,\n",
+        "  metric=DISTANCE_METRIC,\n",
+        "  spec=ServerlessSpec(\n",
+        "    cloud=CLOUD_PROVIDER,\n",
+        "    region=CLOUD_REGION\n",
+        "  ),\n",
+        "  deletion_protection=\"disabled\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8y8CecIrFapb"
+      },
+      "source": [
+        "##### Pod-based Index\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2vV4WsehFsOD"
+      },
+      "outputs": [],
+      "source": [
+        "# Index Configs\n",
+        "INDEX_NAME = ''  # @param {type:\"string\"}\n",
+        "\n",
+        "# Choose a distance metric\n",
+        "DISTANCE_METRIC = 'cosine'  # @param [\"cosine\", \"euclidean\", \"dotproduct\"] {allow-input: true}\n",
+        "\n",
+        "# This number should match the dimension size of the embedding model you choose\n",
+        "# for your RAG Corpus.\n",
+        "EMBEDDING_DIMENSION_SIZE = 768  # @param {\"type\":\"number\",\"placeholder\":\"768\"}\n",
+        "\n",
+        "# Choose the right environment for your cloud provider of choice\n",
+        "# Refer https://docs.pinecone.io/guides/indexes/understanding-indexes#pod-environments\n",
+        "ENVIRONMENT = \"us-central1-gcp\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Choose the pod type\n",
+        "# Refer to https://docs.pinecone.io/guides/indexes/understanding-indexes#pod-based-indexes\n",
+        "POD_TYPE = \"p1.x1\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Explore all the parameters you can play with for creating a pod index by\n",
+        "# following this page:\n",
+        "# https://docs.pinecone.io/reference/api/2024-07/control-plane/create_index\n",
+        "pc.create_index(\n",
+        "  name=INDEX_NAME,\n",
+        "  dimension=EMBEDDING_DIMENSION_SIZE,\n",
+        "  metric=DISTANCE_METRIC,\n",
+        "  spec=PodSpec(\n",
+        "    environment=ENVIRONMENT,\n",
+        "    pod_type=POD_TYPE,\n",
+        "    pods=1,\n",
+        "    metadata_config={\n",
+        "        \"indexed\": [\"file_id\"]  # This field is required for pod-based indexes.\n",
+        "    }\n",
+        "  ),\n",
+        "  deletion_protection=\"disabled\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HopsGrARUogC"
+      },
+      "source": [
+        "### Setup Secret Manager\n",
+        "\n",
+        "[**OPTIONAL**]\n",
+        "\n",
+        "This section helps you add your Pinecone API key to your Google Cloud Secret Manager. This section is not required is you already have a secret with the API key ready to use."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ltqANphsUqfT"
+      },
+      "outputs": [],
+      "source": [
+        "# GCP project ID and the Pinecone API key will be used from the above sections.\n",
+        "# Choose your secret ID\n",
+        "SECRET_ID = \"\"   # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fKt7BzDaVIu9"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import secretmanager\n",
+        "\n",
+        "client = secretmanager.SecretManagerServiceClient()\n",
+        "\n",
+        "# Build the resource name of the parent project.\n",
+        "project_parent = f\"projects/{PROJECT_ID}\"\n",
+        "secret_payload = PINECONE_API_KEY.encode(\"UTF-8\")\n",
+        "\n",
+        "\n",
+        "# Create the secret.\n",
+        "secret_creation_response = client.create_secret(\n",
+        "    request={\n",
+        "        \"parent\": project_parent,\n",
+        "        \"secret_id\": SECRET_ID,\n",
+        "        \"secret\": {\n",
+        "            \"replication\": {\"automatic\": {}},\n",
+        "        },\n",
+        "    }\n",
+        ")\n",
+        "\n",
+        "# Add API key to the secret payload.\n",
+        "# Build the resource name of the parent secret.\n",
+        "secret_parent = client.secret_path(PROJECT_ID, SECRET_ID)\n",
+        "secret_version_response = client.add_secret_version(\n",
+        "    request={\"parent\": secret_parent, \"payload\": {\"data\": secret_payload}}\n",
+        ")\n",
+        "\n",
+        "print(f\"Created secret and added first version: {secret_version_response.name}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rJD9HkbPzehR"
+      },
+      "source": [
+        "##### Give access to Knowledge Engine service account permissions to the secret"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vKZ8voDkx9SU"
+      },
+      "outputs": [],
+      "source": [
+        "def add_secret_accessor_permission(secret_resource_name: str, service_account: str):\n",
+        "    from google.iam.v1 import policy_pb2\n",
+        "\n",
+        "    iam_member = f\"serviceAccount:{service_account}\"\n",
+        "\n",
+        "    # Get the current IAM policy.\n",
+        "    policy = client.get_iam_policy(request={\"resource\": secret_resource_name})\n",
+        "\n",
+        "    # Add the new binding.\n",
+        "    binding = policy_pb2.Binding()\n",
+        "    binding.role = \"roles/secretmanager.secretAccessor\"\n",
+        "    binding.members.append(iam_member)\n",
+        "    policy.bindings.append(binding)\n",
+        "\n",
+        "    # Update the IAM policy.\n",
+        "    client.set_iam_policy(\n",
+        "        request={\"resource\": secret_resource_name, \"policy\": policy}\n",
+        "    )\n",
+        "\n",
+        "    print(f\"Updated IAM policy for secret: {secret_resource_name}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3K03iis-736q"
+      },
+      "source": [
+        "**NOTE**: Skip this below section if this is your first RAG Corpus in the project. This is because your Knowledge Engine service account only gets provisioned during your first resource creation and hence you cannot assign it permissions. The later sections will help you add this after creating your RAG Corpus."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Po9a7k0A8MaA"
+      },
+      "outputs": [],
+      "source": [
+        "# Construct your Knowledge Engine service account name\n",
+        "# Do not update this string since this is the name assigned to your service\n",
+        "# account.\n",
+        "SERVICE_ACCOUNT = (\n",
+        "    f\"service-{PROJECT_NUMBER}@gcp-sa-vertex-rag.iam.gserviceaccount.com\"\n",
+        ")\n",
+        "\n",
+        "add_secret_accessor_permission(\n",
+        "    secret_resource_name=secret_parent, service_account=SERVICE_ACCOUNT\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "K5B_fHBO6dQP"
+      },
+      "source": [
+        "## Create your RAG Corpus"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9EehTOAf678L"
+      },
+      "source": [
+        "### First RAG Corpus Only\n",
+        "\n",
+        "If this is the first RAG Corpus in your project, you might not be able to provide the Knowledge Engine service account access to your secret resource. So this section first creates a RAG Corpus with an empty Pinecone config. With this call, the service account for your project is provisioned.\n",
+        "\n",
+        "Next, it assigns the service account permissions to read your secret. Finally, it updates your RAG Corpus with the Pinecone index name and the secret resource name."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mc7uP9p19niL"
+      },
+      "source": [
+        "##### Create your RAG Corpus without Pinecone information"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "FY7EfGib9TF2"
+      },
+      "outputs": [],
+      "source": [
+        "# Name your corpus\n",
+        "DISPLAY_NAME = \"\"  # @param  {type:\"string\"}\n",
+        "\n",
+        "# Start with empty Pinecone config.\n",
+        "vector_db = rag.Pinecone()\n",
+        "\n",
+        "# Create RAG Corpus\n",
+        "rag_corpus = rag.create_corpus(display_name=DISPLAY_NAME, vector_db=vector_db)\n",
+        "print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rs3sZ2Jh-RRA"
+      },
+      "source": [
+        "##### Grant your Knowledge Engine service account access to your API key secret"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wajzbA_y-dcS"
+      },
+      "outputs": [],
+      "source": [
+        "# Construct your Knowledge Engine service account name\n",
+        "# Do not update this string since this is the name assigned to your service\n",
+        "# account.\n",
+        "SERVICE_ACCOUNT = (\n",
+        "    f\"service-{PROJECT_NUMBER}@gcp-sa-vertex-rag.iam.gserviceaccount.com\"\n",
+        ")\n",
+        "\n",
+        "from google.cloud import secretmanager\n",
+        "\n",
+        "client = secretmanager.SecretManagerServiceClient()\n",
+        "\n",
+        "secret_parent = client.secret_path(PROJECT_ID, SECRET_ID)\n",
+        "add_secret_accessor_permission(\n",
+        "    secret_resource_name=secret_parent, service_account=SERVICE_ACCOUNT\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BnYR4m4k_Qib"
+      },
+      "source": [
+        "##### Call the UpdateRagCorpus API to add the Pinecone index name and API key secret to your RAG Corpus"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "zJQ_zcqm_AsE"
+      },
+      "outputs": [],
+      "source": [
+        "# Full resource name of your created secret version.\n",
+        "# Format: projects/{PROJECT_ID}/secrets/{SECRET_ID}/versions/{VERSION_ID}\n",
+        "SECRET_RESOURCE_NAME = \"\"  # @param  {type:\"string\"}\n",
+        "\n",
+        "# Name of your created Pinecone Index\n",
+        "PINECONE_INDEX_NAME = \"\"  # @param  {type:\"string\"}\n",
+        "\n",
+        "# Construct your updated Pinecone config.\n",
+        "vector_db = rag.Pinecone(\n",
+        "    index_name=PINECONE_INDEX_NAME, api_key=SECRET_RESOURCE_NAME\n",
+        "    )\n",
+        "\n",
+        "updated_rag_corpora = rag.update_corpus(\n",
+        "    corpus_name=rag_corpus.name, vector_db=vector_db\n",
+        "    )\n",
+        "print(f\"Updated RAG Corpus: {rag_corpus.name}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Kx6jHiYABfA8"
+      },
+      "source": [
+        "### Second RAG Corpus Onwards\n",
+        "\n",
+        "In this case, since your service account is already generated, you can directly grant it permissions to access your secret resource containing the Pinecone API key as covered by the steps in the Setup Secret Manager section.\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "h3X36s6dCGJS"
+      },
+      "outputs": [],
+      "source": [
+        "# Full resource name of your created secret version.\n",
+        "# Format: projects/{PROJECT_ID}/secrets/{SECRET_ID}/versions/{VERSION_ID}\n",
+        "SECRET_RESOURCE_NAME = \"\"  # @param  {type:\"string\"}\n",
+        "\n",
+        "# Name of your created Pinecone Index\n",
+        "PINECONE_INDEX_NAME = \"\"  # @param  {type:\"string\"}\n",
+        "\n",
+        "# Name your corpus\n",
+        "DISPLAY_NAME = \"\"  # @param  {type:\"string\"}\n",
+        "\n",
+        "# Construct your Pinecone config.\n",
+        "vector_db = rag.Pinecone(\n",
+        "    index_name=PINECONE_INDEX_NAME, api_key=SECRET_RESOURCE_NAME\n",
+        ")\n",
+        "\n",
+        "# Create RAG Corpus\n",
+        "rag_corpus = rag.create_corpus(display_name=DISPLAY_NAME, vector_db=vector_db)\n",
+        "print(f\"Created RAG Corpus resource: {rag_corpus.name}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "pFXjxTHdDtPG"
+      },
+      "source": [
+        "## Upload a file to your RAG Corpus"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5rrUha3WD0ax"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile test.txt\n",
+        "\n",
+        "Here's a demo for testing RAG upload."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "elxW45WHD6FH"
+      },
+      "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",
+        "\n",
+        "print(f\"Uploaded file to resource: {rag_file.name}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "O8E_RI-VEMQX"
+      },
+      "source": [
+        "## Import files from Google Cloud Storage into your RAG Corpus\n",
+        "\n",
+        "Remember to grant \"Viewer\" access to the Knowledge Engine service account for your Google Cloud Storage bucket.\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JhFEBxcuERDI"
+      },
+      "outputs": [],
+      "source": [
+        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "response = await rag.import_files_async(  # noqa: F704\n",
+        "    corpus_name=rag_corpus.name,\n",
+        "    paths=[GCS_BUCKET],\n",
+        "    chunk_size=512,\n",
+        "    chunk_overlap=50,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WB5spmIRFA5Z"
+      },
+      "outputs": [],
+      "source": [
+        "# Check the files just imported. It may take a few seconds to process the imported files.\n",
+        "list(rag.list_files(corpus_name=rag_corpus.name))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g4mLQIFvFU-T"
+      },
+      "source": [
+        "## Import files from Google Drive\n",
+        "Eligible paths can be https://drive.google.com/drive/folders/{folder_id} or https://drive.google.com/file/d/{file_id}.\n",
+        "\n",
+        "Remember to grant \"Viewer\" access to the Knowledge Engine service account for your Drive folder/files."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5itVoYCqFUbh"
+      },
+      "outputs": [],
+      "source": [
+        "FILE_ID = \"\"  # @param {type:\"string\"}\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",
+        "    chunk_size=1024,\n",
+        "    chunk_overlap=100,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "QEU2be1AFjLM"
+      },
+      "outputs": [],
+      "source": [
+        "# Check the files just imported. It may take a few seconds to process the imported files.\n",
+        "list(rag.list_files(corpus_name=rag_corpus.name))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hS1tzvEVF8ee"
+      },
+      "source": [
+        "## Use your RAG Corpus to add context to your Gemini queries\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xWTbOrMwGKK5"
+      },
+      "outputs": [],
+      "source": [
+        "rag_resource = rag.RagResource(\n",
+        "    rag_corpus=rag_corpus.name,\n",
+        ")\n",
+        "\n",
+        "rag_retrieval_tool = Tool.from_retrieval(\n",
+        "    retrieval=rag.Retrieval(\n",
+        "        source=rag.VertexRagStore(\n",
+        "            rag_resources=[rag_resource],  # Currently only 1 corpus is allowed.\n",
+        "            similarity_top_k=10,\n",
+        "            vector_distance_threshold=0.4,\n",
+        "        ),\n",
+        "    )\n",
+        ")\n",
+        "\n",
+        "rag_model = GenerativeModel(\n",
+        "    \"gemini-1.5-pro-001\",\n",
+        "    tools=[rag_retrieval_tool]\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "roUTPKbBGMoX"
+      },
+      "outputs": [],
+      "source": [
+        "GENERATE_CONTENT_PROMPT = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "response = rag_model.generate_content(GENERATE_CONTENT_PROMPT)\n",
+        "\n",
+        "response"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rblfP8niHOEt"
+      },
+      "source": [
+        "## Cleaning Up\n",
+        "\n",
+        "Clean up the resources created in this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "iyf5xgLfGS3O"
+      },
+      "outputs": [],
+      "source": [
+        "delete_rag_corpus = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if delete_rag_corpus:\n",
+        "    rag.delete_corpus(name=rag_corpus.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EIN5b8_ZHmuq"
+      },
+      "source": [
+        "## API reference\n",
+        "\n",
+        "For more details on RAG corpus/file management and detailed support please visit https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/rag-api"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}