{
  "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": [
        "# Building a Gen AI RAG application with Vertex AI Feature Store and BigQuery\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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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%2Fuse-cases%2Fretrieval-augmented-generation%2Frag_qna_with_bq_and_featurestore.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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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/use-cases/retrieval-augmented-generation/rag_qna_with_bq_and_featurestore.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",
        "| Authors | [Elia Secchi](https://github.com/eliasecchig) [Lorenzo Spataro](https://github.com/lspataroG) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook guides you through building a low-latency vector search system for your Gen AI application using [BigQuery Vector Search](https://cloud.google.com/bigquery/docs/vector-search-intro) and [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/latest/overview). We'll leverage the [`BigQueryVectorStore` LangChain integration]([https://github.com/langchain-ai/langchain-google/blob/main/libs/community/langchain_google_community/bq_storage_vectorstores/featurestore.py#L33]) and [`VertexFSVectorStore` LangChain integration]([https://github.com/langchain-ai/langchain-google/blob/main/libs/community/langchain_google_community/bq_storage_vectorstores/bigquery.py#L26]) to streamline this process.\n",
        "\n",
        "Vertex AI Feature Store seamlessly integrates with BigQuery, providing a unified data storage and flexible vector search options:\n",
        "\n",
        "- **BigQuery Vector Search**: with **`BigQueryVectorStore`** LangChain class, ideal for batch retrieval and prototyping, as it requires no infrastructure setup.\n",
        "- **Feature Store Online Store**: with **`VertexFSVectorStore`** LangChain class, enables low-latency retrieval with manual or scheduled data sync. Perfect for production-ready user-facing Gen AI applications. \n",
        "\n",
        "As part of this notebook you will learn how to:\n",
        "1. Ingest data and embedding using BigQuery Vector Search with the class `BigQueryVectorStore`\n",
        "2. Perform retrieval leveraging BigQuery Vector Search with the class `BigQueryVectorStore`\n",
        "3. Transition to Vertex AI Feature Store with the class `VertexFSVectorStore` for low-latency retrieval\n",
        "4. Understand pros and cons of both options through a performance deep dive"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b2556b5937ab"
      },
      "source": [
        "![bq_fs_diagram_journey.png](https://storage.googleapis.com/github-repo/generative-ai/gemini/use-cases/retrieval-augmented-generation/bq_fs_diagram_journey.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet google-cloud-aiplatform \"langchain-google-vertexai\" \"langchain-google-community[featurestore]\" pypdf==4.2.0"
      ]
    },
    {
      "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>"
      ]
    },
    {
      "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": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## Getting started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "854025f0c407"
      },
      "outputs": [],
      "source": [
        "from langchain.chains import RetrievalQA\n",
        "from langchain.globals import set_debug\n",
        "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
        "from langchain_community.document_loaders import PyPDFLoader\n",
        "from langchain_google_community import BigQueryVectorStore, VertexFSVectorStore\n",
        "from langchain_google_vertexai import VertexAI, VertexAIEmbeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_68282UqlxrR"
      },
      "outputs": [],
      "source": [
        "DATASET = \"sample_app\"  # @param {type:\"string\"}\n",
        "TABLE = \"fixmycar\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LDSM2ZQa8eBs"
      },
      "source": [
        "## Add documents to `BigQueryVectorStore`\n",
        "\n",
        "This step ingests and parse PDF documents, split them, generate embeddings and add the embeddings to the vector store. The document corpus used as dataset is a collection of owners car manual."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dIML1Ttj5nTk"
      },
      "source": [
        "**Summary steps**\n",
        "- Create text embeddings: LangChain `VertexAIEmbeddings`\n",
        "- Ingest PDF files: LangChain `PyPDFLoader`\n",
        "- Chunk documents: LangChain `TextSplitter`\n",
        "- Create Vector Store: LangChain  `VertexAIFeatureStore`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q3hRmLiY2O0w"
      },
      "source": [
        "### Create the Vertex AI Embedding model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2tRYGd_Mv_ix"
      },
      "outputs": [],
      "source": [
        "embedding_model = VertexAIEmbeddings(\n",
        "    model_name=\"text-embedding-005\", project=PROJECT_ID\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pKM_QIU0WPc3"
      },
      "source": [
        "### Ingest PDF file\n",
        "\n",
        "The document is hosted on Cloud Storage bucket (at `gs://github-repo/generative-ai/sample-apps/fixmycar/cymbal-starlight-2024.pdf`) and LangChain provides a convenient document loader [`PyPDFLoader`](https://python.langchain.com/docs/modules/data_connection/document_loaders/pdf/) to load documents from pdfs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UNEClQQjjRcc"
      },
      "outputs": [],
      "source": [
        "GCS_BUCKET_DOCS = \"github-repo/generative-ai/sample-apps/fixmycar\"\n",
        "\n",
        "# Copy the file to the current path\n",
        "!gsutil cp \"gs://$GCS_BUCKET_DOCS/*.pdf\" ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7kYw66EohYza"
      },
      "outputs": [],
      "source": [
        "# Ingest PDF files\n",
        "loader = PyPDFLoader(\"cymbal-starlight-2024.pdf\")\n",
        "documents = loader.load()\n",
        "\n",
        "# Add document name and source to the metadata\n",
        "for document in documents:\n",
        "    doc_md = document.metadata\n",
        "    document_name = doc_md[\"source\"].split(\"/\")[-1]\n",
        "    # derive doc source from Document loader\n",
        "    doc_source_prefix = \"/\".join(GCS_BUCKET_DOCS.split(\"/\")[:3])\n",
        "    doc_source_suffix = \"/\".join(doc_md[\"source\"].split(\"/\")[4:-1])\n",
        "    source = f\"{doc_source_prefix}/{doc_source_suffix}\"\n",
        "    document.metadata = {\"source\": source, \"document_name\": document_name}\n",
        "\n",
        "print(f\"# of documents loaded (pre-chunking) = {len(documents)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j8J7hQkVmkB1"
      },
      "source": [
        "Verify document metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jn3A13Nfmihf"
      },
      "outputs": [],
      "source": [
        "documents[0].metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oZA9HDA0mnWW"
      },
      "source": [
        "## Chunk documents - `TextSplitter`\n",
        "\n",
        "Split the documents to smaller chunks. When splitting the document, ensure a few chunks can fit within the context length of LLM."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n2QuilhMmxgA"
      },
      "outputs": [],
      "source": [
        "# split the documents into chunks\n",
        "text_splitter = RecursiveCharacterTextSplitter(\n",
        "    chunk_size=1000,\n",
        "    chunk_overlap=50,\n",
        "    separators=[\"\\n\\n\", \"\\n\", \".\", \"!\", \"?\", \",\", \" \", \"\"],\n",
        ")\n",
        "doc_splits = text_splitter.split_documents(documents)\n",
        "\n",
        "# Add chunk number to metadata\n",
        "for idx, split in enumerate(doc_splits):\n",
        "    split.metadata[\"chunk\"] = idx\n",
        "\n",
        "print(f\"# of documents = {len(doc_splits)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fv91BaKCmxP0"
      },
      "outputs": [],
      "source": [
        "doc_splits[0].metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8QC-w_t3m9Fe"
      },
      "source": [
        "## Configure `BigQueryVectorStore` as Vector Store\n",
        "\n",
        "We will start the journey using BigQuery Vector Store as it requires no startup time, making the class great for prototyping.\n",
        "\n",
        "You can initialize the class by providing:\n",
        "- `project_id`\n",
        "- `location`\n",
        "- `dataset_name`\n",
        "- `table_name`\n",
        "\n",
        "The table will be used to store embeddings and metadata. You can also point to an existing table. The class will use [BigQuery Vector Search](https://cloud.google.com/bigquery/docs/vector-search-intro) to perform vector search.\n",
        "\n",
        "See [here](https://github.com/langchain-ai/langchain-google/blob/main/libs/community/langchain_google_community/bq_storage_vectorstores/bigquery.py#L26) for the full list of parameters of the class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gLxHuLY1lxrS"
      },
      "outputs": [],
      "source": [
        "bq_store = BigQueryVectorStore(\n",
        "    project_id=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    dataset_name=DATASET,\n",
        "    table_name=TABLE,\n",
        "    embedding=embedding_model,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17cgWCFRlxrS"
      },
      "source": [
        "### Add documents to the store\n",
        "\n",
        "Note: If you have precomputed embeddings, you can add text, embeddings and potential metadata using the method `add_texts_with_embeddings`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_TXggBP9lxrT"
      },
      "outputs": [],
      "source": [
        "doc_ids = bq_store.add_documents(doc_splits)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NwTSeEsGlxrT"
      },
      "source": [
        "Verify the `BigQueryVectorSearch` with similarity search"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0YYiWAOrlxrT"
      },
      "outputs": [],
      "source": [
        "bq_store.similarity_search(\n",
        "    \"What should I do when I call the emergency roadside assistance?\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4Zz3PUFlxrT"
      },
      "source": [
        "### Get a langchain retriever\n",
        "The retriever will be used in a LangChain Chain to find the most similar documents for a given query."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F9mSL5o-lxrT"
      },
      "outputs": [],
      "source": [
        "langchain_retriever = bq_store.as_retriever()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lWroxQMflxrT"
      },
      "source": [
        "### Compose a LangChain Chain\n",
        "\n",
        "We are going to use the [`RetrievalQA` chain](https://python.langchain.com/docs/modules/chains/popular/vector_db_qa)\n",
        "There are several different chain types available, listed [here](https://docs.langchain.com/docs/components/chains/index_related_chains)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-eZsLYIPlxrT"
      },
      "outputs": [],
      "source": [
        "# Set high verbosity\n",
        "set_debug(True)\n",
        "\n",
        "llm = VertexAI(model_name=\"gemini-2.0-flash\")\n",
        "\n",
        "search_query = \"What should I do when calling the emergency roadside assistance?\"  # @param {type:\"string\"}\n",
        "\n",
        "retrieval_qa = RetrievalQA.from_chain_type(\n",
        "    llm=llm, chain_type=\"stuff\", retriever=langchain_retriever\n",
        ")\n",
        "response = retrieval_qa.invoke(search_query)\n",
        "print(\"\\n################ Final Answer ################\\n\")\n",
        "print(response[\"result\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ltOScj77lxrT"
      },
      "source": [
        "## Low latency Vector Search with FeatureStore with `VertexFSVectorStore`\n",
        "\n",
        "We are now ready to perform low latency serving with Feature Store!\n",
        "\n",
        "To do that, you can simply use the method `.to_vertex_fs_vector_store()`, to get a `VertexFSVectorStore` object.\n",
        "\n",
        "See the [class definition](https://github.com/langchain-ai/langchain-google/blob/main/libs/community/langchain_google_community/bq_storage_vectorstores/featurestore.py#L33) for all the parameters you can use.\n",
        "Moving back to `BigQueryVectorStore` is equivalently easy with the `.to_bq_vector_search()` method.\n",
        "\n",
        "Note: Any method we run earlier can be equivalently called on both `BigQueryVectorStore` and `VertexFSVectorStore`. For instance it is possible to add new documents to an instance of `VertexFSVectorStore` as both stores share the same underlying BQ source."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fmocdGDVlxrT"
      },
      "outputs": [],
      "source": [
        "vertex_fs = bq_store.to_vertex_fs_vector_store()  # pass optional parameters here"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UKDhZDh4lxrT"
      },
      "source": [
        "You can also initialize the `VertexFSVectorStore` class directly"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OjWdX1x7lxrT"
      },
      "outputs": [],
      "source": [
        "vertex_fs = VertexFSVectorStore(\n",
        "    project_id=PROJECT_ID,\n",
        "    location=LOCATION,\n",
        "    dataset_name=DATASET,\n",
        "    table_name=TABLE,\n",
        "    embedding=embedding_model,\n",
        "    # pass optional parameters here\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WpPh-_F9lxrT"
      },
      "source": [
        "#### Kick off a synchronization process\n",
        "\n",
        "We use the `sync_data` method to synchronize the data from BigQuery to the Feature Online Store, to achieve low latency serving.\n",
        "\n",
        "> Note: The first synchronization process will take around ~20 minutes because of Feature Online Store creation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TDaEy4PQlxrT"
      },
      "outputs": [],
      "source": [
        "vertex_fs.sync_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a07a6f67360"
      },
      "source": [
        "When in a production environment, you can also use the `cron_schedule` class parameter to setup an automatic scheduled synchronization. For example:\n",
        "```python\n",
        "store = VertexFSVectorStore(cron_schedule=\"TZ=America/Los_Angeles 00 13 11 8 *\", ...)\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ts3D7GSrRPJa"
      },
      "outputs": [],
      "source": [
        "vertex_fs.similarity_search(\"Hello world\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TMMFTX9ClxrT"
      },
      "source": [
        "You can monitor the synchronization process from the Google Cloud Console: [Vertex AI Feature Store Tab](https://console.cloud.google.com/vertex-ai/feature-store/online-stores)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vpREPwQilxrX"
      },
      "source": [
        "#### Serve with Feature Online Store\n",
        "\n",
        "You are now ready to use Vertex AI Feature Store as part of your chain through a retriever object!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nYCnP8nlxrX"
      },
      "outputs": [],
      "source": [
        "langchain_retriever = vertex_fs.as_retriever()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3o60735MlxrX"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "results = langchain_retriever.invoke(\"Leaks under the vehicle\")\n",
        "results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VcyOdT0tlxrX"
      },
      "source": [
        "### Filtering by metadata\n",
        "\n",
        "It is possible to post-filter results by metadata by passing the filter parameter to any search method\n",
        "\n",
        "`VertexFSVectorStore` also support metadata filter while performing search, for this to work:\n",
        "- the `filter_columns` parameter must be passed to `VertexFSVectorStore` when the online feature store feature view is created (first time the class is initialised with a given online store name and feature view name).\n",
        "\n",
        "- the `string_filters` parameter must be passed to any search method. Note only string fields are supported at the moment. See [here](https://github.com/googleapis/python-aiplatform/blob/8a4a41afe47aaff2f69a73e5011b34bcba5cd2e9/google/cloud/aiplatform_v1beta1/types/feature_online_store_service.py#L345)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "swpQFcdclxrX"
      },
      "outputs": [],
      "source": [
        "vertex_fs.similarity_search(search_query, filter={\"chunk\": 28})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iR9K3mJO0wzL"
      },
      "source": [
        "# When should I use which class? A performance deep dive\n",
        "\n",
        "We precompute the embedding so that we exclude that latency from the equation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FDTEgfNqbzWd"
      },
      "outputs": [],
      "source": [
        "my_embedding = embedding_model.embed(search_query)[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJj-WHVIlxrX"
      },
      "source": [
        "## Batch search with `BigQueryVectorStore`\n",
        "\n",
        "For some use cases it is necessary to run batch searches (ie. when running a retrieval evaluation).\n",
        "\n",
        "Leveraging the power of scale of BigQuery, we can run efficient batch searches using the `BigQueryVectorStore` which offers a specialized `batch_search` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m3m60d0SlxrX"
      },
      "outputs": [],
      "source": [
        "results = bq_store.batch_search(\n",
        "    embeddings=None,  # can pass embeddings or\n",
        "    queries=[\"search_query\", \"search_query\"],  # can pass queries\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCZyUzwwbKp6"
      },
      "source": [
        "### Batch search with 10.000 embeddings\n",
        "\n",
        "We can run 10.000 batched searches with `BigQueryVectorStore` in ~20 seconds!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJ3lSPXibOhW"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "fake_embeddings = [my_embedding] * 10000\n",
        "results = bq_store.batch_search(embeddings=fake_embeddings)\n",
        "results[:2]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9KKmPQHf2mmp"
      },
      "source": [
        "## Low latency serving with Feature Store\n",
        "If you are instead looking at powering an online application, Vertex Feature Store might be a good solution as it offers low latency serving.\n",
        "\n",
        "We run a small load test composed by 10 requests to demonstrate the latency reduction."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sQs-bRdH1FEv"
      },
      "source": [
        "### BigQuery single request"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d-WHG6MT06LK"
      },
      "outputs": [],
      "source": [
        "%%timeit -r10\n",
        "bq_store.similarity_search_by_vector(my_embedding)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "56IN1HUp1cEW"
      },
      "source": [
        "### Feature Store single request"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sMDDxBc61bG6"
      },
      "outputs": [],
      "source": [
        "%%timeit -r10\n",
        "vertex_fs.similarity_search_by_vector(my_embedding)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o5-pmxyyahB8"
      },
      "source": [
        "Feature store is faster than BigQuery over single requests!\n",
        "\n",
        "\n",
        "> Note: for server side latency estimate we suggest leveraging the Feature Store dashboards.\n",
        "You can do it by:\n",
        "1. Visit the [Feature Store console](https://console.cloud.google.com/vertex-ai/feature-store/online-stores)\n",
        "2. Click on your newly created Feature Online Store\n",
        "3. Scroll down to \"Serving Latency\" dashboard!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B5PdyfQDlxrX"
      },
      "source": [
        "# Appendix"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wkl6IaO5lxrY"
      },
      "source": [
        "### Get documents by ID\n",
        "\n",
        "For both Vector Stores you can also use the function `get_documents` to retrieve a set of documents given a document ID:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N0Omhx2tlxrY"
      },
      "outputs": [],
      "source": [
        "vertex_fs.get_documents(ids=[\"my_id1\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vuB-CA4glxrY"
      },
      "source": [
        "### Remove documents by ID\n",
        "\n",
        "You can also use the function `delete` to remove a set of documents given a document ID:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xvymq39ClxrY"
      },
      "outputs": [],
      "source": [
        "vertex_fs.delete(ids=[\"my_id1\", \"my_id2\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q3Yes1y-DXPM"
      },
      "source": [
        "## Maximal marginal relevance search\n",
        "\n",
        "You can also use [maximal marginal relevance search](https://python.langchain.com/v0.1/docs/modules/model_io/prompts/example_selectors/mmr/) for both Vector Stores."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "unAvu2mpDWju"
      },
      "outputs": [],
      "source": [
        "mmr_retriever = vertex_fs.as_retriever(search_type=\"mmr\")\n",
        "mmr_retriever.invoke(\"Lane departure warning?\")[1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21f2a7432655"
      },
      "outputs": [],
      "source": [
        "# Delete BigQuery dataset. Uncomment and run the command below if you want to delete the BigQuery set.\n",
        "# from google.cloud import bigquery\n",
        "# Do this only if the dataset is created for this demo.\n",
        "# dataset = f\"{PROJECT_ID}.{DATASET_ID}\"\n",
        "# dataset_object = bigquery.Dataset(dataset)\n",
        "# client.delete_dataset(dataset_object, delete_contents=True, not_found_ok=True)\n",
        "\n",
        "vertex_fs.feature_view.delete()\n",
        "vertex_fs.online_store.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "rag_qna_with_bq_and_featurestore.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
