gemini/rag-engine/rag_engine_feature_store.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 Feature Store\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_feature_store.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_feature_store.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_feature_store.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://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/rag_engine_feature_store.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\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_feature_store.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_feature_store.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_feature_store.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_feature_store.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_feature_store.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_feature_store.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 Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview) as a vector database.\n",
        "\n",
        "RAG Engine uses a built-in vector database powered by Spanner to store and manage vector representations of text documents. The vector database retrieves relevant documents based on the documents' semantic similarity to a given query.\n",
        "\n",
        "By integrating Vertex AI Feature Store as an additional vector database, RAG Engine can use Vertex AI Feature Store to handle large data volumes with low latency, which helps to improve the performance and scalability of your RAG applications.\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": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet google-cloud-aiplatform google-cloud-bigquery"
      ]
    },
    {
      "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": "df91814b0f3a"
      },
      "source": [
        "### Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1cc98c707bfc"
      },
      "outputs": [],
      "source": [
        "from google.cloud import bigquery\n",
        "from vertexai.preview import rag\n",
        "from vertexai.preview.generative_models import GenerativeModel, Tool\n",
        "from vertexai.resources.preview import feature_store"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "adbe5c6b3549"
      },
      "source": [
        "## Set up Vertex AI Feature Store\n",
        "\n",
        "Vertex AI Feature Store, a managed cloud-native service, is an essential component of Vertex AI. It simplifies machine learning (ML) feature management and online serving by letting you manage feature data within a BigQuery table or view. This enables low-latency online feature serving.\n",
        "\n",
        "For `FeatureOnlineStore` instances created with optimized online serving, you\n",
        "can take advantage of a vector similarity search to retrieve a list of\n",
        "semantically similar or related entities, which are known as\n",
        "*approximate nearest neighbors*.\n",
        "\n",
        "The following sections show you how to set up a Vertex AI Feature Store instance for your RAG application.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9d54be277ff0"
      },
      "source": [
        "### Create a BigQuery table schema\n",
        "\n",
        "Use the Cloud Console or the code below to create a BigQuery table schema. It\n",
        "must contain the following fields to serve as the data source.\n",
        "\n",
        "| Field name | Data type | Status |\n",
        "|-------------|-----------|--------|\n",
        "| `corpus_id` | `String` | Required |\n",
        "| `file_id` | `String` | Required |\n",
        "| `chunk_id` | `String` | Required |\n",
        "| `chunk_data_type` |`String` | Nullable |\n",
        "| `chunk_data` | `String` | Nullable |\n",
        "| `file_original_uri` | `String` | Nullable |\n",
        "| `embeddings` | `Float` | Repeated |\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c65a3aef9e43"
      },
      "outputs": [],
      "source": [
        "client = bigquery.Client(project=PROJECT_ID)\n",
        "\n",
        "# Define dataset and table name\n",
        "dataset_id = \"input_us_central1\"  # @param {type:\"string\"}\n",
        "table_id = \"rag_source_new\"  # @param {type:\"string\"}\n",
        "\n",
        "schema = [\n",
        "    bigquery.SchemaField(\"corpus_id\", \"STRING\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"file_id\", \"STRING\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"chunk_id\", \"STRING\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"chunk_data_type\", \"STRING\", mode=\"NULLABLE\"),\n",
        "    bigquery.SchemaField(\"chunk_data\", \"STRING\", mode=\"NULLABLE\"),\n",
        "    bigquery.SchemaField(\"file_original_uri\", \"STRING\", mode=\"NULLABLE\"),\n",
        "    bigquery.SchemaField(\"embeddings\", \"FLOAT64\", mode=\"REPEATED\"),\n",
        "]\n",
        "\n",
        "dataset_ref = bigquery.DatasetReference(PROJECT_ID, dataset_id)\n",
        "\n",
        "try:\n",
        "    dataset = client.get_dataset(dataset_ref)\n",
        "    print(f\"Dataset {dataset_id} already exists.\")\n",
        "except Exception:\n",
        "    dataset = bigquery.Dataset(dataset_ref)\n",
        "    dataset.location = \"US\"  # Set the location (optional, adjust if needed)\n",
        "    dataset = client.create_dataset(dataset)\n",
        "    print(f\"Created dataset {dataset.dataset_id}\")\n",
        "\n",
        "table_ref = dataset_ref.table(table_id)\n",
        "table = client.create_table(bigquery.Table(table_ref, schema=schema))\n",
        "print(f\"Created table {PROJECT_ID}.{dataset_id}.{table_id}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f3ed9c71ea1f"
      },
      "outputs": [],
      "source": [
        "BIGQUERY_TABLE = f'bq://{table.full_table_id.replace(\":\", \".\")}'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93b8a1ce48a9"
      },
      "source": [
        "### Provision a `FeatureOnlineStore` instance\n",
        "\n",
        "To enable online serving of features, use the `CreateFeatureOnlineStore` API to set up a `FeatureOnlineStore` instance. If you\n",
        "are provisioning a `FeatureOnlineStore` for the first time, the operation might\n",
        "take approximately five minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ee177c9bc175"
      },
      "outputs": [],
      "source": [
        "FEATURE_ONLINE_STORE_ID = \"your_feature_online_store_id\"  # @param {type: \"string\"}\n",
        "\n",
        "fos = feature_store.FeatureOnlineStore.create_optimized_store(FEATURE_ONLINE_STORE_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "07030f075c06"
      },
      "source": [
        "### Create a `FeatureView` resource\n",
        "\n",
        "To connect the BigQuery table, which stores the feature data source, to\n",
        "the `FeatureOnlineStore` instance, call the `CreateFeatureView` API to create a\n",
        "`FeatureView` resource. When you create a `FeatureView` resource, choose the\n",
        "default distance metric `DOT_PRODUCT_DISTANCE`, which is defined as the\n",
        "negative of the dot product (smaller `DOT_PRODUCT_DISTANCE` indicates higher\n",
        "similarity)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "013dc5a45561"
      },
      "outputs": [],
      "source": [
        "FEATURE_VIEW_ID = \"your_feature_view_id\"  # @param {type: \"string\"}\n",
        "fv = fos.create_feature_view(\n",
        "    name=FEATURE_VIEW_ID,\n",
        "    source=feature_store.utils.FeatureViewVertexRagSource(uri=BIGQUERY_TABLE),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2fdbc1e83aff"
      },
      "outputs": [],
      "source": [
        "# Check that Feature View was created\n",
        "print(fv)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Use Vertex AI Feature Store in RAG Engine\n",
        "\n",
        "After the Feature Store instance is set up, the following\n",
        "sections show you how to set it up as the vector database to use with the RAG\n",
        "application."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cd05469c3e71"
      },
      "source": [
        "### Set the vector database to create a RAG corpus\n",
        "\n",
        "To create the RAG corpus, you must use `FEATURE_VIEW_RESOURCE_NAME`. The\n",
        "RAG corpus is created and automatically associated with the\n",
        "Vertex AI Feature Store instance.\n",
        "\n",
        "RAG APIs use the generated `rag_corpus_id` to handle the data upload to the Vertex AI Feature Store\n",
        "instance and to retrieve relevant contexts from the `rag_corpus_id`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53865b3ea33e"
      },
      "outputs": [],
      "source": [
        "vector_db = rag.VertexFeatureStore(resource_name=fv.resource_name)\n",
        "\n",
        "# Name your corpus\n",
        "DISPLAY_NAME = \"Feature Store Corpus\"  # @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": "code",
      "execution_count": null,
      "metadata": {
        "id": "834931d4b187"
      },
      "outputs": [],
      "source": [
        "rag_corpus"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93a3296647a2"
      },
      "source": [
        "## Import files into the BigQuery table using the RAG API\n",
        "\n",
        "Use the `ImportRagFiles` API to import files from Google Cloud Storage or\n",
        "Google Drive into the BigQuery table of the Vertex AI Feature Store\n",
        "instance. The files are embedded and stored in the BigQuery table.\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 = \"cloud-samples-data/gen-app-builder/search/cymbal-bank-employee\"  # @param {type:\"string\"}\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": "b524b1e3fa56"
      },
      "source": [
        "### Run a synchronization process to construct a `FeatureOnlineStore` index {:#run-sync-process}\n",
        "\n",
        "After uploading your data into the BigQuery table, run a\n",
        "synchronization process to make your data available for online serving. You must\n",
        "generate a `FeatureOnlineStore` index using the `FeatureView`, and the\n",
        "synchronization process might take 20 minutes to complete.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "479d2e1f7725"
      },
      "outputs": [],
      "source": [
        "feature_view_sync = fv.sync()\n",
        "feature_view_sync"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "63ef904c92b4"
      },
      "outputs": [],
      "source": [
        "# Optional: Wait for sync to complete\n",
        "feature_view_sync.wait()"
      ]
    },
    {
      "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_retrieval_tool = Tool.from_retrieval(\n",
        "    retrieval=rag.Retrieval(\n",
        "        source=rag.VertexRagStore(\n",
        "            rag_resources=[\n",
        "                rag.RagResource(\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",
        ")\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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8b9b84c58cb0"
      },
      "outputs": [],
      "source": [
        "response.text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2899daa12fac"
      },
      "source": [
        "## Using other generation APIs 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",
        "response = rag.retrieval_query(\n",
        "    rag_resources=[\n",
        "        rag.RagResource(\n",
        "            rag_corpus=rag_corpus.name,  # Currently only 1 corpus is allowed.\n",
        "        )\n",
        "    ],\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_feature_store.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}