{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Google Vertex AI Search\n",
    "\n",
    "[Vertex AI Search](https://cloud.google.com/enterprise-search) (formerly known as Enterprise Search on Generative AI App Builder) is a part of the [Vertex AI](https://cloud.google.com/vertex-ai) machine learning platform offered by Google Cloud.\n",
    "\n",
    "Vertex AI Search lets organizations quickly build generative AI powered search engines for customers and employees. It's underpinned by a variety of Google Search technologies, including semantic search, which helps deliver more relevant results than traditional keyword-based search techniques by using natural language processing and machine learning techniques to infer relationships within the content and intent from the user’s query input. Vertex AI Search also benefits from Google’s expertise in understanding how users search and factors in content relevance to order displayed results.\n",
    "\n",
    "Vertex AI Search is available in the Google Cloud Console and via an API for enterprise workflow integration.\n",
    "\n",
    "This notebook demonstrates how to configure Vertex AI Search and use the Vertex AI Search retriever. The Vertex AI Search retriever encapsulates the [Python client library](https://cloud.google.com/generative-ai-app-builder/docs/libraries#client-libraries-install-python) and uses it to access the [Search Service API](https://cloud.google.com/python/docs/reference/discoveryengine/latest/google.cloud.discoveryengine_v1beta.services.search_service).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install pre-requisites\n",
    "\n",
    "You need to install the `google-cloud-discoveryengine` package to use the Vertex AI Search retriever.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install google-cloud-discoveryengine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure access to Google Cloud and Vertex AI Search\n",
    "\n",
    "Vertex AI Search is generally available without allowlist as of August 2023.\n",
    "\n",
    "Before you can use the retriever, you need to complete the following steps:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a search engine and populate an unstructured data store\n",
    "\n",
    "- Follow the instructions in the [Vertex AI Search Getting Started guide](https://cloud.google.com/generative-ai-app-builder/docs/try-enterprise-search) to set up a Google Cloud project and Vertex AI Search.\n",
    "- [Use the Google Cloud Console to create an unstructured data store](https://cloud.google.com/generative-ai-app-builder/docs/create-engine-es#unstructured-data)\n",
    "  - Populate it with the example PDF documents from the `gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs` Cloud Storage folder.\n",
    "  - Make sure to use the `Cloud Storage (without metadata)` option.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set credentials to access Vertex AI Search API\n",
    "\n",
    "The [Vertex AI Search client libraries](https://cloud.google.com/generative-ai-app-builder/docs/libraries) used by the Vertex AI Search retriever provide high-level language support for authenticating to Google Cloud programmatically.\n",
    "Client libraries support [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials); the libraries look for credentials in a set of defined locations and use those credentials to authenticate requests to the API.\n",
    "With ADC, you can make credentials available to your application in a variety of environments, such as local development or production, without needing to modify your application code.\n",
    "\n",
    "If running in [Google Colab](https://colab.google) authenticate with `google.colab.google.auth` otherwise follow one of the [supported methods](https://cloud.google.com/docs/authentication/application-default-credentials) to make sure that you Application Default Credentials are properly set.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure and use the Vertex AI Search retriever\n",
    "\n",
    "The Vertex AI Search retriever is implemented in the `langchain.retriever.GoogleVertexAISearchRetriever` class. The `get_relevant_documents` method returns a list of `langchain.schema.Document` documents where the `page_content` field of each document is populated the document content.\n",
    "Depending on the data type used in Vertex AI Search (website, structured or unstructured) the `page_content` field is populated as follows:\n",
    "\n",
    "- Website with advanced indexing: an `extractive answer` that matches a query. The `metadata` field is populated with metadata (if any) of the document from which the segments or answers were extracted.\n",
    "- Unstructured data source: either an `extractive segment` or an `extractive answer` that matches a query. The `metadata` field is populated with metadata (if any) of the document from which the segments or answers were extracted.\n",
    "- Structured data source: a string json containing all the fields returned from the structured data source. The `metadata` field is populated with metadata (if any) of the document\n",
    "\n",
    "### Extractive answers & extractive segments\n",
    "\n",
    "An extractive answer is verbatim text that is returned with each search result. It is extracted directly from the original document. Extractive answers are typically displayed near the top of web pages to provide an end user with a brief answer that is contextually relevant to their query. Extractive answers are available for website and unstructured search.\n",
    "\n",
    "An extractive segment is verbatim text that is returned with each search result. An extractive segment is usually more verbose than an extractive answer. Extractive segments can be displayed as an answer to a query, and can be used to perform post-processing tasks and as input for large language models to generate answers or new text. Extractive segments are available for unstructured search.\n",
    "\n",
    "For more information about extractive segments and extractive answers refer to [product documentation](https://cloud.google.com/generative-ai-app-builder/docs/snippets).\n",
    "\n",
    "NOTE: Extractive segments require the [Enterprise edition](https://cloud.google.com/generative-ai-app-builder/docs/about-advanced-features#enterprise-features) features to be enabled.\n",
    "\n",
    "When creating an instance of the retriever you can specify a number of parameters that control which data store to access and how a natural language query is processed, including configurations for extractive answers and segments.\n",
    "\n",
    "### The mandatory parameters are:\n",
    "\n",
    "- `project_id` - Your Google Cloud Project ID.\n",
    "- `location_id` - The location of the data store.\n",
    "  - `global` (default)\n",
    "  - `us`\n",
    "  - `eu`\n",
    "- `data_store_id` - The ID of the data store you want to use.\n",
    "  - Note: This was called `search_engine_id` in previous versions of the retriever.\n",
    "\n",
    "The `project_id` and `data_store_id` parameters can be provided explicitly in the retriever's constructor or through the environment variables - `PROJECT_ID` and `DATA_STORE_ID`.\n",
    "\n",
    "You can also configure a number of optional parameters, including:\n",
    "\n",
    "- `max_documents` - The maximum number of documents used to provide extractive segments or extractive answers\n",
    "- `get_extractive_answers` - By default, the retriever is configured to return extractive segments.\n",
    "  - Set this field to `True` to return extractive answers. This is used only when `engine_data_type` set to `0` (unstructured)\n",
    "- `max_extractive_answer_count` - The maximum number of extractive answers returned in each search result.\n",
    "  - At most 5 answers will be returned. This is used only when `engine_data_type` set to `0` (unstructured).\n",
    "- `max_extractive_segment_count` - The maximum number of extractive segments returned in each search result.\n",
    "  - Currently one segment will be returned. This is used only when `engine_data_type` set to `0` (unstructured).\n",
    "- `filter` - The filter expression for the search results based on the metadata associated with the documents in the data store.\n",
    "- `query_expansion_condition` - Specification to determine under which conditions query expansion should occur.\n",
    "  - `0` - Unspecified query expansion condition. In this case, server behavior defaults to disabled.\n",
    "  - `1` - Disabled query expansion. Only the exact search query is used, even if SearchResponse.total_size is zero.\n",
    "  - `2` - Automatic query expansion built by the Search API.\n",
    "- `engine_data_type` - Defines the Vertex AI Search data type\n",
    "  - `0` - Unstructured data\n",
    "  - `1` - Structured data\n",
    "  - `2` - Website data with [Advanced Website Indexing](https://cloud.google.com/generative-ai-app-builder/docs/about-advanced-features#advanced-website-indexing)\n",
    "\n",
    "### Migration guide for `GoogleCloudEnterpriseSearchRetriever`\n",
    "\n",
    "In previous versions, this retriever was called `GoogleCloudEnterpriseSearchRetriever`. Some backwards-incompatible changes had to be made to the retriever after the General Availability launch due to changes in the product behavior.\n",
    "\n",
    "To update to the new retriever, make the following changes:\n",
    "\n",
    "- Change the import from: `from langchain.retrievers import GoogleCloudEnterpriseSearchRetriever` -> `from langchain.retrievers import GoogleVertexAISearchRetriever`.\n",
    "- Change all class references from `GoogleCloudEnterpriseSearchRetriever` -> `GoogleVertexAISearchRetriever`.\n",
    "- Upon class initialization, change the `search_engine_id` parameter name to `data_store_id`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and use the retriever for **unstructured** data with extractive segments\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import (\n",
    "    GoogleVertexAIMultiTurnSearchRetriever,\n",
    "    GoogleVertexAISearchRetriever,\n",
    ")\n",
    "\n",
    "PROJECT_ID = \"<YOUR PROJECT ID>\"  # Set to your Project ID\n",
    "LOCATION_ID = \"<YOUR LOCATION>\"  # Set to your data store location\n",
    "DATA_STORE_ID = \"<YOUR DATA STORE ID>\"  # Set to your data store ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = GoogleVertexAISearchRetriever(\n",
    "    project_id=PROJECT_ID,\n",
    "    location_id=LOCATION_ID,\n",
    "    data_store_id=DATA_STORE_ID,\n",
    "    max_documents=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What are Alphabet's Other Bets?\"\n",
    "\n",
    "result = retriever.get_relevant_documents(query)\n",
    "for doc in result:\n",
    "    print(doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and use the retriever for **unstructured** data with extractive answers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = GoogleVertexAISearchRetriever(\n",
    "    project_id=PROJECT_ID,\n",
    "    location_id=LOCATION_ID,\n",
    "    data_store_id=DATA_STORE_ID,\n",
    "    max_documents=3,\n",
    "    max_extractive_answer_count=3,\n",
    "    get_extractive_answers=True,\n",
    ")\n",
    "\n",
    "result = retriever.get_relevant_documents(query)\n",
    "for doc in result:\n",
    "    print(doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and use the retriever for **structured** data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = GoogleVertexAISearchRetriever(\n",
    "    project_id=PROJECT_ID,\n",
    "    location_id=LOCATION_ID,\n",
    "    data_store_id=DATA_STORE_ID,\n",
    "    max_documents=3,\n",
    "    engine_data_type=1,\n",
    ")\n",
    "\n",
    "result = retriever.get_relevant_documents(query)\n",
    "for doc in result:\n",
    "    print(doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and use the retriever for **website** data with Advanced Website Indexing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = GoogleVertexAISearchRetriever(\n",
    "    project_id=PROJECT_ID,\n",
    "    location_id=LOCATION_ID,\n",
    "    data_store_id=DATA_STORE_ID,\n",
    "    max_documents=3,\n",
    "    max_extractive_answer_count=3,\n",
    "    get_extractive_answers=True,\n",
    "    engine_data_type=2,\n",
    ")\n",
    "\n",
    "result = retriever.get_relevant_documents(query)\n",
    "for doc in result:\n",
    "    print(doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure and use the retriever for multi-turn search\n",
    "\n",
    "[Search with follow-ups](https://cloud.google.com/generative-ai-app-builder/docs/multi-turn-search) is based on generative AI models and it is different from the regular unstructured data search.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = GoogleVertexAIMultiTurnSearchRetriever(\n",
    "    project_id=PROJECT_ID, location_id=LOCATION_ID, data_store_id=DATA_STORE_ID\n",
    ")\n",
    "\n",
    "result = retriever.get_relevant_documents(query)\n",
    "for doc in result:\n",
    "    print(doc)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
