gemini/rag-engine/rag_engine_pinecone.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 Pinecone\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_pinecone.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_pinecone.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_pinecone.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_pinecone.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_pinecone.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_pinecone.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_pinecone.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_pinecone.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_pinecone.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) | [Darshan Mehta](https://github.com/darshanmehta17) |"
      ]
    },
    {
      "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 [Pinecone](https://www.pinecone.io/) as a vector database.\n",
        "\n",
        "For more information, refer to the [official documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/use-pinecone).\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": [],
      "source": [
        "%pip install --upgrade --user --quiet google-cloud-aiplatform google-cloud-secret-manager \"pinecone-client[grpc]\""
      ]
    },
    {
      "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",
        "import vertexai\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\", isTemplate: true}\n",
        "if 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",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "adbe5c6b3549"
      },
      "source": [
        "## (Optional) Setup Pinecone Instance\n",
        "\n",
        "In this section, we have some helper methods to help you setup your Pinecone instance.\n",
        "\n",
        "This section is not required if you already have a Pinecone instance ready to use."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Initialize the Pinecone Python client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "95f06ed30a65"
      },
      "outputs": [],
      "source": [
        "from pinecone import PodSpec, ServerlessSpec\n",
        "from pinecone.grpc import PineconeGRPC as Pinecone\n",
        "\n",
        "# Set API Key\n",
        "# Copy this value from your pinecone.io console\n",
        "PINECONE_API_KEY = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "pc = Pinecone(api_key=PINECONE_API_KEY)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e821e27f013e"
      },
      "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": "9f521dc70361"
      },
      "source": [
        "#### Serverless Index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8b465b3d0eca"
      },
      "outputs": [],
      "source": [
        "# Index Configs\n",
        "INDEX_NAME = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Choose a distance metric\n",
        "DISTANCE_METRIC = (\n",
        "    \"cosine\"  # @param [\"cosine\", \"euclidean\", \"dotproduct\"] {allow-input: true}\n",
        ")\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",
        "\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",
        "# Create the index\n",
        "pc.create_index(\n",
        "    name=INDEX_NAME,\n",
        "    dimension=EMBEDDING_DIMENSION_SIZE,\n",
        "    metric=DISTANCE_METRIC,\n",
        "    spec=ServerlessSpec(cloud=CLOUD_PROVIDER, region=CLOUD_REGION),\n",
        "    deletion_protection=\"disabled\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daec49200393"
      },
      "source": [
        "#### Pod-based Index\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a575d67150b3"
      },
      "outputs": [],
      "source": [
        "# Index Configs\n",
        "INDEX_NAME = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Choose a distance metric\n",
        "DISTANCE_METRIC = (\n",
        "    \"cosine\"  # @param [\"cosine\", \"euclidean\", \"dotproduct\"] {allow-input: true}\n",
        ")\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": "968f535385be"
      },
      "source": [
        "## (Optional) Setup Secret Manager\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": "535dbdfff303"
      },
      "outputs": [],
      "source": [
        "# Google Cloud 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": "217cf0e74015"
      },
      "outputs": [],
      "source": [
        "from google.cloud import secretmanager\n",
        "\n",
        "client = secretmanager.SecretManagerServiceClient()\n",
        "\n",
        "# Create the secret.\n",
        "secret = client.create_secret(\n",
        "    parent=client.common_project_path(PROJECT_ID),\n",
        "    secret_id=SECRET_ID,\n",
        "    secret=secretmanager.Secret(\n",
        "        replication=secretmanager.Replication(\n",
        "            automatic=secretmanager.Replication.Automatic()\n",
        "        )\n",
        "    ),\n",
        ")\n",
        "\n",
        "# Add API key to the secret payload.\n",
        "secret_version = client.add_secret_version(\n",
        "    parent=secret.name,\n",
        "    payload=secretmanager.SecretPayload(data=PINECONE_API_KEY.encode(\"UTF-8\")),\n",
        ")\n",
        "\n",
        "print(f\"Created secret and added first version: {secret_version.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9769f99f15de"
      },
      "source": [
        "## Get Service Account Information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "04e493024500"
      },
      "outputs": [],
      "source": [
        "project_numbers = !gcloud projects list --filter=\"PROJECT_ID={PROJECT_ID}\" --format=\"value(PROJECT_NUMBER)\"\n",
        "PROJECT_NUMBER = project_numbers[0]\n",
        "\n",
        "# Construct your RAG Engine service account name\n",
        "# Do not update this string since this is the name assigned to your service account.\n",
        "SERVICE_ACCOUNT = f\"service-{PROJECT_NUMBER}@gcp-sa-vertex-rag.iam.gserviceaccount.com\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Create a RAG corpus"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b9ad5442bd4e"
      },
      "outputs": [],
      "source": [
        "from vertexai.preview import rag\n",
        "from vertexai.preview.generative_models import GenerativeModel, Tool"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b63b525cee5d"
      },
      "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 RAG 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": "f46f7f34b390"
      },
      "source": [
        "#### Create a RAG Corpus without Pinecone information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9a0e7d8bafeb"
      },
      "outputs": [],
      "source": [
        "# Start with empty Pinecone config.\n",
        "vector_db = rag.Pinecone()\n",
        "\n",
        "# Name your corpus\n",
        "DISPLAY_NAME = \"\"  # @param  {type:\"string\"}\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": "2422ddb6acd4"
      },
      "source": [
        "#### Grant your RAG Engine service account access to your API key secret"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "493b75f6dee0"
      },
      "outputs": [],
      "source": [
        "!gcloud secrets add-iam-policy-binding {secret.name} \\\n",
        "  --member=\"serviceAccount:{SERVICE_ACCOUNT}\" \\\n",
        "  --role=\"roles/secretmanager.secretAccessor\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4e911dd1dcab"
      },
      "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": "84f0c2df5238"
      },
      "outputs": [],
      "source": [
        "# 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(index_name=PINECONE_INDEX_NAME, api_key=secret_version.name)\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": "638af13047a1"
      },
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87459cdb45c6"
      },
      "source": [
        "#### Grant your RAG Engine service account access to your API key secret"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "25f9149874fa"
      },
      "outputs": [],
      "source": [
        "!gcloud secrets add-iam-policy-binding {secret.name} \\\n",
        "  --member=\"serviceAccount:{SERVICE_ACCOUNT}\" \\\n",
        "  --role=\"roles/secretmanager.secretAccessor\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c1396e74d1cb"
      },
      "source": [
        "#### Create a RAG Corpus with Pinecone information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53865b3ea33e"
      },
      "outputs": [],
      "source": [
        "# Name of your created Pinecone Index\n",
        "PINECONE_INDEX_NAME = \"\"  # @param  {type:\"string\"}\n",
        "# Construct your Pinecone config.\n",
        "vector_db = rag.Pinecone(index_name=PINECONE_INDEX_NAME, api_key=secret_version.name)\n",
        "\n",
        "# Name your corpus\n",
        "DISPLAY_NAME = \"\"  # @param  {type:\"string\"}\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": "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 Pinecone 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",
        "    chunk_size=512,\n",
        "    chunk_overlap=50,\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",
        "    chunk_size=1024,\n",
        "    chunk_overlap=100,\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": [
        "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(\"gemini-1.5-flash\", tools=[rag_retrieval_tool])"
      ]
    },
    {
      "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 = rag_model.generate_content(GENERATE_CONTENT_PROMPT)\n",
        "\n",
        "response"
      ]
    },
    {
      "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",
        "    similarity_top_k=10,\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_pinecone.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}