search/create_datastore_and_search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "829cc6c6-f8c3-4fc6-942b-f698be5fa1a2",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Create a Vertex AI Datastore and Search Engine\n",
    "\n",
    "<table align=\"left\">\n",
    "\n",
    "  <td>\n",
    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/search/create_datastore_and_search.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/search/reate_datastore_and_search.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
    "      View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/blob/main/search/reate_datastore_and_search.ipynb\">\n",
    "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
    "      Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0460bdbe-05db-42f8-822c-82195de8329a",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "\n",
    "* Authors: Kara Greenfield\n",
    "* Created: 22 Nov 2023\n",
    "\n",
    "---\n",
    "\n",
    "## Objective\n",
    "\n",
    "This notebook shows how to create and populate a Vertex AI Search Datastore, how to create a search app connected to that datastore, and how to submit queries through the search engine.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b4b1050-7113-487e-aeaf-55690c831a1d",
   "metadata": {},
   "source": [
    "Services used in the notebook:\n",
    "\n",
    "- ✅ Vertex AI Search for document search and retrieval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43828625-a130-449c-ba5f-6a948220f559",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Install pre-requisites\n",
    "\n",
    "If running in Colab install the pre-requisites into the runtime. Otherwise it is assumed that the notebook is running in Vertex AI Workbench. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "8ea5db8a-dccc-4442-b5d7-7088d5ffb5ac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pip install --upgrade google-cloud-discoveryengine -q --user"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de8c61e1-7c51-4fda-92c7-95c356321740",
   "metadata": {},
   "source": [
    "---\n",
    "#### ⚠️ Do not forget to restart the kernel!\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f6a5de5-156e-4f14-99e6-3ab33e076c81",
   "metadata": {},
   "source": [
    "## Authenticate\n",
    "\n",
    "If running in Colab authenticate with `google.colab.google.auth` otherwise assume that running on Vertex AI Workbench.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "df3b7309-c51c-44a0-a466-b0b2733e0c28",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    from google.colab import auth as google_auth\n",
    "\n",
    "    google_auth.authenticate_user()\n",
    "\n",
    "from google.auth import default\n",
    "\n",
    "creds, _ = default()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae7a4925-145a-40f3-9fa1-3b69a42d488d",
   "metadata": {},
   "source": [
    "## Configure notebook environment\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "389b9e51-1ce4-4bdf-80b5-fcdc1882f853",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from google.cloud import discoveryengine_v1alpha as discoveryengine\n",
    "\n",
    "PROJECT_ID = \"<YOUR_PROJECT_ID>\"  # @param {type:\"string\"}\n",
    "LOCATION = \"global\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07f1aecd-4633-4451-b5a7-1f26e4cb2631",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Create and Populate a Datastore\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9a1d0021-e299-49d4-a657-2e101ae49eb6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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": null,
   "id": "8a4db726-8da1-4c76-8934-944aaf5f9b53",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The datastore name can only contain lowercase letters, numbers, and hyphens\n",
    "DATASTORE_NAME = \"alphabet-contracts\"\n",
    "DATASTORE_ID = f\"{DATASTORE_NAME}-id\"\n",
    "\n",
    "create_data_store(PROJECT_ID, LOCATION, DATASTORE_NAME, DATASTORE_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "03121270-5d2f-403b-81ea-c1d241357bd1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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": "code",
   "execution_count": null,
   "id": "1ddfe66a-acb4-4fdb-b9ed-76a332bb0f0c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "source_documents_gs_uri = (\n",
    "    \"gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs\"\n",
    ")\n",
    "\n",
    "import_documents(PROJECT_ID, LOCATION, DATASTORE_ID, source_documents_gs_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a957202-b67e-47ca-84c3-b8a62cfbe405",
   "metadata": {},
   "source": [
    "## Create a Search Engine\n",
    "\n",
    "This is used to set the `search_tier` to enterprise and to enable advanced LLM features.\n",
    "\n",
    "Enterprise tier is required to get extractive answers from a search query and advanced LLM features are required to sumarize search results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "0e39e9bb-f381-44f6-a45e-3322669a171f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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": 26,
   "id": "4a853982-8c2e-402a-a808-5364bf932619",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "create_engine(PROJECT_ID, LOCATION, DATASTORE_NAME, DATASTORE_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9f4d978-9164-4de3-b01a-179051706313",
   "metadata": {},
   "source": [
    "## Query your Datastore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1c4dfb62-7846-43d0-9cba-fd8886ce5546",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\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=True,\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": 29,
   "id": "ad41e18b-38d2-4f4c-98ae-df14eda900ae",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sundar Pichai is the CEO of Google and Alphabet [1]. He has served as CEO of Google since 2015 and as CEO of Alphabet since 2019 [1].\n"
     ]
    }
   ],
   "source": [
    "query = \"Who is the CEO of Google?\"\n",
    "\n",
    "print(search_sample(PROJECT_ID, LOCATION, DATASTORE_ID, query).summary.summary_text)"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m113",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m113"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}