From cc5497660551bfce95311955c4dadef6a3fe619b Mon Sep 17 00:00:00 2001
From: Tyler Ohlsen <ohltyler@amazon.com>
Date: Tue, 20 Aug 2024 15:51:53 -0700
Subject: [PATCH] Add basic preset queries; onboard hybrid/multimodal search
 use cases (#302)

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
---
 common/constants.ts                           |  96 +++++++++++++--
 common/interfaces.ts                          |   5 +
 .../pages/workflow_detail/workflow_detail.tsx |   6 +-
 .../configure_search_request.tsx              |  37 +-----
 .../search_inputs/edit_query_modal.tsx        | 112 ++++++++++++++++++
 .../workflow_inputs/workflow_inputs.tsx       |   5 +-
 .../import_workflow/import_workflow_modal.tsx |   4 +-
 public/pages/workflows/new_workflow/utils.ts  |  47 +++++++-
 .../workflows/workflow_list/workflow_list.tsx |  10 ++
 public/pages/workflows/workflows.tsx          |   8 +-
 server/resources/templates/hybrid_search.json |  14 +++
 .../templates/multimodal_search.json          |  14 +++
 12 files changed, 303 insertions(+), 55 deletions(-)
 create mode 100644 public/pages/workflow_detail/workflow_inputs/search_inputs/edit_query_modal.tsx
 create mode 100644 server/resources/templates/hybrid_search.json
 create mode 100644 server/resources/templates/multimodal_search.json

diff --git a/common/constants.ts b/common/constants.ts
index 338964be..45843895 100644
--- a/common/constants.ts
+++ b/common/constants.ts
@@ -3,7 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { WORKFLOW_STATE } from './interfaces';
+import { QueryPreset, WORKFLOW_STATE } from './interfaces';
+import { customStringify } from './utils';
 
 export const PLUGIN_ID = 'flow-framework';
 export const SEARCH_STUDIO = 'Search Studio';
@@ -59,6 +60,8 @@ export const SEARCH_MODELS_NODE_API_PATH = `${BASE_MODEL_NODE_API_PATH}/search`;
 // frontend-specific workflow types, derived from the available preset templates
 export enum WORKFLOW_TYPE {
   SEMANTIC_SEARCH = 'Semantic search',
+  MULTIMODAL_SEARCH = 'Multimodal search',
+  HYBRID_SEARCH = 'Hybrid search',
   CUSTOM = 'Custom',
   UNKNOWN = 'Unknown',
 }
@@ -142,6 +145,91 @@ export const FIXED_TOKEN_LENGTH_OPTIONAL_FIELDS = [
 export const DELIMITER_OPTIONAL_FIELDS = ['delimiter'];
 export const SHARED_OPTIONAL_FIELDS = ['max_chunk_limit', 'description', 'tag'];
 
+/**
+ * QUERIES
+ */
+export const FETCH_ALL_QUERY = {
+  query: {
+    match_all: {},
+  },
+  size: 1000,
+};
+export const SEMANTIC_SEARCH_QUERY = {
+  _source: {
+    excludes: [`{{vector_field}}`],
+  },
+  query: {
+    neural: {
+      [`{{vector_field}}`]: {
+        query_text: `{{query_text}}`,
+        model_id: `{{model_id}}`,
+        k: 100,
+      },
+    },
+  },
+};
+export const MULTIMODAL_SEARCH_QUERY = {
+  _source: {
+    excludes: [`{{vector_field}}`],
+  },
+  query: {
+    neural: {
+      [`{{vector_field}}`]: {
+        query_text: `{{query_text}}`,
+        query_image: `{{query_image}}`,
+        model_id: `{{model_id}}`,
+        k: 100,
+      },
+    },
+  },
+};
+export const HYBRID_SEARCH_QUERY = {
+  _source: {
+    excludes: [`{{vector_field}}`],
+  },
+  query: {
+    hybrid: {
+      queries: [
+        {
+          match: {
+            [`{{text_field}}`]: {
+              query: `{{query_text}}`,
+            },
+          },
+        },
+        {
+          neural: {
+            [`{{vector_field}}`]: {
+              query_text: `{{query_text}}`,
+              model_id: `{{model_id}}`,
+              k: 5,
+            },
+          },
+        },
+      ],
+    },
+  },
+};
+
+export const QUERY_PRESETS = [
+  {
+    name: 'Fetch all',
+    query: customStringify(FETCH_ALL_QUERY),
+  },
+  {
+    name: WORKFLOW_TYPE.SEMANTIC_SEARCH,
+    query: customStringify(SEMANTIC_SEARCH_QUERY),
+  },
+  {
+    name: WORKFLOW_TYPE.MULTIMODAL_SEARCH,
+    query: customStringify(MULTIMODAL_SEARCH_QUERY),
+  },
+  {
+    name: WORKFLOW_TYPE.HYBRID_SEARCH,
+    query: customStringify(HYBRID_SEARCH_QUERY),
+  },
+] as QueryPreset[];
+
 /**
  * MISCELLANEOUS
  */
@@ -152,12 +240,6 @@ export const DEFAULT_NEW_WORKFLOW_STATE = WORKFLOW_STATE.NOT_STARTED;
 export const DEFAULT_NEW_WORKFLOW_STATE_TYPE = ('NOT_STARTED' as any) as typeof WORKFLOW_STATE;
 export const DATE_FORMAT_PATTERN = 'MM/DD/YY hh:mm A';
 export const EMPTY_FIELD_STRING = '--';
-export const FETCH_ALL_QUERY_BODY = {
-  query: {
-    match_all: {},
-  },
-  size: 1000,
-};
 export const INDEX_NOT_FOUND_EXCEPTION = 'index_not_found_exception';
 export const ERROR_GETTING_WORKFLOW_MSG = 'Failed to retrieve template';
 export const NO_MODIFICATIONS_FOUND_TEXT =
diff --git a/common/interfaces.ts b/common/interfaces.ts
index 57481697..6128c6fe 100644
--- a/common/interfaces.ts
+++ b/common/interfaces.ts
@@ -456,6 +456,11 @@ export type WorkflowDict = {
   [workflowId: string]: Workflow;
 };
 
+export type QueryPreset = {
+  name: string;
+  query: string;
+};
+
 /**
  ********** OPENSEARCH TYPES/INTERFACES ************
  */
diff --git a/public/pages/workflow_detail/workflow_detail.tsx b/public/pages/workflow_detail/workflow_detail.tsx
index 1b337cca..29b87dab 100644
--- a/public/pages/workflow_detail/workflow_detail.tsx
+++ b/public/pages/workflow_detail/workflow_detail.tsx
@@ -16,7 +16,7 @@ import {
   EuiPage,
   EuiPageBody,
 } from '@elastic/eui';
-import { APP_PATH, BREADCRUMBS, SHOW_ACTIONS_IN_HEADER} from '../../utils';
+import { APP_PATH, BREADCRUMBS, SHOW_ACTIONS_IN_HEADER } from '../../utils';
 import { getCore } from '../../services';
 import { WorkflowDetailHeader } from './components';
 import {
@@ -28,7 +28,7 @@ import {
 import { ResizableWorkspace } from './resizable_workspace';
 import {
   ERROR_GETTING_WORKFLOW_MSG,
-  FETCH_ALL_QUERY_BODY,
+  FETCH_ALL_QUERY,
   MAX_WORKFLOW_NAME_TO_DISPLAY,
   getCharacterLimitedString,
 } from '../../../common';
@@ -102,7 +102,7 @@ export function WorkflowDetail(props: WorkflowDetailProps) {
   // - fetch available models as their IDs may be used when building flows
   useEffect(() => {
     dispatch(getWorkflow({ workflowId, dataSourceId }));
-    dispatch(searchModels({ apiBody: FETCH_ALL_QUERY_BODY, dataSourceId }));
+    dispatch(searchModels({ apiBody: FETCH_ALL_QUERY, dataSourceId }));
   }, []);
 
   return errorMessage.includes(ERROR_GETTING_WORKFLOW_MSG) ? (
diff --git a/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx b/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx
index 92712fc6..528554eb 100644
--- a/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/search_inputs/configure_search_request.tsx
@@ -12,11 +12,6 @@ import {
   EuiFlexGroup,
   EuiFlexItem,
   EuiFormRow,
-  EuiModal,
-  EuiModalBody,
-  EuiModalFooter,
-  EuiModalHeader,
-  EuiModalHeaderTitle,
   EuiSuperSelect,
   EuiSuperSelectOption,
   EuiText,
@@ -31,6 +26,7 @@ import {
   useAppDispatch,
 } from '../../../../store';
 import { getDataSourceId } from '../../../../utils/utils';
+import { EditQueryModal } from './edit_query_modal';
 
 interface ConfigureSearchRequestProps {
   setQuery: (query: string) => void;
@@ -88,33 +84,10 @@ export function ConfigureSearchRequest(props: ConfigureSearchRequestProps) {
   return (
     <>
       {isEditModalOpen && (
-        <EuiModal
-          onClose={() => setIsEditModalOpen(false)}
-          style={{ width: '70vw' }}
-        >
-          <EuiModalHeader>
-            <EuiModalHeaderTitle>
-              <p>{`Edit query`}</p>
-            </EuiModalHeaderTitle>
-          </EuiModalHeader>
-          <EuiModalBody>
-            <JsonField
-              label="Query"
-              fieldPath={'search.request'}
-              editorHeight="25vh"
-              readOnly={false}
-            />
-          </EuiModalBody>
-          <EuiModalFooter>
-            <EuiButton
-              onClick={() => setIsEditModalOpen(false)}
-              fill={false}
-              color="primary"
-            >
-              Close
-            </EuiButton>
-          </EuiModalFooter>
-        </EuiModal>
+        <EditQueryModal
+          setModalOpen={setIsEditModalOpen}
+          queryFieldPath="search.request"
+        />
       )}
       <EuiFlexGroup direction="column">
         <EuiFlexItem grow={false}>
diff --git a/public/pages/workflow_detail/workflow_inputs/search_inputs/edit_query_modal.tsx b/public/pages/workflow_detail/workflow_inputs/search_inputs/edit_query_modal.tsx
new file mode 100644
index 00000000..e6f2cb2e
--- /dev/null
+++ b/public/pages/workflow_detail/workflow_inputs/search_inputs/edit_query_modal.tsx
@@ -0,0 +1,112 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React, { useState, useEffect } from 'react';
+import { getIn, useFormikContext } from 'formik';
+import {
+  EuiButton,
+  EuiModal,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiModalHeader,
+  EuiModalHeaderTitle,
+  EuiSpacer,
+  EuiSuperSelect,
+  EuiSuperSelectOption,
+  EuiText,
+} from '@elastic/eui';
+import { JsonField } from '../input_fields';
+import {
+  QUERY_PRESETS,
+  QueryPreset,
+  WorkflowFormValues,
+} from '../../../../../common';
+
+interface EditQueryModalProps {
+  queryFieldPath: string;
+  setModalOpen(isOpen: boolean): void;
+}
+
+/**
+ * Basic modal for configuring a query. Provides a dropdown to select from
+ * a set of pre-defined queries targeted for different use cases.
+ */
+export function EditQueryModal(props: EditQueryModalProps) {
+  // Form state
+  const { values, setFieldValue } = useFormikContext<WorkflowFormValues>();
+
+  // selected preset state
+  const [queryPreset, setQueryPreset] = useState<string | undefined>(undefined);
+
+  // if the current query matches some preset, display the preset name as the selected
+  // option in the dropdown. only execute when first rendering so it isn't triggered
+  // when users are updating the underlying value in the JSON editor.
+  useEffect(() => {
+    setQueryPreset(
+      QUERY_PRESETS.find(
+        (preset) => preset.query === getIn(values, props.queryFieldPath)
+      )?.name
+    );
+  }, []);
+
+  return (
+    <EuiModal
+      onClose={() => props.setModalOpen(false)}
+      style={{ width: '70vw' }}
+    >
+      <EuiModalHeader>
+        <EuiModalHeaderTitle>
+          <p>{`Edit query`}</p>
+        </EuiModalHeaderTitle>
+      </EuiModalHeader>
+      <EuiModalBody>
+        <EuiText color="subdued">
+          Start with a preset or enter manually.
+        </EuiText>{' '}
+        <EuiSpacer size="s" />
+        <EuiSuperSelect
+          options={QUERY_PRESETS.map(
+            (preset: QueryPreset) =>
+              ({
+                value: preset.name,
+                inputDisplay: (
+                  <>
+                    <EuiText size="s">{preset.name}</EuiText>
+                  </>
+                ),
+                dropdownDisplay: <EuiText size="s">{preset.name}</EuiText>,
+                disabled: false,
+              } as EuiSuperSelectOption<string>)
+          )}
+          valueOfSelected={queryPreset || ''}
+          onChange={(option: string) => {
+            setQueryPreset(option);
+            setFieldValue(
+              props.queryFieldPath,
+              QUERY_PRESETS.find((preset) => preset.name === option)?.query
+            );
+          }}
+          isInvalid={false}
+        />
+        <EuiSpacer size="s" />
+        <JsonField
+          label="Query"
+          fieldPath={props.queryFieldPath}
+          editorHeight="25vh"
+          readOnly={false}
+        />
+      </EuiModalBody>
+      <EuiModalFooter>
+        <EuiButton
+          onClick={() => props.setModalOpen(false)}
+          fill={false}
+          color="primary"
+        >
+          Close
+        </EuiButton>
+      </EuiModalFooter>
+    </EuiModal>
+  );
+}
diff --git a/public/pages/workflow_detail/workflow_inputs/workflow_inputs.tsx b/public/pages/workflow_detail/workflow_inputs/workflow_inputs.tsx
index 6c178901..2279ccb9 100644
--- a/public/pages/workflow_detail/workflow_inputs/workflow_inputs.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/workflow_inputs.tsx
@@ -4,7 +4,6 @@
  */
 
 import React, { useCallback, useEffect, useState } from 'react';
-import { useSelector } from 'react-redux';
 import { getIn, useFormikContext } from 'formik';
 import { debounce, isEmpty, isEqual } from 'lodash';
 import {
@@ -814,7 +813,9 @@ export function WorkflowInputs(props: WorkflowInputsProps) {
                       <EuiFlexItem grow={false}>
                         <EuiButton
                           disabled={
-                            isRunningSearch || isProposingNoSearchResources
+                            isRunningSearch ||
+                            (isProposingNoSearchResources &&
+                              hasProvisionedSearchResources(props.workflow))
                           }
                           isLoading={isRunningSearch}
                           fill={false}
diff --git a/public/pages/workflows/import_workflow/import_workflow_modal.tsx b/public/pages/workflows/import_workflow/import_workflow_modal.tsx
index 1e1d4bf8..436044dd 100644
--- a/public/pages/workflows/import_workflow/import_workflow_modal.tsx
+++ b/public/pages/workflows/import_workflow/import_workflow_modal.tsx
@@ -30,7 +30,7 @@ import {
   searchWorkflows,
   useAppDispatch,
 } from '../../../store';
-import { FETCH_ALL_QUERY_BODY, Workflow } from '../../../../common';
+import { FETCH_ALL_QUERY, Workflow } from '../../../../common';
 import { WORKFLOWS_TAB } from '../workflows';
 import { getDataSourceId } from '../../../utils/utils';
 
@@ -151,7 +151,7 @@ export function ImportWorkflowModal(props: ImportWorkflowModalProps) {
                 const { workflow } = result;
                 dispatch(
                   searchWorkflows({
-                    apiBody: FETCH_ALL_QUERY_BODY,
+                    apiBody: FETCH_ALL_QUERY,
                     dataSourceId,
                   })
                 );
diff --git a/public/pages/workflows/new_workflow/utils.ts b/public/pages/workflows/new_workflow/utils.ts
index 1d88febe..b9675bdc 100644
--- a/public/pages/workflows/new_workflow/utils.ts
+++ b/public/pages/workflows/new_workflow/utils.ts
@@ -11,8 +11,11 @@ import {
   DEFAULT_NEW_WORKFLOW_NAME,
   UIState,
   WORKFLOW_TYPE,
-  FETCH_ALL_QUERY_BODY,
+  FETCH_ALL_QUERY,
   customStringify,
+  SEMANTIC_SEARCH_QUERY,
+  MULTIMODAL_SEARCH_QUERY,
+  HYBRID_SEARCH_QUERY,
 } from '../../../../common';
 import { generateId } from '../../../utils';
 
@@ -26,7 +29,14 @@ export function enrichPresetWorkflowWithUiMetadata(
       uiMetadata = fetchSemanticSearchMetadata();
       break;
     }
-    // TODO: add more presets
+    case WORKFLOW_TYPE.MULTIMODAL_SEARCH: {
+      uiMetadata = fetchMultimodalSearchMetadata();
+      break;
+    }
+    case WORKFLOW_TYPE.HYBRID_SEARCH: {
+      uiMetadata = fetchHybridSearchMetadata();
+      break;
+    }
     default: {
       uiMetadata = fetchEmptyMetadata();
       break;
@@ -83,7 +93,7 @@ function fetchEmptyMetadata(): UIState {
         request: {
           id: 'request',
           type: 'json',
-          value: customStringify(FETCH_ALL_QUERY_BODY),
+          value: customStringify(FETCH_ALL_QUERY),
         },
         pipelineName: {
           id: 'pipelineName',
@@ -108,8 +118,6 @@ function fetchEmptyMetadata(): UIState {
 }
 
 function fetchSemanticSearchMetadata(): UIState {
-  // We can reuse the base state. Only need to override a few things,
-  // such as preset ingest processors.
   let baseState = fetchEmptyMetadata();
   baseState.type = WORKFLOW_TYPE.SEMANTIC_SEARCH;
   baseState.config.ingest.enrich.processors = [new MLIngestProcessor().toObj()];
@@ -117,6 +125,35 @@ function fetchSemanticSearchMetadata(): UIState {
   baseState.config.ingest.index.settings.value = customStringify({
     [`index.knn`]: true,
   });
+  baseState.config.search.request.value = customStringify(
+    SEMANTIC_SEARCH_QUERY
+  );
+  return baseState;
+}
+
+function fetchMultimodalSearchMetadata(): UIState {
+  let baseState = fetchEmptyMetadata();
+  baseState.type = WORKFLOW_TYPE.MULTIMODAL_SEARCH;
+  baseState.config.ingest.enrich.processors = [new MLIngestProcessor().toObj()];
+  baseState.config.ingest.index.name.value = 'my-knn-index';
+  baseState.config.ingest.index.settings.value = customStringify({
+    [`index.knn`]: true,
+  });
+  baseState.config.search.request.value = customStringify(
+    MULTIMODAL_SEARCH_QUERY
+  );
+  return baseState;
+}
+
+function fetchHybridSearchMetadata(): UIState {
+  let baseState = fetchEmptyMetadata();
+  baseState.type = WORKFLOW_TYPE.HYBRID_SEARCH;
+  baseState.config.ingest.enrich.processors = [new MLIngestProcessor().toObj()];
+  baseState.config.ingest.index.name.value = 'my-knn-index';
+  baseState.config.ingest.index.settings.value = customStringify({
+    [`index.knn`]: true,
+  });
+  baseState.config.search.request.value = customStringify(HYBRID_SEARCH_QUERY);
   return baseState;
 }
 
diff --git a/public/pages/workflows/workflow_list/workflow_list.tsx b/public/pages/workflows/workflow_list/workflow_list.tsx
index 16f52b4d..543e7e62 100644
--- a/public/pages/workflows/workflow_list/workflow_list.tsx
+++ b/public/pages/workflows/workflow_list/workflow_list.tsx
@@ -51,6 +51,16 @@ const filterOptions = [
     checked: 'on',
   } as EuiFilterSelectItem,
   // @ts-ignore
+  {
+    name: WORKFLOW_TYPE.MULTIMODAL_SEARCH,
+    checked: 'on',
+  } as EuiFilterSelectItem,
+  // @ts-ignore
+  {
+    name: WORKFLOW_TYPE.HYBRID_SEARCH,
+    checked: 'on',
+  } as EuiFilterSelectItem,
+  // @ts-ignore
   {
     name: WORKFLOW_TYPE.CUSTOM,
     checked: 'on',
diff --git a/public/pages/workflows/workflows.tsx b/public/pages/workflows/workflows.tsx
index 401c0b41..6025e402 100644
--- a/public/pages/workflows/workflows.tsx
+++ b/public/pages/workflows/workflows.tsx
@@ -25,7 +25,7 @@ import { WorkflowList } from './workflow_list';
 import { NewWorkflow } from './new_workflow';
 import { AppState, searchWorkflows, useAppDispatch } from '../../store';
 import { EmptyListMessage } from './empty_list_message';
-import { FETCH_ALL_QUERY_BODY, SEARCH_STUDIO } from '../../../common';
+import { FETCH_ALL_QUERY, SEARCH_STUDIO } from '../../../common';
 import { ImportWorkflowModal } from './import_workflow';
 import { MountPoint } from '../../../../../src/core/public';
 import { DataSourceSelectableConfig } from '../../../../../src/plugins/data_source_management/public';
@@ -119,7 +119,7 @@ export function Workflows(props: WorkflowsProps) {
     if (selectedTabId === WORKFLOWS_TAB.MANAGE) {
       dispatch(
         searchWorkflows({
-          apiBody: FETCH_ALL_QUERY_BODY,
+          apiBody: FETCH_ALL_QUERY,
           dataSourceId: dataSourceId,
         })
       );
@@ -141,7 +141,7 @@ export function Workflows(props: WorkflowsProps) {
   useEffect(() => {
     dispatch(
       searchWorkflows({
-        apiBody: FETCH_ALL_QUERY_BODY,
+        apiBody: FETCH_ALL_QUERY,
         dataSourceId: dataSourceId,
       })
     );
@@ -161,7 +161,7 @@ export function Workflows(props: WorkflowsProps) {
     }
     dispatch(
       searchWorkflows({
-        apiBody: FETCH_ALL_QUERY_BODY,
+        apiBody: FETCH_ALL_QUERY,
         dataSourceId: dataSourceId,
       })
     );
diff --git a/server/resources/templates/hybrid_search.json b/server/resources/templates/hybrid_search.json
new file mode 100644
index 00000000..f1768c16
--- /dev/null
+++ b/server/resources/templates/hybrid_search.json
@@ -0,0 +1,14 @@
+{
+    "name": "Hybrid Search",
+    "description": "A basic workflow containing the ingest pipeline and index configurations for performing hybrid search",
+    "version": {
+        "template": "1.0.0",
+        "compatibility": [
+            "2.17.0",
+            "3.0.0"
+        ]
+    },
+    "ui_metadata": {
+        "type": "Hybrid search"
+    }
+}
\ No newline at end of file
diff --git a/server/resources/templates/multimodal_search.json b/server/resources/templates/multimodal_search.json
new file mode 100644
index 00000000..409a303d
--- /dev/null
+++ b/server/resources/templates/multimodal_search.json
@@ -0,0 +1,14 @@
+{
+    "name": "Multimodal Search",
+    "description": "A basic workflow containing the ingest pipeline and index configurations for performing multimodal search",
+    "version": {
+        "template": "1.0.0",
+        "compatibility": [
+            "2.17.0",
+            "3.0.0"
+        ]
+    },
+    "ui_metadata": {
+        "type": "Multimodal search"
+    }
+}
\ No newline at end of file