part5_RAG_VertexSearch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bc6fc05e-9237-4af1-aaf0-db2e6756f861",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2023 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",
   "id": "15f3fc1c-7e0b-4df2-b03d-424e8eb23c83",
   "metadata": {},
   "source": [
    "# Retrieval Augmented Generation (RAG) with Vertex AI Search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "575d4a65-9354-424b-8e85-4373dcc8c30e",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Before you begin\n",
    "\n",
    "### Install dependencies\n",
    "\n",
    "We will first need to install and upgrade the Google Cloud Discovery Engine library. Once the library is installed, restart the kernel so that you can use the updated packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cf4fb700-fb8f-4fba-9bba-4f7ee1ef232b",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install -q --upgrade google-cloud-discoveryengine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa55de9-cd40-4a95-a371-250f27bc914f",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Import libraries and set environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a20fd3ef-f378-4c5b-9010-cb704d0c4f86",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "from google.api_core.client_options import ClientOptions\n",
    "from google.cloud import discoveryengine\n",
    "\n",
    "from google.cloud import discoveryengine_v1alpha as discoveryengine\n",
    "from google.api_core.client_options import ClientOptions\n",
    "\n",
    "import socket\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f40d712d-39ae-4846-925b-7b772d05e397",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Project ID: my-project-0004-346516\n",
      "GCS Bucket URI: gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1\n"
     ]
    }
   ],
   "source": [
    "UNIQUE_PREFIX = socket.gethostname()\n",
    "UNIQUE_PREFIX = re.sub('[^A-Za-z0-9]+', '', UNIQUE_PREFIX)[:-6]\n",
    "\n",
    "# Cloud project id.\n",
    "PROJECT_IDS = !(gcloud config get-value core/project)\n",
    "PROJECT_ID = PROJECT_IDS[0]  # @param {type:\"string\"}\n",
    "\n",
    "# The Cloud Storage bucket for storing experiments output.\n",
    "# Remove prefix gs://, e.g. foo_bucket.\n",
    "GCS_BUCKET_LOCATION = \"asia-southeast1\"\n",
    "REGION = 'asia-southeast1'\n",
    "\n",
    "PROJECT_ID_PREFIX = PROJECT_ID[:-6]\n",
    "GCS_BUCKET_URI = f\"gs://hp-bk-{PROJECT_ID_PREFIX}-{UNIQUE_PREFIX}-{REGION}\"\n",
    "\n",
    "# print variables for verification\n",
    "print(f\"Project ID: {PROJECT_ID}\")\n",
    "print(f\"GCS Bucket URI: {GCS_BUCKET_URI}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59759092-c1ba-431f-9afa-733602acfc22",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7aa77d52-b4dd-4402-a7aa-ec1bf0b00afe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import vertexai\n",
    "from vertexai.preview.language_models import TextGenerationModel\n",
    "from vertexai.preview.generative_models import GenerativeModel, Part\n",
    "\n",
    "LOCATION = REGION = 'us-central1'\n",
    "\n",
    "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
    "parameters = {\n",
    "    \"candidate_count\": 1,\n",
    "    \"max_output_tokens\": 1024,\n",
    "    \"temperature\": 1,\n",
    "    \"top_k\": 40\n",
    "}\n",
    "\n",
    "# @st.cache_resource\n",
    "def get_model():\n",
    "    generation_model = TextGenerationModel.from_pretrained(\"text-bison@002\")\n",
    "    return generation_model\n",
    "\n",
    "\n",
    "def get_text_generation(prompt=\"\", **parameters):\n",
    "    generation_model = get_model()\n",
    "    response = generation_model.predict(prompt=prompt, **parameters)\n",
    "\n",
    "    return response.text\n",
    "\n",
    "def generate_palm_unicorn_v1(input_prompt):\n",
    "    \n",
    "    model = TextGenerationModel.from_pretrained(\"text-unicorn@001\")\n",
    "\n",
    "    response = model.predict(\n",
    "        input_prompt,\n",
    "        **parameters\n",
    "    )\n",
    "    print(f\"Response from Model: {response.text}\")\n",
    "    \n",
    "    return(response.text)\n",
    "\n",
    "def generate_palm_bison32k(input_prompt):\n",
    "    \n",
    "    model = TextGenerationModel.from_pretrained(\"text-bison-32k\")\n",
    "\n",
    "    response = model.predict(\n",
    "        input_prompt,\n",
    "        **parameters\n",
    "    )\n",
    "    print(f\"Response from Model: {response.text}\")\n",
    "    \n",
    "    return(response.text)\n",
    "\n",
    "def generate(input_prompt):\n",
    "    model = GenerativeModel(\"gemini-ultra\")\n",
    "    responses = model.generate_content(\n",
    "        input_prompt ,\n",
    "    generation_config={\n",
    "        \"max_output_tokens\": 2048,\n",
    "        \"temperature\": 0.2,\n",
    "        \"top_p\": 1,\n",
    "        \"top_k\": 32\n",
    "    },\n",
    "        safety_settings=[],\n",
    "        stream=True,\n",
    "    )\n",
    "    \n",
    "    all_response  = []\n",
    "    \n",
    "    for response in responses:\n",
    "        # print(response.text, end=\"\")\n",
    "        all_response.append(response.text)\n",
    "    \n",
    "    # print (all_response)\n",
    "    \n",
    "    return(\" \".join(all_response))\n",
    "    \n",
    "\n",
    "def generate_pro(input_prompt):\n",
    "    model = GenerativeModel(\"gemini-pro\")\n",
    "    responses = model.generate_content(\n",
    "    input_prompt,\n",
    "    generation_config={\n",
    "        \"max_output_tokens\": 2048,\n",
    "        \"temperature\": 0.2,\n",
    "        \"top_p\": 1\n",
    "    },stream=True,)\n",
    "    \n",
    "    all_response  = []\n",
    "    \n",
    "    for response in responses:\n",
    "        all_response.append(response.text)\n",
    "    \n",
    "    # print (all_response)\n",
    "    \n",
    "    return(\" \".join(all_response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "69facd95-3cb7-436f-9c5c-e4d790f37ca8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Helper Function to create data store\n",
    "def create_data_store(\n",
    "    project_id: str, location: str, data_store_name: str, data_store_id: str\n",
    "):\n",
    "    # Create a client\n",
    "    client_options = (\n",
    "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
    "        if location != \"global\"\n",
    "        else None\n",
    "    )\n",
    "    client = discoveryengine.DataStoreServiceClient(client_options=client_options)\n",
    "\n",
    "    # Initialize request argument(s)\n",
    "    data_store = discoveryengine.DataStore(\n",
    "        display_name=data_store_name,\n",
    "        industry_vertical=\"GENERIC\",\n",
    "        content_config=\"CONTENT_REQUIRED\",\n",
    "    )\n",
    "\n",
    "    request = discoveryengine.CreateDataStoreRequest(\n",
    "        parent=discoveryengine.DataStoreServiceClient.collection_path(\n",
    "            project_id, location, \"default_collection\"\n",
    "        ),\n",
    "        data_store=data_store,\n",
    "        data_store_id=data_store_id,\n",
    "    )\n",
    "    operation = client.create_data_store(request=request)\n",
    "\n",
    "    # Make the request\n",
    "    # The try block is necessary to prevent execution from haulting due to an error being thrown when the datastore takes a while to instantiate\n",
    "    try:\n",
    "        response = operation.result(timeout=90)\n",
    "    except:\n",
    "        print(\"long-running operation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "707b0999-27c8-4a67-99e3-466d0eb582e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Helper Function to import documents from GCS bucket into datastore\n",
    "def import_documents(\n",
    "    project_id: str,\n",
    "    location: str,\n",
    "    data_store_id: str,\n",
    "    gcs_uri: str,\n",
    "):\n",
    "    # Create a client\n",
    "    client_options = (\n",
    "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
    "        if location != \"global\"\n",
    "        else None\n",
    "    )\n",
    "    client = discoveryengine.DocumentServiceClient(client_options=client_options)\n",
    "\n",
    "    # The full resource name of the search engine branch.\n",
    "    # e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}\n",
    "    parent = client.branch_path(\n",
    "        project=project_id,\n",
    "        location=location,\n",
    "        data_store=data_store_id,\n",
    "        branch=\"default_branch\",\n",
    "    )\n",
    "\n",
    "    source_documents = [f\"{gcs_uri}/*\"]\n",
    "\n",
    "    request = discoveryengine.ImportDocumentsRequest(\n",
    "        parent=parent,\n",
    "        gcs_source=discoveryengine.GcsSource(\n",
    "            input_uris=source_documents, data_schema=\"content\"\n",
    "        ),\n",
    "        # Options: `FULL`, `INCREMENTAL`\n",
    "        reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,\n",
    "    )\n",
    "\n",
    "    # Make the request\n",
    "    operation = client.import_documents(request=request)\n",
    "\n",
    "    response = operation.result()\n",
    "\n",
    "    # Once the operation is complete,\n",
    "    # get information from operation metadata\n",
    "    metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)\n",
    "\n",
    "    # Handle the response\n",
    "    return operation.operation.name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "832e95c1-6439-4c62-a112-b0c5d1ed3d47",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 1. Configure a Datastore for our documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f92fe42f-f917-4f9e-9072-6e3105fcb21d",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Upload the book PDFs onto Google Cloud Storage (GCS) bucket\n",
    "In this section, we will be doing the following:\n",
    "1. Create a Cloud Storage bucket\n",
    "2. Upload the PDFs of the Harry Potter series into the Cloud Storage bucket\n",
    "3. Verify that the PDFs have been uploaded successfully\n",
    "4. Create an empty datastore and import the PDFs documents into the datastore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e228f9be-b258-4e9f-abf9-7eca814d1b92",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/...\n",
      "Copying file://./books/Book1_HarryPotter_and_the_Sorcerers_Stone_pg15.pdf [Content-Type=application/pdf]...\n",
      "Copying file://./books/Book1_The_Sorcerers_Stone.pdf [Content-Type=application/pdf]...\n",
      "Copying file://./books/Book2_The_Chamber_of_Secrets.pdf [Content-Type=application/pdf]...\n",
      "Copying file://./books/Book3_The_Prisoner_of_Azkaban.pdf [Content-Type=application/pdf]...\n",
      "/ [4 files][  4.0 MiB/  4.0 MiB]  450.4 KiB/s                                   \n",
      "==> NOTE: You are performing a sequence of gsutil operations that may\n",
      "run significantly faster if you instead use gsutil -m cp ... Please\n",
      "see the -m section under \"gsutil help options\" for further information\n",
      "about when gsutil -m can be advantageous.\n",
      "\n",
      "Copying file://./books/Book4_The_Goblet_of_Fire.pdf [Content-Type=application/pdf]...\n",
      "Copying file://./books/Book5_The_Order_of_the_Phoenix.pdf [Content-Type=application/pdf]...\n",
      "Copying file://./books/Book6_The_HalfBlood_Prince.pdf [Content-Type=application/pdf]...\n",
      "Copying file://./books/Book7_The_Deathly_Hallows.pdf [Content-Type=application/pdf]...\n",
      "\\ [8 files][ 21.9 MiB/ 21.9 MiB]  615.5 KiB/s                                   \n",
      "Operation completed over 8 objects/21.9 MiB.                                     \n"
     ]
    }
   ],
   "source": [
    "# Create a Cloud Storage Bucket\n",
    "!gcloud storage buckets create $GCS_BUCKET_URI --location=$GCS_BUCKET_LOCATION\n",
    "\n",
    "# Upload the PDFs located in the books/ directory into the GCS bucket that you created\n",
    "!gsutil cp -r ./books/* $GCS_BUCKET_URI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7157d0d3-af2f-4d57-9f32-de1ad57e7ad2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book1_HarryPotter_and_the_Sorcerers_Stone_pg15.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book1_The_Sorcerers_Stone.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book2_The_Chamber_of_Secrets.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book3_The_Prisoner_of_Azkaban.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book4_The_Goblet_of_Fire.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book5_The_Order_of_the_Phoenix.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book6_The_HalfBlood_Prince.pdf\n",
      "gs://hp-books-my-project-0004-346516-pytorch112kag-asia-southeast1/Book7_The_Deathly_Hallows.pdf\n"
     ]
    }
   ],
   "source": [
    "# Verify that all Books 1 to 7 are uploaded to the GCS bucket (8 files in total, 2 for Part 1)\n",
    "!gsutil ls $GCS_BUCKET_URI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f6ccca82-4f07-4c3f-983e-3f5a82453cbb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datastore name: pytorch112kag-harry-potter-datastore\n",
      "Datastore ID: pytorch112kag-harry-potter-datastore-id\n"
     ]
    }
   ],
   "source": [
    "# The datastore name can only contain lowercase letters, numbers, and hyphens\n",
    "DATASTORE_NAME = f\"{UNIQUE_PREFIX}-harry-potter-datastore\"\n",
    "DATASTORE_ID = f\"{DATASTORE_NAME}-id\"\n",
    "LOCATION = 'global'\n",
    "\n",
    "# print variables for verification\n",
    "print(f\"Datastore name: {DATASTORE_NAME}\")\n",
    "print(f\"Datastore ID: {DATASTORE_ID}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "62465d05-8632-419f-939a-303bc78a9dc2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datastore pytorch112kag-harry-potter-datastore-id successfully created\n"
     ]
    }
   ],
   "source": [
    "# Create the datastore\n",
    "try:\n",
    "    create_data_store(PROJECT_ID, LOCATION, DATASTORE_NAME, DATASTORE_ID)\n",
    "    print(f\"Datastore {DATASTORE_ID} successfully created\")\n",
    "except:\n",
    "    print(\"if not running first time, DATASTORE_ID may already exist\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddcf6570-fb7b-451a-8216-8a8094a98614",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Start the import of documents into datastore\n",
    "import_documents(PROJECT_ID, LOCATION, DATASTORE_ID, GCS_BUCKET_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0d660a2-22b0-4ebd-b8ff-be6c8519ed74",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 2. Create a Search Engine for your datastore\n",
    "In this section, you will be creating a Search app which we will connect to the Harry Potter datastore that we have created earlier. \n",
    "\n",
    "For the search app, we will set the search_tier to Enterprise tier and to enable advanced LLM features. Enterprise tier is required to get extractive answers from a search query and advanced LLM features are required to sumarize search results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0cf35ab-abb7-491c-a544-1f485a4c7833",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Helper function to create a Vertex Search Engine\n",
    "def create_engine(\n",
    "    project_id: str, location: str, data_store_name: str, data_store_id: str\n",
    "):\n",
    "    # Create a client\n",
    "    client_options = (\n",
    "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
    "        if location != \"global\"\n",
    "        else None\n",
    "    )\n",
    "    client = discoveryengine.EngineServiceClient(client_options=client_options)\n",
    "\n",
    "    # Initialize request argument(s)\n",
    "    config = discoveryengine.Engine.SearchEngineConfig(\n",
    "        search_tier=\"SEARCH_TIER_ENTERPRISE\", search_add_ons=[\"SEARCH_ADD_ON_LLM\"]\n",
    "    )\n",
    "\n",
    "    engine = discoveryengine.Engine(\n",
    "        display_name=data_store_name,\n",
    "        solution_type=\"SOLUTION_TYPE_SEARCH\",\n",
    "        industry_vertical=\"GENERIC\",\n",
    "        data_store_ids=[data_store_id],\n",
    "        search_engine_config=config,\n",
    "    )\n",
    "\n",
    "    request = discoveryengine.CreateEngineRequest(\n",
    "        parent=discoveryengine.DataStoreServiceClient.collection_path(\n",
    "            project_id, location, \"default_collection\"\n",
    "        ),\n",
    "        engine=engine,\n",
    "        engine_id=engine.display_name,\n",
    "    )\n",
    "\n",
    "    # Make the request\n",
    "    operation = client.create_engine(request=request)\n",
    "    response = operation.result(timeout=90)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2487d037-9b14-4c2f-865f-4c42a3b16a6d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create the Vertex Search Engine\n",
    "try:\n",
    "    create_engine(PROJECT_ID, LOCATION, DATASTORE_NAME, DATASTORE_ID)\n",
    "except:\n",
    "    print(\"if not running first time, create_engine may already exist\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "274d9dad-5758-4447-925a-5fc59d3cb520",
   "metadata": {},
   "source": [
    "## 3. Query your datastore through the Search app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee26540f-72d3-418c-954e-859cdb54bb3f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "def search_sample(\n",
    "    project_id: str,\n",
    "    location: str,\n",
    "    data_store_id: str,\n",
    "    search_query: str,\n",
    ") -> List[discoveryengine.SearchResponse]:\n",
    "    #  For more information, refer to:\n",
    "    # https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store\n",
    "    client_options = (\n",
    "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
    "        if LOCATION != \"global\"\n",
    "        else None\n",
    "    )\n",
    "\n",
    "    # Create a client\n",
    "    client = discoveryengine.SearchServiceClient(client_options=client_options)\n",
    "\n",
    "    # The full resource name of the search engine serving config\n",
    "    # e.g. projects/{project_id}/locations/{location}/dataStores/{data_store_id}/servingConfigs/{serving_config_id}\n",
    "    serving_config = client.serving_config_path(\n",
    "        project=project_id,\n",
    "        location=location,\n",
    "        data_store=data_store_id,\n",
    "        serving_config=\"default_config\",\n",
    "    )\n",
    "\n",
    "    # Optional: Configuration options for search\n",
    "    # Refer to the `ContentSearchSpec` reference for all supported fields:\n",
    "    # https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.types.SearchRequest.ContentSearchSpec\n",
    "    content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec(\n",
    "        # For information about snippets, refer to:\n",
    "        # https://cloud.google.com/generative-ai-app-builder/docs/snippets\n",
    "        snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(\n",
    "            return_snippet=True\n",
    "        ),\n",
    "        # For information about search summaries, refer to:\n",
    "        # https://cloud.google.com/generative-ai-app-builder/docs/get-search-summaries\n",
    "        summary_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(\n",
    "            summary_result_count=5,\n",
    "            include_citations=True,\n",
    "            ignore_adversarial_query=True,\n",
    "            ignore_non_summary_seeking_query=False,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    # Refer to the `SearchRequest` reference for all supported fields:\n",
    "    # https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1.types.SearchRequest\n",
    "    request = discoveryengine.SearchRequest(\n",
    "        serving_config=serving_config,\n",
    "        query=search_query,\n",
    "        page_size=10,\n",
    "        content_search_spec=content_search_spec,\n",
    "        query_expansion_spec=discoveryengine.SearchRequest.QueryExpansionSpec(\n",
    "            condition=discoveryengine.SearchRequest.QueryExpansionSpec.Condition.AUTO,\n",
    "        ),\n",
    "        spell_correction_spec=discoveryengine.SearchRequest.SpellCorrectionSpec(\n",
    "            mode=discoveryengine.SearchRequest.SpellCorrectionSpec.Mode.AUTO\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    response = client.search(request)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df8ba78c-a51a-4349-b07a-66ff0f4ca81a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Ask a sample query to get an answer from the search engine!\n",
    "query = \"Who is the best friend of harry potter?\"\n",
    "\n",
    "print(search_sample(PROJECT_ID, LOCATION, DATASTORE_ID, query).summary.summary_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9820d9d3-5267-4107-be49-055ce53e84b1",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 3. Let's try asking a list of questions located in a CSV file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b46b0ec-0fc4-42e9-9fff-8a8a30f7654a",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Run through all the sample questions in the CSV file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b093465c-64b5-4202-b311-75db280e6e5f",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import time\n",
    "\n",
    "filename = \"./harry_potte_qa.csv\"\n",
    "df_qa_VertexSearch = pd.read_csv(filename, sep =\"|\")\n",
    "\n",
    "System_Prompts = \"\"\" You are an expert in reading harry potter books, but only provide evidences from the information provide and do not use an other information\n",
    "so here are some search results : \n",
    "\"\"\"\n",
    "\n",
    "Question_Prompts = \"\"\" -- Based on information above help to answer following user question\n",
    "\"\"\"\n",
    "\n",
    "# Iterate through the first 10 questions\n",
    "for i in range(0, 10):\n",
    "    \n",
    "    print(\"iteration #\", i)\n",
    "    time.sleep(2)  \n",
    "    \n",
    "    RAG_query = df_qa_VertexSearch.loc[i,'Question'] \n",
    "    \n",
    "    print(RAG_query)\n",
    "    try: \n",
    "        RAG_results = search_sample(PROJECT_ID, LOCATION, DATASTORE_ID, RAG_query).summary.summary_text\n",
    "    except:\n",
    "        RAG_results = \"No results\"\n",
    "        \n",
    "    Gemini_query = System_Prompts + \" \" + RAG_results + \" \" + Question_Prompts + \" \" + df_qa_VertexSearch.loc[i,'Question']\n",
    "    \n",
    "    try:        \n",
    "        df_qa_VertexSearch.loc[i, \"Gemini_pro_model_output_v1\"] = generate_pro(Gemini_query)\n",
    "        df_qa_VertexSearch.loc[i, \"palm_bison32k_output_v1\"] = generate_palm_bison32k(Gemini_query)\n",
    "\n",
    "    except:\n",
    "        df_qa_VertexSearch.loc[i, \"Gemini_pro_model_output_v1\"] = \"No answer found\"\n",
    "        df_qa_VertexSearch.loc[i, \"palm_bison32k_output_v1\"] = \"No answer found\"\n",
    "        print(\"long-running operation\")\n",
    "\n",
    "\n",
    "\n",
    "output1 = \"./results/harry_potter_qa_OOTB-RAG_output.csv\"\n",
    "\n",
    "df_qa_VertexSearch.to_csv(output1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6136097e-e01f-46fa-8cc7-d0a2b5f4a6a1",
   "metadata": {},
   "source": [
    "# 4. Deploy the Search app as a Streamlit app to Cloud Run\n",
    "For this section, please head to the [Google Cloud console](https://console.cloud.google.com/gen-app-builder/engines) to continue!"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-env-pytorch-pytorch",
   "name": "workbench-notebooks.m125",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m125"
  },
  "kernelspec": {
   "display_name": "PyTorch 1-13 (Local)",
   "language": "python",
   "name": "conda-env-pytorch-pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}