Onboard search flows for remaining use cases

Signed-off-by: Tyler Ohlsen <[email protected]>
opensearch-project · Apr 22, 2024 · 78731db · 78731db
1 parent 75d8c8f
commit 78731db
Show file tree

Hide file tree

Showing 4 changed files with 205 additions and 6 deletions.
diff --git a/public/component_types/transformer/sparse_encoder_transformer.ts b/public/component_types/transformer/sparse_encoder_transformer.ts
@@ -25,6 +25,12 @@ export class SparseEncoderTransformer extends MLTransformer {
         baseClass: COMPONENT_CLASS.DOCUMENT,
         acceptMultiple: false,
       },
+      {
+        id: 'query',
+        label: 'Query',
+        baseClass: COMPONENT_CLASS.QUERY,
+        acceptMultiple: false,
+      },
     ];
     this.createFields = [
       {
@@ -59,6 +65,10 @@ export class SparseEncoderTransformer extends MLTransformer {
         label: 'Transformed Document',
         baseClasses: [COMPONENT_CLASS.DOCUMENT],
       },
+      {
+        label: 'Transformed Query',
+        baseClasses: [COMPONENT_CLASS.QUERY],
+      },
     ];
   }
 }
diff --git a/public/pages/workflow_detail/utils/data_extractor_utils.ts b/public/pages/workflow_detail/utils/data_extractor_utils.ts
@@ -85,8 +85,8 @@ function getModelId(workflow: Workflow): string | undefined {
       if (model.category === MODEL_CATEGORY.PRETRAINED) {
         const modelResource = workflow.resourcesCreated?.find(
           (resource) => resource.type === WORKFLOW_RESOURCE_TYPE.MODEL_ID
-        ) as WorkflowResource;
-        return modelResource.id;
+        );
+        return modelResource?.id;
       } else {
         return model.id;
       }

diff --git a/public/pages/workflow_detail/workspace/reactflow-styles.scss b/public/pages/workflow_detail/workspace/reactflow-styles.scss
@@ -38,8 +38,10 @@ $handle-color-invalid: $euiColorDanger;
 // Overriding the styling for the reactflow node when it is selected.
 // We need to use important tag to override ReactFlow's wrapNode that sets the box-shadow.
 // Ref: https://github.com/wbkd/react-flow/blob/main/packages/core/src/components/Nodes/wrapNode.tsx#L187
-// TODO: when the node sizing is dynamic (e.g., 'min-height', the shadow only covers the min height
-// instead of the node's final rendered height.
+// TODO: when the node sizing is dynamic (e.g., 'min-height', it causes several issues:
+// 1. the shadow only covers the min height instead of the node's final rendered height
+// 2. the bounding edges of the parent node only fit with the 'min-height' amount, causing
+//    the node to look like it can be drug slightly out of the parent node
 .reactflow-workspace .react-flow__node-custom.selected {
   box-shadow: 0 0 2px 2px rgba(0, 0, 0, 0.5);
   border-radius: 5px;

diff --git a/public/pages/workflows/new_workflow/utils.ts b/public/pages/workflows/new_workflow/utils.ts
@@ -74,8 +74,14 @@ function fetchSemanticSearchWorkspaceFlow(): WorkspaceFlowState {
   const ingestId1 = generateId(COMPONENT_CLASS.TEXT_EMBEDDING_TRANSFORMER);
   const ingestId2 = generateId(COMPONENT_CLASS.KNN_INDEXER);
   const ingestGroupId = generateId(COMPONENT_CATEGORY.INGEST);
+  const searchGroupId = generateId(COMPONENT_CATEGORY.SEARCH);
+  const searchId0 = generateId(COMPONENT_CLASS.NEURAL_QUERY);
+  const searchId1 = generateId(COMPONENT_CLASS.SPARSE_ENCODER_TRANSFORMER);
+  const searchId2 = generateId(COMPONENT_CLASS.KNN_INDEXER);
   const edgeId0 = generateId('edge');
   const edgeId1 = generateId('edge');
+  const edgeId2 = generateId('edge');
+  const edgeId3 = generateId('edge');
 
   const ingestNodes = [
     {
@@ -126,9 +132,58 @@ function fetchSemanticSearchWorkspaceFlow(): WorkspaceFlowState {
       deletable: false,
     },
   ] as ReactFlowComponent[];
+  const searchNodes = [
+    {
+      id: searchGroupId,
+      position: { x: 400, y: 1000 },
+      type: NODE_CATEGORY.SEARCH_GROUP,
+      data: { label: COMPONENT_CATEGORY.SEARCH },
+      style: {
+        width: 1300,
+        height: 400,
+      },
+      className: 'reactflow__group-node__search',
+      selectable: true,
+      draggable: true,
+      deletable: false,
+    },
+    {
+      id: searchId0,
+      position: { x: 100, y: 70 },
+      data: initComponentData(new NeuralQuery().toObj(), searchId0),
+      type: NODE_CATEGORY.CUSTOM,
+      parentNode: searchGroupId,
+      extent: 'parent',
+      draggable: true,
+      deletable: false,
+    },
+    {
+      id: searchId1,
+      position: { x: 500, y: 70 },
+      data: initComponentData(
+        new TextEmbeddingTransformer().toPlaceholderObj(),
+        searchId1
+      ),
+      type: NODE_CATEGORY.CUSTOM,
+      parentNode: searchGroupId,
+      extent: 'parent',
+      draggable: true,
+      deletable: false,
+    },
+    {
+      id: searchId2,
+      position: { x: 900, y: 70 },
+      data: initComponentData(new KnnIndexer().toPlaceholderObj(), searchId2),
+      type: NODE_CATEGORY.CUSTOM,
+      parentNode: searchGroupId,
+      extent: 'parent',
+      draggable: true,
+      deletable: false,
+    },
+  ] as ReactFlowComponent[];
 
   return {
-    nodes: ingestNodes,
+    nodes: [...ingestNodes, ...searchNodes],
     edges: [
       {
         id: edgeId0,
@@ -164,6 +219,44 @@ function fetchSemanticSearchWorkspaceFlow(): WorkspaceFlowState {
         zIndex: 2,
         deletable: false,
       },
+      {
+        id: edgeId2,
+        key: edgeId2,
+        source: searchId0,
+        target: searchId1,
+        sourceClasses: ingestNodes.find((node) => node.id === searchId0)?.data
+          .baseClasses,
+        targetClasses: ingestNodes.find((node) => node.id === searchId1)?.data
+          .baseClasses,
+        sourceHandle: COMPONENT_CLASS.QUERY,
+        targetHandle: COMPONENT_CLASS.QUERY,
+        markerEnd: {
+          type: MarkerType.ArrowClosed,
+          width: 20,
+          height: 20,
+        },
+        zIndex: 2,
+        deletable: false,
+      },
+      {
+        id: edgeId3,
+        key: edgeId3,
+        source: searchId1,
+        target: searchId2,
+        sourceClasses: ingestNodes.find((node) => node.id === searchId1)?.data
+          .baseClasses,
+        targetClasses: ingestNodes.find((node) => node.id === searchId2)?.data
+          .baseClasses,
+        sourceHandle: COMPONENT_CLASS.QUERY,
+        targetHandle: COMPONENT_CLASS.QUERY,
+        markerEnd: {
+          type: MarkerType.ArrowClosed,
+          width: 20,
+          height: 20,
+        },
+        zIndex: 2,
+        deletable: false,
+      },
     ] as ReactFlowEdge[],
   };
 }
@@ -173,8 +266,14 @@ function fetchNeuralSparseSearchWorkspaceFlow(): WorkspaceFlowState {
   const ingestId1 = generateId(COMPONENT_CLASS.SPARSE_ENCODER_TRANSFORMER);
   const ingestId2 = generateId(COMPONENT_CLASS.KNN_INDEXER);
   const ingestGroupId = generateId(COMPONENT_CATEGORY.INGEST);
+  const searchGroupId = generateId(COMPONENT_CATEGORY.SEARCH);
+  const searchId0 = generateId(COMPONENT_CLASS.NEURAL_QUERY);
+  const searchId1 = generateId(COMPONENT_CLASS.SPARSE_ENCODER_TRANSFORMER);
+  const searchId2 = generateId(COMPONENT_CLASS.KNN_INDEXER);
   const edgeId0 = generateId('edge');
   const edgeId1 = generateId('edge');
+  const edgeId2 = generateId('edge');
+  const edgeId3 = generateId('edge');
 
   const ingestNodes = [
     {
@@ -226,8 +325,58 @@ function fetchNeuralSparseSearchWorkspaceFlow(): WorkspaceFlowState {
     },
   ] as ReactFlowComponent[];
 
+  const searchNodes = [
+    {
+      id: searchGroupId,
+      position: { x: 400, y: 1000 },
+      type: NODE_CATEGORY.SEARCH_GROUP,
+      data: { label: COMPONENT_CATEGORY.SEARCH },
+      style: {
+        width: 1300,
+        height: 400,
+      },
+      className: 'reactflow__group-node__search',
+      selectable: true,
+      draggable: true,
+      deletable: false,
+    },
+    {
+      id: searchId0,
+      position: { x: 100, y: 70 },
+      data: initComponentData(new NeuralQuery().toObj(), searchId0),
+      type: NODE_CATEGORY.CUSTOM,
+      parentNode: searchGroupId,
+      extent: 'parent',
+      draggable: true,
+      deletable: false,
+    },
+    {
+      id: searchId1,
+      position: { x: 500, y: 70 },
+      data: initComponentData(
+        new SparseEncoderTransformer().toPlaceholderObj(),
+        searchId1
+      ),
+      type: NODE_CATEGORY.CUSTOM,
+      parentNode: searchGroupId,
+      extent: 'parent',
+      draggable: true,
+      deletable: false,
+    },
+    {
+      id: searchId2,
+      position: { x: 900, y: 70 },
+      data: initComponentData(new KnnIndexer().toPlaceholderObj(), searchId2),
+      type: NODE_CATEGORY.CUSTOM,
+      parentNode: searchGroupId,
+      extent: 'parent',
+      draggable: true,
+      deletable: false,
+    },
+  ] as ReactFlowComponent[];
+
   return {
-    nodes: ingestNodes,
+    nodes: [...ingestNodes, ...searchNodes],
     edges: [
       {
         id: edgeId0,
@@ -263,6 +412,44 @@ function fetchNeuralSparseSearchWorkspaceFlow(): WorkspaceFlowState {
         zIndex: 2,
         deletable: false,
       },
+      {
+        id: edgeId2,
+        key: edgeId2,
+        source: searchId0,
+        target: searchId1,
+        sourceClasses: ingestNodes.find((node) => node.id === searchId0)?.data
+          .baseClasses,
+        targetClasses: ingestNodes.find((node) => node.id === searchId1)?.data
+          .baseClasses,
+        sourceHandle: COMPONENT_CLASS.QUERY,
+        targetHandle: COMPONENT_CLASS.QUERY,
+        markerEnd: {
+          type: MarkerType.ArrowClosed,
+          width: 20,
+          height: 20,
+        },
+        zIndex: 2,
+        deletable: false,
+      },
+      {
+        id: edgeId3,
+        key: edgeId3,
+        source: searchId1,
+        target: searchId2,
+        sourceClasses: ingestNodes.find((node) => node.id === searchId1)?.data
+          .baseClasses,
+        targetClasses: ingestNodes.find((node) => node.id === searchId2)?.data
+          .baseClasses,
+        sourceHandle: COMPONENT_CLASS.QUERY,
+        targetHandle: COMPONENT_CLASS.QUERY,
+        markerEnd: {
+          type: MarkerType.ArrowClosed,
+          width: 20,
+          height: 20,
+        },
+        zIndex: 2,
+        deletable: false,
+      },
     ] as ReactFlowEdge[],
   };
 }