Finish onboarding and edge case handling

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>

Author: ohltyler

public/pages/workflow_detail/prototype/ingestor.tsx

@@ -16,28 +16,24 @@ import {
   USE_CASE,
   Workflow,
   getIndexName,
-  getSemanticSearchValues,
+  getNeuralSearchValues,
 } from '../../../../common';
 import { ingest, useAppDispatch } from '../../../store';
 import { getCore } from '../../../services';
-import { getFormattedJSONString } from './utils';
+import {
+  NeuralSparseValues,
+  SemanticSearchValues,
+  WorkflowValues,
+  getFormattedJSONString,
+} from './utils';
 
 interface IngestorProps {
   workflow: Workflow;
 }
 
-type WorkflowValues = {
-  modelId: string;
-};
-
-type SemanticSearchValues = WorkflowValues & {
-  inputField: string;
-  vectorField: string;
-};
-
 type DocGeneratorFn = (
   queryText: string,
-  workflowValues: SemanticSearchValues
+  workflowValues: SemanticSearchValues | NeuralSparseValues
 ) => {};
 
 /**
@@ -188,8 +184,9 @@ function getDocGeneratorFn(workflow: Workflow): DocGeneratorFn {
   let fn;
   switch (workflow.use_case) {
     case USE_CASE.SEMANTIC_SEARCH:
+    case USE_CASE.NEURAL_SPARSE_SEARCH:
     default: {
-      fn = () => generateSemanticSearchDoc;
+      fn = () => generateNeuralSearchDoc;
     }
   }
   return fn;
@@ -200,17 +197,18 @@ function getWorkflowValues(workflow: Workflow): WorkflowValues {
   let values;
   switch (workflow.use_case) {
     case USE_CASE.SEMANTIC_SEARCH:
+    case USE_CASE.NEURAL_SPARSE_SEARCH:
     default: {
-      values = getSemanticSearchValues(workflow);
+      values = getNeuralSearchValues(workflow);
     }
   }
   return values;
 }
 
-// utility fn to generate a document suited for semantic search
-function generateSemanticSearchDoc(
+// utility fn to generate a document suited for neural search use cases
+function generateNeuralSearchDoc(
   docValue: string,
-  workflowValues: SemanticSearchValues
+  workflowValues: SemanticSearchValues | NeuralSparseValues
 ): {} {
   return {
     [workflowValues.inputField]: docValue,

public/pages/workflow_detail/prototype/query_executor.tsx

@@ -16,28 +16,24 @@ import {
   USE_CASE,
   Workflow,
   getIndexName,
-  getSemanticSearchValues,
+  getNeuralSearchValues,
 } from '../../../../common';
 import { searchIndex, useAppDispatch } from '../../../store';
 import { getCore } from '../../../services';
-import { getFormattedJSONString } from './utils';
+import {
+  NeuralSparseValues,
+  SemanticSearchValues,
+  WorkflowValues,
+  getFormattedJSONString,
+} from './utils';
 
 interface QueryExecutorProps {
   workflow: Workflow;
 }
 
-type WorkflowValues = {
-  modelId: string;
-};
-
-type SemanticSearchValues = WorkflowValues & {
-  inputField: string;
-  vectorField: string;
-};
-
 type QueryGeneratorFn = (
   queryText: string,
-  workflowValues: SemanticSearchValues
+  workflowValues: SemanticSearchValues | NeuralSparseValues
 ) => {};
 
 /**
@@ -187,9 +183,13 @@ export function QueryExecutor(props: QueryExecutorProps) {
 function getQueryGeneratorFn(workflow: Workflow): QueryGeneratorFn {
   let fn;
   switch (workflow.use_case) {
-    case USE_CASE.SEMANTIC_SEARCH:
-    default: {
+    case USE_CASE.SEMANTIC_SEARCH: {
       fn = () => generateSemanticSearchQuery;
+      break;
+    }
+    case USE_CASE.NEURAL_SPARSE_SEARCH:
+    default: {
+      fn = () => generateNeuralSparseQuery;
     }
   }
   return fn;
@@ -201,7 +201,7 @@ function getWorkflowValues(workflow: Workflow): WorkflowValues {
   switch (workflow.use_case) {
     case USE_CASE.SEMANTIC_SEARCH:
     default: {
-      values = getSemanticSearchValues(workflow);
+      values = getNeuralSearchValues(workflow);
     }
   }
   return values;
@@ -213,6 +213,7 @@ function generateSemanticSearchQuery(
   workflowValues: SemanticSearchValues
 ): {} {
   return {
+    // TODO: can make this configurable
     _source: {
       excludes: [`${workflowValues.vectorField}`],
     },
@@ -221,13 +222,35 @@ function generateSemanticSearchQuery(
         [workflowValues.vectorField]: {
           query_text: queryText,
           model_id: workflowValues.modelId,
+          // TODO: expose k as configurable
           k: 5,
         },
       },
     },
   };
 }
 
+// utility fn to generate a neural sparse search query
+function generateNeuralSparseQuery(
+  queryText: string,
+  workflowValues: NeuralSparseValues
+): {} {
+  return {
+    // TODO: can make this configurable
+    _source: {
+      excludes: [`${workflowValues.vectorField}`],
+    },
+    query: {
+      neural_sparse: {
+        [workflowValues.vectorField]: {
+          query_text: queryText,
+          model_id: workflowValues.modelId,
+        },
+      },
+    },
+  };
+}
+
 function processHits(hits: any[]): {}[] {
   return hits.map((hit) => hit._source);
 }

public/pages/workflow_detail/prototype/utils.ts

@@ -4,9 +4,22 @@
  */
 
 /**
- * Shared utility fns used in the prototyping page.
+ * Shared utility fns and constants used in the prototyping page.
  */
 
+// UTILITY FNS
 export function getFormattedJSONString(obj: {}): string {
   return Object.values(obj).length > 0 ? JSON.stringify(obj, null, '\t') : '';
 }
+
+// CONSTANTS
+export type WorkflowValues = {
+  modelId: string;
+};
+
+export type SemanticSearchValues = WorkflowValues & {
+  inputField: string;
+  vectorField: string;
+};
+
+export type NeuralSparseValues = SemanticSearchValues;

public/pages/workflow_detail/utils/data_extractor_utils.ts

@@ -33,7 +33,9 @@ export function getIndexName(workflow: Workflow): string | undefined {
   }
 }
 
-export function getSemanticSearchValues(
+// Returns values for neural search use cases. Note many of them
+// persist the same values to use during ingest and search, so we keep the naming general
+export function getNeuralSearchValues(
   workflow: Workflow
 ): { modelId: string; inputField: string; vectorField: string } {
   const modelId = getModelId(workflow) as string;

public/pages/workflow_detail/utils/workflow_to_template_utils.ts

@@ -154,6 +154,12 @@ function transformerToTemplateNodes(
       const modelId = model.id;
       const ingestPipelineName = generateId('ingest_pipeline');
 
+      // register model workflow step type is different per use case
+      const registerModelStepType =
+        flowNode.data.type === COMPONENT_CLASS.TEXT_EMBEDDING_TRANSFORMER
+          ? REGISTER_LOCAL_PRETRAINED_MODEL_STEP_TYPE
+          : REGISTER_LOCAL_SPARSE_ENCODING_MODEL_STEP_TYPE;
+
       let registerModelStep = undefined as
         | RegisterPretrainedModelNode
         | undefined;
@@ -170,15 +176,9 @@ function transformerToTemplateNodes(
           (model) => model.name === modelId
         ) as PretrainedSentenceTransformer;
 
-        // workflow step type is different per use case
-        const stepType =
-          flowNode.data.type === COMPONENT_CLASS.TEXT_EMBEDDING_TRANSFORMER
-            ? REGISTER_LOCAL_PRETRAINED_MODEL_STEP_TYPE
-            : REGISTER_LOCAL_SPARSE_ENCODING_MODEL_STEP_TYPE;
-
         registerModelStep = {
-          id: stepType,
-          type: stepType,
+          id: registerModelStepType,
+          type: registerModelStepType,
           user_inputs: {
             name: pretrainedModel.name,
             description: pretrainedModel.description,
@@ -193,7 +193,7 @@ function transformerToTemplateNodes(
       // or directly from the user
       const finalModelId =
         registerModelStep !== undefined
-          ? `\${{${REGISTER_LOCAL_PRETRAINED_MODEL_STEP_TYPE}.model_id}}`
+          ? `\${{${registerModelStepType}.model_id}}`
           : modelId;
 
       // processor is different per use case
@@ -236,6 +236,12 @@ function transformerToTemplateNodes(
           },
         },
       } as CreateIngestPipelineNode;
+      if (registerModelStep !== undefined) {
+        createIngestPipelineStep.previous_node_inputs = {
+          ...createIngestPipelineStep.previous_node_inputs,
+          [registerModelStepType]: 'model_id',
+        };
+      }
 
       return registerModelStep !== undefined
         ? [registerModelStep, createIngestPipelineStep]

public/pages/workflows/workflow_list/columns.tsx

@@ -16,6 +16,7 @@ export const columns = (actions: any[]) => [
   {
     field: 'name',
     name: 'Name',
+    width: '20%',
     sortable: true,
     render: (name: string, workflow: Workflow) => (
       <EuiLink href={`${PLUGIN_ID}#/workflows/${workflow.id}`}>{name}</EuiLink>
@@ -29,6 +30,7 @@ export const columns = (actions: any[]) => [
   {
     field: 'use_case',
     name: 'Type',
+    width: '30%',
     sortable: true,
   },
   {