fix: Updates to VAIS Web App (GoogleCloudPlatform#1333)

- Changed Data Store to Engine
- Change all references from genappbuilder to vais

search/web-app/README.md

@@ -72,7 +72,7 @@ Additional features include how to search the public Cloud Knowledge Graph using
      - ![configId](img/configId.png)
      - Be sure to set authorization type as `Public Access` and add your web application url to the `Allowed Domains` once it's deployed.
      - Add the `configId` for your Search Engines to `WIDGET_CONFIGS` in `consts.py`
-   - To use the Custom UI, add the engine id for your search engine to `CUSTOM_UI_DATASTORE_IDS` in `consts.py`
+   - To use the Custom UI, add the engine id for your search engine to `CUSTOM_UI_ENGINE_IDS` in `consts.py`
      - This is the string after `/engines/` in the Cloud Console URL.
        - `https://console.cloud.google.com/gen-app-builder/engines/website-search-engine_1681248733152/...`
        - Engine ID is `website-search-engine_1681248733152`

search/web-app/consts.py

@@ -19,18 +19,28 @@
 
 WIDGET_CONFIGS = [
     {
-        "name": "Contracts (Unstructured)",
+        "name": "Google Cloud Website",
+        "config_id": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
+        "placeholder": "Vertex AI",
+    },
+    {
+        "name": "Contracts",
         "config_id": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
         "placeholder": "What is the SLA?",
     },
     {
-        "name": "Alphabet Earnings Reports (Unstructured)",
+        "name": "Alphabet Earnings Reports",
         "config_id": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
         "placeholder": "What was Google's revenue in 2021?",
     },
+    {
+        "name": "Kaggle Movies",
+        "config_id": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
+        "placeholder": "Harry Potter",
+    },
 ]
 
-CUSTOM_UI_DATASTORE_IDS = [
+CUSTOM_UI_ENGINE_IDS = [
     {
         "name": "Google Cloud Website",
         "engine_id": "google-cloud-site-search",
@@ -39,10 +49,14 @@
         "name": "Google Merchandise Store (Advanced Indexing)",
         "engine_id": "google-merch-store",
     },
+    {
+        "name": "Cymbal Bank",
+        "engine_id": "cymbal-bank-onboarding",
+    },
 ]
 
 
-IMAGE_SEARCH_DATASTORE_IDs = [
+IMAGE_SEARCH_ENGINE_IDs = [
     {
         "name": "Google Merchandise Store",
         "engine_id": "google-merch-store",

search/web-app/main.py

@@ -20,18 +20,18 @@
 from urllib.parse import urlparse
 
 from consts import (
-    CUSTOM_UI_DATASTORE_IDS,
+    CUSTOM_UI_ENGINE_IDS,
     LOCATION,
     PROJECT_ID,
     SUMMARY_MODELS,
     VALID_LANGUAGES,
     WIDGET_CONFIGS,
-    IMAGE_SEARCH_DATASTORE_IDs,
+    IMAGE_SEARCH_ENGINE_IDs,
     RECOMMENDATIONS_DATASTORE_IDs,
 )
 from ekg_utils import search_public_kg
 from flask import Flask, render_template, request
-from genappbuilder_utils import (
+from vais_utils import (
     list_documents,
     recommend_personalize,
     search_enterprise_search,
@@ -49,7 +49,7 @@
     "default_language": VALID_LANGUAGES[0],
 }
 
-CUSTOM_UI_SEARCH_ENGINES = [d["name"] for d in CUSTOM_UI_DATASTORE_IDS]
+CUSTOM_UI_SEARCH_ENGINES = [d["name"] for d in CUSTOM_UI_ENGINE_IDS]
 
 NAV_LINKS = [
     {"link": "/", "name": "Widgets", "icon": "widgets"},
@@ -119,8 +119,8 @@ def search() -> str:
     )
 
 
-@app.route("/search_genappbuilder", methods=["POST"])
-def search_genappbuilder() -> str:
+@app.route("/search_vais", methods=["POST"])
+def search_vais() -> str:
     """
     Handle Search Vertex AI Search Request
     """
@@ -155,7 +155,7 @@ def search_genappbuilder() -> str:
     results, summary, request_url, raw_request, raw_response = search_enterprise_search(
         project_id=PROJECT_ID,
         location=LOCATION,
-        engine_id=CUSTOM_UI_DATASTORE_IDS[int(search_engine)]["engine_id"],
+        engine_id=CUSTOM_UI_ENGINE_IDS[int(search_engine)]["engine_id"],
         search_query=search_query,
         summary_model=summary_model,
         summary_preamble=summary_preamble,
@@ -188,8 +188,8 @@ def image_search() -> str:
     )
 
 
-@app.route("/imagesearch_genappbuilder", methods=["POST"])
-def imagesearch_genappbuilder() -> str:
+@app.route("/imagesearch_vais", methods=["POST"])
+def imagesearch_vais() -> str:
     """
     Handle Image Search Vertex AI Search Request
     """
@@ -232,7 +232,7 @@ def imagesearch_genappbuilder() -> str:
         results, _, request_url, raw_request, raw_response = search_enterprise_search(
             project_id=PROJECT_ID,
             location=LOCATION,
-            engine_id=IMAGE_SEARCH_DATASTORE_IDs[0]["engine_id"],
+            engine_id=IMAGE_SEARCH_ENGINE_IDs[0]["engine_id"],
             search_query=search_query,
             image_bytes=image_bytes,
             params={"search_type": 1},
@@ -270,8 +270,8 @@ def recommend() -> str:
     )
 
 
-@app.route("/recommend_genappbuilder", methods=["POST"])
-def recommend_genappbuilder() -> str:
+@app.route("/recommend_vais", methods=["POST"])
+def recommend_vais() -> str:
     """
     Handle Recommend Vertex AI Search Request
     """

search/web-app/static/index.css

@@ -231,6 +231,6 @@ a:active {
 }
 
 .search-form .mdc-radio {
-  margin-right: 3rem;
+  margin-right: 4rem;
   vertical-align: top;
 }

search/web-app/templates/image-search.html

@@ -19,7 +19,7 @@
 {% block content %}
 <form
   class="search-form"
-  action="/imagesearch_genappbuilder"
+  action="/imagesearch_vais"
   method="post"
   enctype="multipart/form-data">
   <!-- Search Box -->

search/web-app/templates/recommend.html

@@ -13,7 +13,7 @@
 {% block description %}Select a document from the list to view recommendations of similar documents in the dataset.{% endblock %}
 
 {% block content %}
-<form class="search-form" action="/recommend_genappbuilder" method="post">
+<form class="search-form" action="/recommend_vais" method="post">
   <!-- Attribution Token -->
   <input type="hidden" name="attribution_token" value="{{ attribution_token }}" />
   <!-- Document Selection Menu -->

search/web-app/templates/search.html

@@ -11,7 +11,7 @@
 
 {% block title %}Vertex AI Search Demo - {{title}}{% endblock %}
 {% block content %}
-<form class="search-form" action="/search_genappbuilder" method="post">
+<form class="search-form" action="/search_vais" method="post">
   <!-- Search Box -->
   <div class="mdc-text-field mdc-text-field--outlined query-field">
     <input class="mdc-text-field__input" id="query-input" name="search_query" required />

search/web-app/genappbuilder_utils.py → search/web-app/vais_utils.py

@@ -47,8 +47,7 @@ def list_documents(
 def search_enterprise_search(
     project_id: str,
     location: str,
-    data_store_id: str | None = None,
-    engine_id: str | None = None,
+    engine_id: str,
     page_size: int = 50,
     search_query: str | None = None,
     image_bytes: bytes | None = None,
@@ -62,17 +61,7 @@ def search_enterprise_search(
     # Create a client
     client = discoveryengine.SearchServiceClient()
 
-    if data_store_id:
-        serving_config = client.serving_config_path(
-            project=project_id,
-            location=location,
-            data_store=data_store_id,
-            serving_config="default_config",
-        )
-    elif engine_id:
-        serving_config = f"projects/{project_id}/locations/{location}/collections/default_collection/engines/{engine_id}/servingConfigs/default_config"
-    else:
-        raise ValueError("Either `data_store_id` or `engine_id` must be provided.")
+    serving_config = f"projects/{project_id}/locations/{location}/collections/default_collection/engines/{engine_id}/servingConfigs/default_config"
 
     # Configuration options for search
     content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec(