From 58b301338cfae2af168d30c015b11d23c64eb7d8 Mon Sep 17 00:00:00 2001 From: Tyler Ohlsen Date: Mon, 20 May 2024 07:48:14 -0700 Subject: [PATCH] Add initial logic to generate ReactFlow workspace (#160) Signed-off-by: Tyler Ohlsen (cherry picked from commit 73930abe3bda8c2dff7d55638f13f81392d31b57) --- public/component_types/other/results.tsx | 1 + public/component_types/transformer/index.ts | 1 + .../workflow_detail/workspace/workspace.tsx | 14 +- public/pages/workflows/new_workflow/utils.ts | 339 +++++++------- public/utils/utils.ts | 431 ++++++++++++++++++ 5 files changed, 612 insertions(+), 174 deletions(-) diff --git a/public/component_types/other/results.tsx b/public/component_types/other/results.tsx index e2505682..f787c665 100644 --- a/public/component_types/other/results.tsx +++ b/public/component_types/other/results.tsx @@ -16,5 +16,6 @@ export class Results extends BaseComponent { this.type = COMPONENT_CLASS.RESULTS; this.label = 'Results'; this.description = 'OpenSearch results'; + this.inputs = [{ id: 'input', label: 'Input', acceptMultiple: false }]; } } diff --git a/public/component_types/transformer/index.ts b/public/component_types/transformer/index.ts index df8cf864..ce3afc3e 100644 --- a/public/component_types/transformer/index.ts +++ b/public/component_types/transformer/index.ts @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +export * from './ml_transformer'; export * from './text_embedding_transformer'; export * from './sparse_encoder_transformer'; export * from './results_transformer'; diff --git a/public/pages/workflow_detail/workspace/workspace.tsx b/public/pages/workflow_detail/workspace/workspace.tsx index 0429df4f..4a024918 100644 --- a/public/pages/workflow_detail/workspace/workspace.tsx +++ b/public/pages/workflow_detail/workspace/workspace.tsx @@ -23,6 +23,7 @@ import { ReactFlowComponent, ReactFlowEdge, Workflow, + WorkflowConfig, } from '../../../../common'; import { IngestGroupComponent, @@ -30,6 +31,7 @@ import { WorkspaceComponent, } from './workspace_components'; import { DeletableEdge } from './workspace_edge'; +import { uiConfigToWorkspaceFlow } from '../../../utils'; // styling import 'reactflow/dist/style.css'; @@ -107,12 +109,15 @@ export function Workspace(props: WorkspaceProps) { [setEdges] ); - // Initialization. Set the nodes and edges to an existing workflow state, + // Initialization. Generate the nodes and edges based on the workflow config. useEffect(() => { const workflow = { ...props.workflow }; - if (workflow?.ui_metadata?.workspace_flow) { - setNodes(workflow.ui_metadata.workspace_flow.nodes); - setEdges(workflow.ui_metadata.workspace_flow.edges); + if (workflow?.ui_metadata?.config) { + const proposedWorkspaceFlow = uiConfigToWorkspaceFlow( + workflow.ui_metadata?.config as WorkflowConfig + ); + setNodes(proposedWorkspaceFlow.nodes); + setEdges(proposedWorkspaceFlow.edges); } }, [props.workflow]); @@ -141,6 +146,7 @@ export function Workspace(props: WorkspaceProps) { onConnect={onConnect} className="reactflow-workspace" fitView + minZoom={0.2} edgesUpdatable={!props.readonly} edgesFocusable={!props.readonly} nodesDraggable={!props.readonly} diff --git a/public/pages/workflows/new_workflow/utils.ts b/public/pages/workflows/new_workflow/utils.ts index 4e8ad4e1..9c540589 100644 --- a/public/pages/workflows/new_workflow/utils.ts +++ b/public/pages/workflows/new_workflow/utils.ts @@ -121,181 +121,180 @@ function fetchSemanticSearchMetadata(): UIState { }, }, ] as IModelProcessorConfig; - baseState.workspace_flow = fetchSemanticSearchWorkspaceFlow(); return baseState; } -function fetchSemanticSearchWorkspaceFlow(): WorkspaceFlowState { - const ingestId0 = generateId(COMPONENT_CLASS.DOCUMENT); - 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'); +// function fetchSemanticSearchWorkspaceFlow(): WorkspaceFlowState { +// const ingestId0 = generateId(COMPONENT_CLASS.DOCUMENT); +// 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 = [ - { - id: ingestGroupId, - position: { x: 400, y: 400 }, - type: NODE_CATEGORY.INGEST_GROUP, - data: { label: COMPONENT_CATEGORY.INGEST }, - style: { - width: 1300, - height: 400, - }, - className: 'reactflow__group-node__ingest', - selectable: true, - draggable: true, - deletable: false, - }, - { - id: ingestId0, - position: { x: 100, y: 70 }, - data: initComponentData(new Document().toObj(), ingestId0), - type: NODE_CATEGORY.CUSTOM, - parentNode: ingestGroupId, - extent: 'parent', - draggable: true, - deletable: false, - }, - { - id: ingestId1, - position: { x: 500, y: 70 }, - data: initComponentData( - new TextEmbeddingTransformer().toObj(), - ingestId1 - ), - type: NODE_CATEGORY.CUSTOM, - parentNode: ingestGroupId, - extent: 'parent', - draggable: true, - deletable: false, - }, - { - id: ingestId2, - position: { x: 900, y: 70 }, - data: initComponentData(new KnnIndexer().toObj(), ingestId2), - type: NODE_CATEGORY.CUSTOM, - parentNode: ingestGroupId, - extent: 'parent', - draggable: true, - 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[]; +// const ingestNodes = [ +// { +// id: ingestGroupId, +// position: { x: 400, y: 400 }, +// type: NODE_CATEGORY.INGEST_GROUP, +// data: { label: COMPONENT_CATEGORY.INGEST }, +// style: { +// width: 1300, +// height: 400, +// }, +// className: 'reactflow__group-node__ingest', +// selectable: true, +// draggable: true, +// deletable: false, +// }, +// { +// id: ingestId0, +// position: { x: 100, y: 70 }, +// data: initComponentData(new Document().toObj(), ingestId0), +// type: NODE_CATEGORY.CUSTOM, +// parentNode: ingestGroupId, +// extent: 'parent', +// draggable: true, +// deletable: false, +// }, +// { +// id: ingestId1, +// position: { x: 500, y: 70 }, +// data: initComponentData( +// new TextEmbeddingTransformer().toObj(), +// ingestId1 +// ), +// type: NODE_CATEGORY.CUSTOM, +// parentNode: ingestGroupId, +// extent: 'parent', +// draggable: true, +// deletable: false, +// }, +// { +// id: ingestId2, +// position: { x: 900, y: 70 }, +// data: initComponentData(new KnnIndexer().toObj(), ingestId2), +// type: NODE_CATEGORY.CUSTOM, +// parentNode: ingestGroupId, +// extent: 'parent', +// draggable: true, +// 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, ...searchNodes], - edges: [ - { - id: edgeId0, - key: edgeId0, - source: ingestId0, - target: ingestId1, - markerEnd: { - type: MarkerType.ArrowClosed, - width: 20, - height: 20, - }, - zIndex: 2, - deletable: false, - }, - { - id: edgeId1, - key: edgeId1, - source: ingestId1, - target: ingestId2, - markerEnd: { - type: MarkerType.ArrowClosed, - width: 20, - height: 20, - }, - zIndex: 2, - deletable: false, - }, - { - id: edgeId2, - key: edgeId2, - source: searchId0, - target: searchId1, - markerEnd: { - type: MarkerType.ArrowClosed, - width: 20, - height: 20, - }, - zIndex: 2, - deletable: false, - }, - { - id: edgeId3, - key: edgeId3, - source: searchId1, - target: searchId2, - markerEnd: { - type: MarkerType.ArrowClosed, - width: 20, - height: 20, - }, - zIndex: 2, - deletable: false, - }, - ] as ReactFlowEdge[], - }; -} +// return { +// nodes: [...ingestNodes, ...searchNodes], +// edges: [ +// { +// id: edgeId0, +// key: edgeId0, +// source: ingestId0, +// target: ingestId1, +// markerEnd: { +// type: MarkerType.ArrowClosed, +// width: 20, +// height: 20, +// }, +// zIndex: 2, +// deletable: false, +// }, +// { +// id: edgeId1, +// key: edgeId1, +// source: ingestId1, +// target: ingestId2, +// markerEnd: { +// type: MarkerType.ArrowClosed, +// width: 20, +// height: 20, +// }, +// zIndex: 2, +// deletable: false, +// }, +// { +// id: edgeId2, +// key: edgeId2, +// source: searchId0, +// target: searchId1, +// markerEnd: { +// type: MarkerType.ArrowClosed, +// width: 20, +// height: 20, +// }, +// zIndex: 2, +// deletable: false, +// }, +// { +// id: edgeId3, +// key: edgeId3, +// source: searchId1, +// target: searchId2, +// markerEnd: { +// type: MarkerType.ArrowClosed, +// width: 20, +// height: 20, +// }, +// zIndex: 2, +// deletable: false, +// }, +// ] as ReactFlowEdge[], +// }; +// } // function fetchNeuralSparseSearchWorkspaceFlow(): WorkspaceFlowState { // const ingestId0 = generateId(COMPONENT_CLASS.DOCUMENT); diff --git a/public/utils/utils.ts b/public/utils/utils.ts index ed34a152..e1dba207 100644 --- a/public/utils/utils.ts +++ b/public/utils/utils.ts @@ -27,7 +27,27 @@ import { IConfigField, IndexConfig, IProcessorConfig, + WorkspaceFlowState, + ReactFlowEdge, + ReactFlowComponent, + COMPONENT_CLASS, + COMPONENT_CATEGORY, + NODE_CATEGORY, + IConfig, + IModelProcessorConfig, + PROCESSOR_TYPE, + MODEL_TYPE, } from '../../common'; +import { + Document, + KnnIndexer, + MLTransformer, + NeuralQuery, + Results, + SparseEncoderTransformer, + TextEmbeddingTransformer, +} from '../component_types'; +import { MarkerType } from 'reactflow'; // Append 16 random characters export function generateId(prefix: string): string { @@ -306,3 +326,414 @@ export function getStateOptions(): EuiFilterSelectItem[] { } as EuiFilterSelectItem, ]; } + +/* + **************** ReactFlow workspace utils ********************** + */ + +export function uiConfigToWorkspaceFlow( + config: WorkflowConfig +): WorkspaceFlowState { + const nodes = [] as ReactFlowComponent[]; + const edges = [] as ReactFlowEdge[]; + + const ingestWorkspaceFlow = ingestConfigToWorkspaceFlow(config.ingest); + nodes.push(...ingestWorkspaceFlow.nodes); + edges.push(...ingestWorkspaceFlow.edges); + + const searchWorkspaceFlow = searchConfigToWorkspaceFlow(config.search); + nodes.push(...searchWorkspaceFlow.nodes); + edges.push(...searchWorkspaceFlow.edges); + + return { + nodes: nodes.map((node) => addDefaults(node)), + edges, + }; +} + +function ingestConfigToWorkspaceFlow( + ingestConfig: IngestConfig +): WorkspaceFlowState { + const nodes = [] as ReactFlowComponent[]; + const edges = [] as ReactFlowEdge[]; + + // Parent ingest node + const parentNode = { + id: generateId(COMPONENT_CATEGORY.INGEST), + position: { x: 400, y: 400 }, + type: NODE_CATEGORY.INGEST_GROUP, + data: { label: COMPONENT_CATEGORY.INGEST }, + style: { + width: 1300, + height: 400, + }, + className: 'reactflow__group-node__ingest', + } as ReactFlowComponent; + + nodes.push(parentNode); + + // By default, always include a document node and an index node. + const docNodeId = generateId(COMPONENT_CLASS.DOCUMENT); + const docNode = { + id: docNodeId, + position: { x: 100, y: 70 }, + data: initComponentData(new Document().toObj(), docNodeId), + type: NODE_CATEGORY.CUSTOM, + parentNode: parentNode.id, + extent: 'parent', + } as ReactFlowComponent; + const indexNodeId = generateId(COMPONENT_CLASS.KNN_INDEXER); + const indexNode = { + id: indexNodeId, + position: { x: 900, y: 70 }, + data: initComponentData(new KnnIndexer().toObj(), indexNodeId), + type: NODE_CATEGORY.CUSTOM, + parentNode: parentNode.id, + extent: 'parent', + } as ReactFlowComponent; + nodes.push(docNode, indexNode); + + // Get nodes/edges from the sub-configurations + const enrichWorkspaceFlow = enrichConfigToWorkspaceFlow( + ingestConfig.enrich, + parentNode.id + ); + + nodes.push(...enrichWorkspaceFlow.nodes); + edges.push(...enrichWorkspaceFlow.edges); + + // Link up the set of localized nodes/edges per sub-workflow + edges.push(...getIngestEdges(docNode, enrichWorkspaceFlow, indexNode)); + + return { + nodes, + edges, + }; +} + +function enrichConfigToWorkspaceFlow( + enrichConfig: EnrichConfig, + parentNodeId: string +): WorkspaceFlowState { + const nodes = [] as ReactFlowComponent[]; + const edges = [] as ReactFlowEdge[]; + + // TODO: few assumptions are made here, such as there will always be + // a single model-related processor. In the future make this more flexible and generic. + const modelProcessorConfig = enrichConfig.processors.find( + (processorConfig) => processorConfig.type === PROCESSOR_TYPE.MODEL + ) as IModelProcessorConfig; + + let transformer = {} as MLTransformer; + let transformerNodeId = ''; + switch (modelProcessorConfig.modelType) { + case MODEL_TYPE.TEXT_EMBEDDING: { + transformer = new TextEmbeddingTransformer(); + transformerNodeId = generateId( + COMPONENT_CLASS.TEXT_EMBEDDING_TRANSFORMER + ); + break; + } + case MODEL_TYPE.SPARSE_ENCODER: { + transformer = new SparseEncoderTransformer(); + transformerNodeId = generateId( + COMPONENT_CLASS.SPARSE_ENCODER_TRANSFORMER + ); + break; + } + } + + nodes.push({ + id: transformerNodeId, + position: { x: 500, y: 70 }, + data: initComponentData(transformer, transformerNodeId), + type: NODE_CATEGORY.CUSTOM, + parentNode: parentNodeId, + extent: 'parent', + }); + return { + nodes, + edges, + }; +} + +// Given the set of localized flows per sub-configuration, generate the global ingest-level edges. +// This takes the assumption the flow is linear, and all sub-configuration flows are fully connected. +function getIngestEdges( + docNode: ReactFlowComponent, + enrichFlow: WorkspaceFlowState, + indexNode: ReactFlowComponent +): ReactFlowEdge[] { + const startAndEndNodesEnrich = getStartAndEndNodes(enrichFlow); + + // Users may omit search request processors altogether. Need to handle cases separately. + if (startAndEndNodesEnrich !== undefined) { + const sourceToEnrichEdgeId = generateId('edge'); + const enrichToIndexEdgeId = generateId('edge'); + + return [ + generateReactFlowEdge( + sourceToEnrichEdgeId, + docNode.id, + startAndEndNodesEnrich.startNode.id + ), + generateReactFlowEdge( + enrichToIndexEdgeId, + startAndEndNodesEnrich.endNode.id, + indexNode.id + ), + ] as ReactFlowEdge[]; + } else { + const sourceToIndexEdgeId = generateId('edge'); + return [ + generateReactFlowEdge(sourceToIndexEdgeId, docNode.id, indexNode.id), + ] as ReactFlowEdge[]; + } +} + +function searchConfigToWorkspaceFlow( + searchConfig: SearchConfig +): WorkspaceFlowState { + const nodes = [] as ReactFlowComponent[]; + const edges = [] as ReactFlowEdge[]; + + // Parent search node + const parentNode = { + id: generateId(COMPONENT_CATEGORY.SEARCH), + 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', + } as ReactFlowComponent; + + nodes.push(parentNode); + + // By default, always include a query node, an index node, and a results node. + const queryNodeId = generateId(COMPONENT_CLASS.NEURAL_QUERY); + const queryNode = { + id: queryNodeId, + position: { x: 100, y: 70 }, + data: initComponentData(new NeuralQuery().toObj(), queryNodeId), + type: NODE_CATEGORY.CUSTOM, + parentNode: parentNode.id, + extent: 'parent', + } as ReactFlowComponent; + const indexNodeId = generateId(COMPONENT_CLASS.KNN_INDEXER); + const indexNode = { + id: indexNodeId, + position: { x: 500, y: 70 }, + data: initComponentData(new KnnIndexer().toObj(), indexNodeId), + type: NODE_CATEGORY.CUSTOM, + parentNode: parentNode.id, + extent: 'parent', + } as ReactFlowComponent; + const resultsNodeId = generateId(COMPONENT_CLASS.RESULTS); + const resultsNode = { + id: resultsNodeId, + position: { x: 900, y: 70 }, + data: initComponentData(new Results().toObj(), resultsNodeId), + type: NODE_CATEGORY.CUSTOM, + parentNode: parentNode.id, + extent: 'parent', + } as ReactFlowComponent; + nodes.push(queryNode, indexNode, resultsNode); + + // Get nodes/edges from the sub-configurations + const enrichRequestWorkspaceFlow = enrichRequestConfigToWorkspaceFlow( + searchConfig.enrichRequest, + parentNode.id + ); + const enrichResponseWorkspaceFlow = enrichResponseConfigToWorkspaceFlow( + searchConfig.enrichResponse, + parentNode.id + ); + + nodes.push( + ...enrichRequestWorkspaceFlow.nodes, + ...enrichResponseWorkspaceFlow.nodes + ); + edges.push( + ...enrichRequestWorkspaceFlow.edges, + ...enrichResponseWorkspaceFlow.edges + ); + + // Link up the set of localized nodes/edges per sub-workflow + edges.push( + ...getSearchEdges( + queryNode, + enrichRequestWorkspaceFlow, + indexNode, + enrichResponseWorkspaceFlow, + resultsNode + ) + ); + + return { + nodes, + edges, + }; +} + +// TODO: implement this +function enrichRequestConfigToWorkspaceFlow( + enrichConfig: IConfig, + parentNodeId: string +): WorkspaceFlowState { + const nodes = [] as ReactFlowComponent[]; + const edges = [] as ReactFlowEdge[]; + + return { + nodes, + edges, + }; +} + +// TODO: implement this +function enrichResponseConfigToWorkspaceFlow( + enrichResponseConfig: IConfig, + parentNodeId: string +): WorkspaceFlowState { + const nodes = [] as ReactFlowComponent[]; + const edges = [] as ReactFlowEdge[]; + + return { + nodes, + edges, + }; +} + +// Given the set of localized flows per sub-configuration, generate the global search-level edges. +// This takes the assumption the flow is linear, and all sub-configuration flows are fully connected. +function getSearchEdges( + queryNode: ReactFlowComponent, + enrichRequestFlow: WorkspaceFlowState, + indexNode: ReactFlowComponent, + enrichResponseFlow: WorkspaceFlowState, + resultsNode: ReactFlowComponent +): ReactFlowEdge[] { + const startAndEndNodesEnrichRequest = getStartAndEndNodes(enrichRequestFlow); + const startAndEndNodesEnrichResponse = getStartAndEndNodes( + enrichResponseFlow + ); + const edges = [] as ReactFlowEdge[]; + + // Users may omit search request processors altogether. Need to handle cases separately. + if (startAndEndNodesEnrichRequest !== undefined) { + const requestToEnrichRequestEdgeId = generateId('edge'); + const enrichRequestToIndexEdgeId = generateId('edge'); + edges.push( + ...([ + generateReactFlowEdge( + requestToEnrichRequestEdgeId, + queryNode.id, + startAndEndNodesEnrichRequest.startNode.id + ), + + generateReactFlowEdge( + enrichRequestToIndexEdgeId, + startAndEndNodesEnrichRequest.endNode.id, + indexNode.id + ), + ] as ReactFlowEdge[]) + ); + } else { + const requestToIndexEdgeId = generateId('edge'); + edges.push( + generateReactFlowEdge(requestToIndexEdgeId, queryNode.id, indexNode.id) + ); + } + + // Users may omit search response processors altogether. Need to handle cases separately. + if (startAndEndNodesEnrichResponse !== undefined) { + const indexToEnrichResponseEdgeId = generateId('edge'); + const enrichResponseToResultsEdgeId = generateId('edge'); + + edges.push( + ...([ + generateReactFlowEdge( + indexToEnrichResponseEdgeId, + indexNode.id, + startAndEndNodesEnrichResponse.startNode.id + ), + generateReactFlowEdge( + enrichResponseToResultsEdgeId, + startAndEndNodesEnrichResponse.endNode.id, + resultsNode.id + ), + ] as ReactFlowEdge[]) + ); + } else { + const indexToResultsEdgeId = generateId('edge'); + edges.push( + generateReactFlowEdge(indexToResultsEdgeId, indexNode.id, resultsNode.id) + ); + } + + return edges; +} + +// Get start and end nodes in a flow. This assumes the flow is linear and fully connected, +// such that there will always be a single start and single end node. +function getStartAndEndNodes( + workspaceFlow: WorkspaceFlowState +): + | { + startNode: ReactFlowComponent; + endNode: ReactFlowComponent; + } + | undefined { + if (workspaceFlow.nodes.length === 0) { + return undefined; + } + if (workspaceFlow.nodes.length === 1) { + return { + startNode: workspaceFlow.nodes[0], + endNode: workspaceFlow.nodes[0], + }; + } + + const nodeIdsWithTarget = workspaceFlow.edges.map((edge) => edge.target); + const nodeIdsWithSource = workspaceFlow.edges.map((edge) => edge.source); + + return { + startNode: workspaceFlow.nodes.filter( + (node) => !nodeIdsWithTarget.includes(node.id) + )[0], + endNode: workspaceFlow.nodes.filter( + (node) => !nodeIdsWithSource.includes(node.id) + )[0], + }; +} + +function addDefaults(component: ReactFlowComponent): ReactFlowComponent { + return { + ...component, + draggable: false, + selectable: false, + deletable: false, + }; +} + +function generateReactFlowEdge( + id: string, + source: string, + target: string +): ReactFlowEdge { + return { + id, + key: id, + source, + target, + markerEnd: { + type: MarkerType.ArrowClosed, + width: 20, + height: 20, + }, + zIndex: 2, + deletable: false, + } as ReactFlowEdge; +}