search/vertexai-search-options/vertexai_search_options.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": [
        "# Building Search Applications with Vertex AI Search\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/search/vertexai-search-options/vertexai_search_options.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%2Fsearch%2Fvertexai-search-options%2Fvertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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/search/vertexai-search-options/vertexai_search_options.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) | [Megha Agarwal](https://github.com/agarwal22megha/) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Vertex AI Search leverages decades of expertise Google has in information retrieval and brings together the power of deep information retrieval, state-of-the-art natural language processing, and the latest in large language model (LLM) processing to understand user intent and return the most relevant results for the user.\n",
        "\n",
        "Based on where Developers are in their journey, their orchestration framework of choice,  they can select Vertex AI Search out-of-the-box capabilities or customize their search solutions with Vertex Retrievers or use the Vertex AI DIY APIs to build the end to end RAG application.\n",
        "\n",
        "This notebook explores how you can leverage Vertex AI Search out-of-the-box capabilities/ customize your search application with Vertex AI Search Retriever via LangChain or Grounding Service.\n",
        "\n",
        "\n",
        "We create a Vertex AI Agent Builder Search App with unstructured data store on Google Cloud Console.\n",
        "\n",
        "Once the Search App is created, this notebook walks you through how you can leverage:\n",
        "\n",
        "1. Vertex AI Search out-of-the-box capabilities by leveraging the [Vertex AI Agent Builder Python API reference documentation](https://cloud.google.com/python/docs/reference/discoveryengine/latest)\n",
        "\n",
        "2. Leverage Vertex AI Search Datastore as a [Grounding Source](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/grounding) with Gemini Models directly to ground responses in your data\n",
        "\n",
        "3. Leverage Vertex DataStore as [LangChain Retriever](https://python.langchain.com/v0.2/docs/integrations/retrievers/google_vertex_ai_search/)\n"
      ]
    },
    {
      "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\n",
        "%pip install google-cloud-discoveryengine\n",
        "%pip install langchain_google_community\n",
        "%pip install langchain langchain-google-vertexai\n",
        "%pip install langchain-google-community[vertexaisearch]"
      ]
    },
    {
      "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 has 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 is 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": [
        "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": [
        "### Create a Vertex AI Search App On Google Cloud\n",
        "[Follow These Steps](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search#unstructured-data_1) to leverage Vertex AI Agent Builder to create a Search App with unstructured data store on Google Cloud Console.\n",
        "\n",
        "Once the search app is created successfully, make a Note of the Search App Id, Datastore ID, location from Vertex AI Agent Builder console."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CaXr_Fz9OXkQ"
      },
      "outputs": [],
      "source": [
        "# Set to your data store location\n",
        "VERTEX_AI_SEARCH_LOCATION = \"global\"  # @param {type:\"string\"}\n",
        "# Set to your search app ID\n",
        "VERTEX_AI_SEARCH_APP_ID = \"your_search_app_id\"  # @param {type:\"string\"}\n",
        "# Set to your data store ID\n",
        "VERTEX_AI_SEARCH_DATASTORE_ID = \"your_datastore_id\"  # @param {type:\"string\"}\n",
        "\n",
        "MODEL = \"gemini-1.5-pro-001\"  # @param {type:\"string\"}\n",
        "\n",
        "SEARCH_QUERY = \"When does Alphabet plan to get to net zero?\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "from google.api_core.client_options import ClientOptions\n",
        "from google.cloud import discoveryengine_v1 as discoveryengine\n",
        "from langchain.chains import RetrievalQAWithSourcesChain\n",
        "from langchain_google_community import VertexAISearchRetriever\n",
        "from langchain_google_vertexai import VertexAI\n",
        "from vertexai.generative_models import (\n",
        "    GenerationConfig,\n",
        "    GenerativeModel,\n",
        "    SafetySetting,\n",
        "    Tool,\n",
        ")\n",
        "import vertexai.preview.generative_models as generative_models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QhOI5q1LIzmU"
      },
      "source": [
        "## Part 1: Vertex AI Search Out-of-the-box capabilities\n",
        "Leverage the [Vertex AI Agent Builder Python API reference documentation](https://cloud.google.com/python/docs/reference/discoveryengine/latest) to get the search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Chkv32AgF9Sx"
      },
      "outputs": [],
      "source": [
        "CUSTOM_PROMPT = \"\"\"\n",
        "    <PERSONA_AND_GOAL>\n",
        "      You are a helpful assistant knowledgeable about Alphabet quarterly earning reports.\n",
        "      Be mindful of the time frame user inputs\n",
        "      Help user with their queries related with Alphabet\n",
        "      Only respond with relevant information available in Grounding Knowledge snippet.\n",
        "      If no relevant snippet is available, respond with you dont know\n",
        "      Do not make up information\n",
        "    </PERSONA_AND_GOAL>\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0eOlSO08IK2Z"
      },
      "outputs": [],
      "source": [
        "def search_spec():\n",
        "    content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec(\n",
        "        snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(\n",
        "            return_snippet=True\n",
        "        ),\n",
        "        summary_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(\n",
        "            summary_result_count=10,\n",
        "            include_citations=True,\n",
        "            ignore_adversarial_query=True,\n",
        "            ignore_non_summary_seeking_query=True,\n",
        "            model_prompt_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec.ModelPromptSpec(\n",
        "                preamble=CUSTOM_PROMPT\n",
        "            ),\n",
        "            model_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec.ModelSpec(\n",
        "                version=\"stable\",\n",
        "            ),\n",
        "        ),\n",
        "    )\n",
        "    return content_search_spec\n",
        "\n",
        "\n",
        "def search_sample(project_id: str, location: str, engine_id: str, search_query: str):\n",
        "    client_options = (\n",
        "        ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
        "        if location != \"global\"\n",
        "        else None\n",
        "    )\n",
        "\n",
        "    client = discoveryengine.SearchServiceClient(client_options=client_options)\n",
        "\n",
        "    serving_config = f\"projects/{project_id}/locations/{location}/collections/default_collection/engines/{engine_id}/servingConfigs/default_config\"\n",
        "\n",
        "    content_search_spec = search_spec()\n",
        "\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",
        "    response = client.search(request)\n",
        "\n",
        "    return response"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IV52oaTJIK-i"
      },
      "outputs": [],
      "source": [
        "search_response = search_sample(\n",
        "    PROJECT_ID, VERTEX_AI_SEARCH_LOCATION, VERTEX_AI_SEARCH_APP_ID, SEARCH_QUERY\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JsbqEIgyILBJ"
      },
      "outputs": [],
      "source": [
        "search_response.summary.summary_with_metadata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_OnZDEfJCzi"
      },
      "source": [
        "## Part 2: Grounding with Gemini on Vertex AI Search Datastore\n",
        "\n",
        "Leverage Vertex AI Search Datastore as a [Grounding Source](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/grounding) with Gemini Models directly to ground responses in your data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uGV_iw5YJMNK"
      },
      "outputs": [],
      "source": [
        "system_instructions = \"\"\"\n",
        "<PERSONA_AND_GOAL>\n",
        "  You are a helpful assistant knowledgeable about Alphabet quarterly earning reports.\n",
        "  Help user with their queries related with Alphabet by following given <INSTRUCTIONS> and <CONTEXT>\n",
        "  only respond with information available in Grounding Knowledge store.\n",
        "</PERSONA_AND_GOAL>\n",
        "\n",
        "<INSTRUCTIONS>\n",
        "- Always refer to the tool and Ground your answers in it\n",
        "- Understand the retrieved snippet by the tool and only use that information to help users\n",
        "- For supporting references, you can provide the Grounding tool snippets verbatim, and any other info like page number\n",
        "- For Information not available in the tool, mention you dont have access to the information.\n",
        "- Output \"answer\" should be I dont know when the user question is irrelevant or outside the <CONTEXT>\n",
        "- Leave \"reference_snippet\" as null if you are not sure about the page and text snippet\n",
        "<INSTRUCTIONS>\n",
        "\n",
        "<CONTEXT>\n",
        "  Grounding tool finds most relevant snippets from the Alphabet earning reports data store.\n",
        "  Use the information provided by the tool as your knowledge base.\n",
        "</CONTEXT>\n",
        "\n",
        "<CONSTRAINTS>\n",
        "- ONLY use information available from the Grounding tool\n",
        "</CONSTRAINTS>\n",
        "\n",
        "<OUTPUT_FORMAT>\n",
        "- Response should ALWAYS be in following JSON Output with answer and reference_snippet as keys, e.g. {\"answer\": , \"reference_snippet\": }\n",
        "</OUTPUT_FORMAT>\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wSo25gr9MBSd"
      },
      "outputs": [],
      "source": [
        "tools = [\n",
        "    Tool.from_retrieval(\n",
        "        retrieval=generative_models.grounding.Retrieval(\n",
        "            source=generative_models.grounding.VertexAISearch(\n",
        "                datastore=VERTEX_AI_SEARCH_DATASTORE_ID,\n",
        "                project=PROJECT_ID,\n",
        "                location=VERTEX_AI_SEARCH_LOCATION,\n",
        "            ),\n",
        "            disable_attribution=False,\n",
        "        )\n",
        "    ),\n",
        "]\n",
        "\n",
        "generation_config = GenerationConfig(max_output_tokens=8192, temperature=1, top_p=0.95)\n",
        "\n",
        "safety_settings = [\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    ),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jCJAOvHdLrGI"
      },
      "outputs": [],
      "source": [
        "model = GenerativeModel(MODEL, tools=tools, system_instruction=[system_instructions])\n",
        "\n",
        "response = model.generate_content(\n",
        "    [SEARCH_QUERY],\n",
        "    generation_config=generation_config,\n",
        "    safety_settings=safety_settings,\n",
        "    stream=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BNrZ-w3gRoLx"
      },
      "outputs": [],
      "source": [
        "for candidate in response:\n",
        "    print(candidate)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lQv61qA5Q71g"
      },
      "source": [
        "##### [Here](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/grounding/intro-grounding-gemini.ipynb) is another Grounding Gemini with Vertex AI Search Example which you might find useful"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V2X81IQUOImp"
      },
      "source": [
        "## Part 3: Vertex AI Search with LangChain\n",
        "\n",
        "Developers have the flexibility to incorporate Vertex AI Search as a [LangChain Retriever](https://python.langchain.com/v0.2/docs/integrations/retrievers/google_vertex_ai_search/) in their existing LangChain applications.\n",
        "\n",
        "This means you can continue leveraging your preferred orchestrator while seamlessly integrating Vertex AI Search data stores into existing RAG pipelines.  Vertex AI Search enables Google-quality search capabilities applied directly to your custom data, elevating the search result quality and relevance of your retrieval-augmented generation workflows.\n",
        "\n",
        "Find more notebook examples of leveraging VertexAISearchRetriever with LangChain [here](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/search/retrieval-augmented-generation/examples/question_answering.ipynb).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kt09YTDmJMYW"
      },
      "outputs": [],
      "source": [
        "llm = VertexAI(model_name=MODEL)\n",
        "\n",
        "retriever = VertexAISearchRetriever(\n",
        "    project_id=PROJECT_ID,\n",
        "    location_id=VERTEX_AI_SEARCH_LOCATION,\n",
        "    data_store_id=VERTEX_AI_SEARCH_DATASTORE_ID,\n",
        "    get_extractive_answers=True,\n",
        "    max_documents=10,\n",
        "    max_extractive_segment_count=1,\n",
        "    max_extractive_answer_count=5,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dkt_l9tmPeYH"
      },
      "outputs": [],
      "source": [
        "retrieval_qa_with_sources = RetrievalQAWithSourcesChain.from_chain_type(\n",
        "    llm=llm, chain_type=\"stuff\", retriever=retriever\n",
        ")\n",
        "\n",
        "retrieval_qa_with_sources.invoke(SEARCH_QUERY, return_only_outputs=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "[Delete the Vertex AI Search App](https://cloud.google.com/generative-ai-app-builder/docs/delete-engine) and\n",
        "\n",
        "[Delete the datastore](https://cloud.google.com/generative-ai-app-builder/docs/delete-a-data-store) your created"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "vertexai_search_options.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}