search/custom-embeddings/custom_embeddings.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z3Gi-zdCeEbE"
      },
      "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": "WcfbGhs1eJF6"
      },
      "source": [
        "# Custom Embeddings 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/custom-embeddings/custom_embeddings.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%2Fcustom-embeddings%2Fcustom_embeddings.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/custom-embeddings/custom_embeddings.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/search/custom-embeddings/custom_embeddings.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/search/custom-embeddings/custom_embeddings.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/custom-embeddings/custom_embeddings.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/custom-embeddings/custom_embeddings.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/custom-embeddings/custom_embeddings.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/custom-embeddings/custom_embeddings.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/custom-embeddings/custom_embeddings.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": "ghQ2aBsbnMyn"
      },
      "source": [
        "---\n",
        "\n",
        "* Author: Holt Skinner\n",
        "\n",
        "---\n",
        "\n",
        "This notebook demonstrates how to:\n",
        "\n",
        "  - Get text embeddings using [`textembedding-gecko` in Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings)\n",
        "  - Convert embeddings into the [format expected by Vertex AI Search](https://cloud.google.com/generative-ai-app-builder/docs/prepare-data#unstructured)\n",
        "  - [Create a search app with custom embeddings](https://cloud.google.com/generative-ai-app-builder/docs/bring-embeddings)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsW5tPDRkT4m"
      },
      "source": [
        "## Getting started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jx1FQVAokWVb"
      },
      "source": [
        "### Install libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJFw23w1kYVj"
      },
      "outputs": [],
      "source": [
        "%pip install -q --upgrade --user google-cloud-aiplatform google-cloud-discoveryengine google-cloud-storage google-cloud-bigquery[pandas] google-cloud-bigquery-storage pandas ipywidgets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a778a925e288"
      },
      "outputs": [],
      "source": [
        "%load_ext google.cloud.bigquery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "42f78e5e31a4"
      },
      "source": [
        "---\n",
        "#### ⚠️ Do not forget to click the \"RESTART RUNTIME\" button above.\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-jwsaMQYkZm8"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, you will need to authenticate your environment. To do this, run the new cell below. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ikOmH4doxOFs"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    # Authenticate user to Google Cloud\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-h0ba4rmkpKW"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YLUml_s7iqBc"
      },
      "outputs": [],
      "source": [
        "import subprocess\n",
        "import time\n",
        "\n",
        "from google.api_core.client_options import ClientOptions\n",
        "from google.api_core.exceptions import GoogleAPICallError\n",
        "from google.cloud import bigquery\n",
        "from google.cloud import discoveryengine_v1alpha as discoveryengine\n",
        "from google.cloud import storage\n",
        "import requests\n",
        "from tqdm import tqdm  # to show a progress bar\n",
        "import vertexai\n",
        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
        "\n",
        "tqdm.pandas()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7933d0d989be"
      },
      "source": [
        "## Configure notebook environment\n",
        "\n",
        "### Set the following constants to reflect your environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "277ff46137d2"
      },
      "outputs": [],
      "source": [
        "# Define project information for Vertex AI\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# Initialize Vertex AI SDK\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cd2d5363e770"
      },
      "source": [
        "## Creating embeddings with Vertex AI\n",
        "\n",
        "### Data Preparation\n",
        "\n",
        "We will be using [the Stack Overflow public dataset](https://console.cloud.google.com/marketplace/product/stack-exchange/stack-overflow) hosted on BigQuery table `bigquery-public-data.stackoverflow.posts_questions`.\n",
        "\n",
        "This is a very big dataset with 23 million rows that doesn't fit into memory. We are going to limit it to 500 rows for this tutorial.\n",
        "\n",
        "- Fetch the data from BigQuery\n",
        "- Concat the Title and Body, and create embeddings from the text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "caeb5dc9d824"
      },
      "outputs": [],
      "source": [
        "bq_client = bigquery.Client(project=PROJECT_ID)\n",
        "query = f\"\"\"\n",
        "SELECT\n",
        "  DISTINCT \n",
        "  q.id,\n",
        "  q.title,\n",
        "  q.body,\n",
        "  q.answer_count,\n",
        "  q.comment_count,\n",
        "  q.creation_date,\n",
        "  q.last_activity_date,\n",
        "  q.score,\n",
        "  q.tags,\n",
        "  q.view_count\n",
        "FROM\n",
        "  `bigquery-public-data.stackoverflow.posts_questions` AS q\n",
        "WHERE\n",
        "  q.score > 0\n",
        "ORDER BY\n",
        "  q.view_count DESC\n",
        "LIMIT\n",
        "  500;\n",
        "\"\"\"\n",
        "\n",
        "# Load the BQ Table into a Pandas DataFrame\n",
        "df = bq_client.query(query).result().to_dataframe()\n",
        "\n",
        "# Convert ID to String\n",
        "df[\"id\"] = df[\"id\"].apply(str)\n",
        "\n",
        "# examine the data\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1e8c38f4cb9a"
      },
      "source": [
        "### Call the API to generate embeddings\n",
        "\n",
        "With the Stack Overflow dataset, we will use the `title` and `body` columns (the question title and description) and generate embedding for it with Embeddings for Text API. The API is available under the [`vertexai`](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai) package of the SDK.\n",
        "\n",
        "From the package, import [`TextEmbeddingModel`](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.language_models.TextEmbeddingModel) and get a model.\n",
        "\n",
        "For more information, refer to:\n",
        "\n",
        "- [Vertex AI: Get Text Embeddings](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings)\n",
        "- [Vertex AI: Model versions and lifecycle](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/model-versioning)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "20498d0ff2bb"
      },
      "outputs": [],
      "source": [
        "# Load the text embeddings model\n",
        "model = TextEmbeddingModel.from_pretrained(\"text-embedding-004\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b5f970232b4e"
      },
      "outputs": [],
      "source": [
        "# Get embeddings for a list of texts\n",
        "\n",
        "\n",
        "def get_embeddings_wrapper(texts, batch_size: int = 5) -> list:\n",
        "    embs = []\n",
        "    for i in tqdm(range(0, len(texts), batch_size)):\n",
        "        # Create embeddings optimized for document retrieval\n",
        "        result = model.get_embeddings(\n",
        "            [\n",
        "                TextEmbeddingInput(text=text, task_type=\"RETRIEVAL_DOCUMENT\")\n",
        "                for text in texts[i : i + batch_size]\n",
        "            ]\n",
        "        )\n",
        "        embs.extend([e.values for e in result])\n",
        "    return embs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ad26e79ea33"
      },
      "source": [
        "Get embeddings for the question titles/body and add them as the `\"embedding\"` column."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "84025146f55e"
      },
      "outputs": [],
      "source": [
        "df[\"title_body\"] = df[\"title\"] + \"\\n\" + df[\"body\"]\n",
        "\n",
        "df = df.assign(embedding=get_embeddings_wrapper(df.title_body))\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TKBmi2BMk_OU"
      },
      "source": [
        "## Scrape HTML from Question Pages\n",
        "\n",
        "- Get the HTML from the StackOverflow Question page\n",
        "   - Upload it to GCS as the Document Store/for displayed search results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tXHmC10IitET"
      },
      "outputs": [],
      "source": [
        "JSONL_MIME_TYPE = \"application/jsonl\"\n",
        "HTML_MIME_TYPE = \"text/html\"\n",
        "\n",
        "BUCKET_NAME = \"ucs-demo\"\n",
        "DIRECTORY = \"embeddings-stackoverflow\"\n",
        "BLOB_PREFIX = f\"{DIRECTORY}/html/\"\n",
        "\n",
        "GCS_URI_PREFIX = f\"gs://{BUCKET_NAME}/{BLOB_PREFIX}\"\n",
        "\n",
        "storage_client = storage.Client()\n",
        "bucket = storage_client.bucket(BUCKET_NAME)\n",
        "\n",
        "\n",
        "def scrape_question(question_url: str) -> str:\n",
        "    response = requests.get(question_url)\n",
        "\n",
        "    if response.status_code != 200 or not response.content:\n",
        "        print(f\"URL: {question_url} Code: {response.status_code}\")\n",
        "        return None\n",
        "\n",
        "    print(f\"Scraping {question_url}\")\n",
        "\n",
        "    link_title = response.url.split(\"/\")[-1] + \".html\"\n",
        "    gcs_uri = f\"{GCS_URI_PREFIX}{link_title}\"\n",
        "\n",
        "    # Upload HTML to Google Cloud Storage\n",
        "    blob = bucket.blob(f\"{BLOB_PREFIX}{link_title}\")\n",
        "    blob.upload_from_string(response.content, content_type=HTML_MIME_TYPE)\n",
        "    time.sleep(1)\n",
        "    return gcs_uri"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b2109f01aef"
      },
      "outputs": [],
      "source": [
        "# Get the published URL from the ID\n",
        "QUESTION_BASE_URL = \"https://stackoverflow.com/questions/\"\n",
        "df[\"question_url\"] = df[\"id\"].apply(lambda x: f\"{QUESTION_BASE_URL}{x}\")\n",
        "\n",
        "# Scrape HTML from stackoverflow.com and upload to GCS\n",
        "df[\"gcs_uri\"] = df[\"question_url\"].apply(scrape_question)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "31fd6d9cf358"
      },
      "source": [
        "Restructure the embeddings data to JSONL to follow the [Vertex AI Search format (Unstructured with Metadata)](https://cloud.google.com/generative-ai-app-builder/docs/prepare-data). This format is required to use custom embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "61e4989dfde3"
      },
      "outputs": [],
      "source": [
        "EMBEDDINGS_FIELD_NAME = \"embedding_vector\"\n",
        "\n",
        "\n",
        "def format_row(row):\n",
        "    return {\n",
        "        \"id\": row[\"id\"],\n",
        "        \"content\": {\"mimeType\": HTML_MIME_TYPE, \"uri\": row[\"gcs_uri\"]},\n",
        "        \"structData\": {\n",
        "            EMBEDDINGS_FIELD_NAME: row[\"embedding\"],\n",
        "            \"title\": row[\"title\"],\n",
        "            \"body\": row[\"body\"],\n",
        "            \"question_url\": row[\"question_url\"],\n",
        "            \"answer_count\": row[\"answer_count\"],\n",
        "            \"creation_date\": row[\"creation_date\"],\n",
        "            \"score\": row[\"score\"],\n",
        "        },\n",
        "    }\n",
        "\n",
        "\n",
        "vais_embeddings = (\n",
        "    df.apply(format_row, axis=1)\n",
        "    .to_json(orient=\"records\", lines=True, force_ascii=False)\n",
        "    .replace(r\"\\/\", \"/\")  # To prevent escaping the / characters\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ca66ea170841"
      },
      "source": [
        "Upload the JSONL file to Google Cloud Storage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "09d9ef95f44f"
      },
      "outputs": [],
      "source": [
        "jsonl_filename = f\"{DIRECTORY}/vais_embeddings.jsonl\"\n",
        "embeddings_file = f\"gs://{BUCKET_NAME}/{jsonl_filename}\"\n",
        "\n",
        "blob = bucket.blob(jsonl_filename)\n",
        "blob.upload_from_string(vais_embeddings, content_type=JSONL_MIME_TYPE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a7db96893f7"
      },
      "source": [
        "## Set up Vertex AI Search & Conversation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1046f3aaa57a"
      },
      "outputs": [],
      "source": [
        "DATA_STORE_LOCATION = \"global\"\n",
        "\n",
        "client_options = (\n",
        "    ClientOptions(api_endpoint=f\"{DATA_STORE_LOCATION}-discoveryengine.googleapis.com\")\n",
        "    if DATA_STORE_LOCATION != \"global\"\n",
        "    else None\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a11ceb9aded5"
      },
      "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 = 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",
        "        solution_types=[\"SOLUTION_TYPE_SEARCH\"],\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",
        "    try:\n",
        "        operation.result()\n",
        "    except GoogleAPICallError:\n",
        "        pass\n",
        "\n",
        "\n",
        "def update_schema(\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    data_store_id: str,\n",
        "):\n",
        "    client = discoveryengine.SchemaServiceClient(client_options=client_options)\n",
        "\n",
        "    schema = discoveryengine.Schema(\n",
        "        name=client.schema_path(project_id, location, data_store_id, \"default_schema\"),\n",
        "        struct_schema={\n",
        "            \"$schema\": \"https://json-schema.org/draft/2020-12/schema\",\n",
        "            \"type\": \"object\",\n",
        "            \"properties\": {\n",
        "                EMBEDDINGS_FIELD_NAME: {\n",
        "                    \"type\": \"array\",\n",
        "                    \"keyPropertyMapping\": \"embedding_vector\",\n",
        "                    \"dimension\": 768,\n",
        "                    \"items\": {\"type\": \"number\"},\n",
        "                }\n",
        "            },\n",
        "        },\n",
        "    )\n",
        "\n",
        "    operation = client.update_schema(\n",
        "        request=discoveryengine.UpdateSchemaRequest(schema=schema)\n",
        "    )\n",
        "\n",
        "    print(\"Waiting for operation to complete...\")\n",
        "\n",
        "    response = operation.result()\n",
        "\n",
        "    # Handle the response\n",
        "    print(response)\n",
        "\n",
        "\n",
        "def import_documents(\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    data_store_id: str,\n",
        "    gcs_uri: str,\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",
        "    request = discoveryengine.ImportDocumentsRequest(\n",
        "        parent=parent,\n",
        "        gcs_source=discoveryengine.GcsSource(input_uris=[gcs_uri]),\n",
        "        # Options: `FULL`, `INCREMENTAL`\n",
        "        reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.FULL,\n",
        "    )\n",
        "\n",
        "    # Make the request\n",
        "    operation = client.import_documents(request=request)\n",
        "\n",
        "\n",
        "def create_engine(\n",
        "    project_id: str, location: str, data_store_name: str, data_store_id: str\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,
      "metadata": {
        "id": "07821e2bf681"
      },
      "outputs": [],
      "source": [
        "DATA_STORE_NAME = \"stackoverflow-embeddings\"\n",
        "DATA_STORE_ID = f\"{DATA_STORE_NAME}-id\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "abacceea9f35"
      },
      "outputs": [],
      "source": [
        "# Create a Data Store\n",
        "create_data_store(PROJECT_ID, DATA_STORE_LOCATION, DATA_STORE_NAME, DATA_STORE_ID)\n",
        "\n",
        "# Update the Data Store Schema for embeddings\n",
        "update_schema(PROJECT_ID, DATA_STORE_LOCATION, DATA_STORE_ID)\n",
        "\n",
        "# Import the embeddings JSONL file\n",
        "import_documents(PROJECT_ID, DATA_STORE_LOCATION, DATA_STORE_ID, embeddings_file)\n",
        "\n",
        "# Create a Search App and attach the Data Store\n",
        "create_engine(PROJECT_ID, DATA_STORE_LOCATION, DATA_STORE_NAME, DATA_STORE_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6066cae8e5b"
      },
      "source": [
        "Next, we need to set the embedding specification for the data store. We will set the same spec for all search requests.\n",
        "\n",
        "`0.5 * relevance_score`\n",
        "\n",
        "- This is not supported in client libraries, so we will use the  `requests` module to make a REST request\n",
        "- Documentation: https://cloud.google.com/generative-ai-app-builder/docs/bring-embeddings#global"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ea9cda23c98"
      },
      "outputs": [],
      "source": [
        "access_token = (\n",
        "    subprocess.check_output([\"gcloud\", \"auth\", \"print-access-token\"])\n",
        "    .decode(\"utf-8\")\n",
        "    .strip()\n",
        ")\n",
        "\n",
        "response = requests.patch(\n",
        "    url=f\"https://discoveryengine.googleapis.com/v1alpha/projects/{PROJECT_ID}/locations/{DATA_STORE_LOCATION}/collections/default_collection/dataStores/{DATA_STORE_ID}/servingConfigs/default_search?updateMask=embeddingConfig,rankingExpression\",\n",
        "    headers={\n",
        "        \"Authorization\": f\"Bearer {access_token}\",\n",
        "        \"Content-Type\": \"application/json; charset=utf-8\",\n",
        "        \"X-Goog-User-Project\": PROJECT_ID,\n",
        "    },\n",
        "    json={\n",
        "        \"name\": f\"projects/{PROJECT_ID}/locations/{DATA_STORE_LOCATION}/collections/default_collection/dataStores/{DATA_STORE_ID}/servingConfigs/default_search\",\n",
        "        \"embeddingConfig\": {\"fieldPath\": EMBEDDINGS_FIELD_NAME},\n",
        "        \"ranking_expression\": \"0.5 * relevance_score\",\n",
        "    },\n",
        ")\n",
        "\n",
        "print(response.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b4dbe6e41fb9"
      },
      "source": [
        "## Test Search App\n",
        "\n",
        "Make a sample query to check how the results look."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f0892ff0bf0e"
      },
      "outputs": [],
      "source": [
        "def search_data_store(\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    data_store_id: str,\n",
        "    search_query: str,\n",
        ") -> list[discoveryengine.SearchResponse]:\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": null,
      "metadata": {
        "id": "64ba2ea463f5"
      },
      "outputs": [],
      "source": [
        "search_query = \"How do I create an array in Java?\"\n",
        "\n",
        "response = search_data_store(\n",
        "    PROJECT_ID, DATA_STORE_LOCATION, DATA_STORE_ID, search_query\n",
        ")\n",
        "\n",
        "print(f\"Summary: {response.summary.summary_text}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d11679ae854d"
      },
      "source": [
        "## Deploy search engine\n",
        "\n",
        "This search engine can now be deployed to a web page using the prebuilt [search widget](https://cloud.google.com/generative-ai-app-builder/docs/add-widget)\n",
        "\n",
        "For a deployed example of the Stack Overflow custom embeddings search engine, go to [vertex-ai-search.web.app](https://vertex-ai-search.web.app/)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "custom_embeddings.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}