From a866cbc6d7e1e28ec2f3b8b2b68efa34b4427262 Mon Sep 17 00:00:00 2001 From: twwu Date: Wed, 25 Jun 2025 10:11:26 +0800 Subject: [PATCH] feat: implement usePipeline hook for managing pipeline variables and refactor input field handling --- .../input-field/field-list/hooks.ts | 14 ++- .../rag-pipeline/hooks/use-pipeline.tsx | 115 ++++++++++++++++++ .../nodes/_base/components/variable/utils.ts | 41 ++++++- 3 files changed, 159 insertions(+), 11 deletions(-) create mode 100644 web/app/components/rag-pipeline/hooks/use-pipeline.tsx diff --git a/web/app/components/rag-pipeline/components/input-field/field-list/hooks.ts b/web/app/components/rag-pipeline/components/input-field/field-list/hooks.ts index aeced64ca7..a2aff2ce3a 100644 --- a/web/app/components/rag-pipeline/components/input-field/field-list/hooks.ts +++ b/web/app/components/rag-pipeline/components/input-field/field-list/hooks.ts @@ -8,9 +8,11 @@ import type { InputVar } from '@/models/pipeline' import type { SortableItem } from './types' import type { MoreInfo, ValueSelector } from '@/app/components/workflow/types' import { ChangeType } from '@/app/components/workflow/types' -import { useWorkflow } from '@/app/components/workflow/hooks' import { useBoolean } from 'ahooks' import Toast from '@/app/components/base/toast' +import { usePipeline } from '../../../hooks/use-pipeline' + +const VARIABLE_PREFIX = 'rag' export const useFieldList = ( initialInputFields: InputVar[], @@ -22,7 +24,7 @@ export const useFieldList = ( const [removedVar, setRemovedVar] = useState([]) const [removedIndex, setRemoveIndex] = useState(0) - const { handleOutVarRenameChange, isVarUsedInNodes, removeUsedVarInNodes } = useWorkflow() + const { handleInputVarRename, isVarUsedInNodes, removeUsedVarInNodes } = usePipeline() const [isShowRemoveVarConfirm, { setTrue: showRemoveVarConfirm, @@ -61,9 +63,9 @@ export const useFieldList = ( const handleRemoveField = useCallback((index: number) => { const itemToRemove = inputFieldsRef.current[index] // Check if the variable is used in other nodes - if (isVarUsedInNodes([nodeId, itemToRemove.variable || ''])) { + if (isVarUsedInNodes([VARIABLE_PREFIX, nodeId, itemToRemove.variable || ''])) { showRemoveVarConfirm() - setRemovedVar([nodeId, itemToRemove.variable || '']) + setRemovedVar([VARIABLE_PREFIX, nodeId, itemToRemove.variable || '']) setRemoveIndex(index as number) return } @@ -99,9 +101,9 @@ export const useFieldList = ( handleInputFieldsChange(newInputFields) // Update variable name in nodes if it has changed if (moreInfo?.type === ChangeType.changeVarName) - handleOutVarRenameChange(nodeId, [nodeId, moreInfo.payload?.beforeKey || ''], [nodeId, moreInfo.payload?.afterKey || '']) + handleInputVarRename(nodeId, [VARIABLE_PREFIX, nodeId, moreInfo.payload?.beforeKey || ''], [VARIABLE_PREFIX, nodeId, moreInfo.payload?.afterKey || '']) handleCloseInputFieldEditor() - }, [editingField?.variable, handleCloseInputFieldEditor, handleInputFieldsChange, handleOutVarRenameChange, nodeId]) + }, [editingField?.variable, handleCloseInputFieldEditor, handleInputFieldsChange, handleInputVarRename, nodeId]) return { inputFields, diff --git a/web/app/components/rag-pipeline/hooks/use-pipeline.tsx b/web/app/components/rag-pipeline/hooks/use-pipeline.tsx new file mode 100644 index 0000000000..8ac73b7342 --- /dev/null +++ b/web/app/components/rag-pipeline/hooks/use-pipeline.tsx @@ -0,0 +1,115 @@ +import { useCallback } from 'react' +import { getOutgoers, useStoreApi } from 'reactflow' +import { BlockEnum, type Node, type ValueSelector } from '../../workflow/types' +import { uniqBy } from 'lodash-es' +import { findUsedVarNodes, updateNodeVars } from '../../workflow/nodes/_base/components/variable/utils' +import type { DataSourceNodeType } from '../../workflow/nodes/data-source/types' + +export const usePipeline = () => { + const store = useStoreApi() + + const getAllDatasourceNodes = useCallback(() => { + const { + getNodes, + } = store.getState() + const nodes = getNodes() as Node[] + const datasourceNodes = nodes.filter(node => node.data.type === BlockEnum.DataSource) + + return datasourceNodes + }, [store]) + + const getAllNodesInSameBranch = useCallback((nodeId: string) => { + const { + getNodes, + edges, + } = store.getState() + const nodes = getNodes() + const list: Node[] = [] + + const traverse = (root: Node, callback: (node: Node) => void) => { + if (root) { + const outgoers = getOutgoers(root, nodes, edges) + + if (outgoers.length) { + outgoers.forEach((node) => { + callback(node) + traverse(node, callback) + }) + } + } + } + + if (nodeId === 'shared') { + const allDatasourceNodes = getAllDatasourceNodes() + + if (allDatasourceNodes.length === 0) + return [] + + list.push(...allDatasourceNodes) + + allDatasourceNodes.forEach((node) => { + traverse(node, (childNode) => { + list.push(childNode) + }) + }) + } + else { + const currentNode = nodes.find(node => node.id === nodeId)! + + if (!currentNode) + return [] + + list.push(currentNode) + + traverse(currentNode, (node) => { + list.push(node) + }) + } + + return uniqBy(list, 'id') + }, [getAllDatasourceNodes, store]) + + const isVarUsedInNodes = useCallback((varSelector: ValueSelector) => { + const nodeId = varSelector[1] // Assuming the first element is always 'VARIABLE_PREFIX'(rag) + const afterNodes = getAllNodesInSameBranch(nodeId) + const effectNodes = findUsedVarNodes(varSelector, afterNodes) + return effectNodes.length > 0 + }, [getAllNodesInSameBranch]) + + const handleInputVarRename = useCallback((nodeId: string, oldValeSelector: ValueSelector, newVarSelector: ValueSelector) => { + const { getNodes, setNodes } = store.getState() + const afterNodes = getAllNodesInSameBranch(nodeId) + const effectNodes = findUsedVarNodes(oldValeSelector, afterNodes) + if (effectNodes.length > 0) { + const newNodes = getNodes().map((node) => { + if (effectNodes.find(n => n.id === node.id)) + return updateNodeVars(node, oldValeSelector, newVarSelector) + + return node + }) + setNodes(newNodes) + } + }, [getAllNodesInSameBranch, store]) + + const removeUsedVarInNodes = useCallback((varSelector: ValueSelector) => { + const nodeId = varSelector[1] // Assuming the first element is always 'VARIABLE_PREFIX'(rag) + const { getNodes, setNodes } = store.getState() + const afterNodes = getAllNodesInSameBranch(nodeId) + const effectNodes = findUsedVarNodes(varSelector, afterNodes) + if (effectNodes.length > 0) { + const newNodes = getNodes().map((node) => { + if (effectNodes.find(n => n.id === node.id)) + return updateNodeVars(node, varSelector, []) + + return node + }) + setNodes(newNodes) + } + }, [getAllNodesInSameBranch, store]) + + return { + handleInputVarRename, + isVarUsedInNodes, + removeUsedVarInNodes, + } +} diff --git a/web/app/components/workflow/nodes/_base/components/variable/utils.ts b/web/app/components/workflow/nodes/_base/components/variable/utils.ts index b68a506c18..63f2527aa8 100644 --- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts +++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts @@ -838,9 +838,9 @@ export const getVarType = ({ }) const targetVarNodeId = (() => { - if(isSystem) + if (isSystem) return startNode?.id - if(isInNodeRagVariable) + if (isInNodeRagVariable) return valueSelector[1] return valueSelector[0] })() @@ -857,14 +857,14 @@ export const getVarType = ({ } else { const targetVar = curr.find((v: any) => { - if(isInNodeRagVariable) + if (isInNodeRagVariable) return v.variable === valueSelector.join('.') return v.variable === valueSelector[1] - }) + }) if (!targetVar) return VarType.string - if(isInNodeRagVariable) + if (isInNodeRagVariable) return targetVar.type const isStructuredOutputVar = !!targetVar.children?.schema?.properties @@ -1084,6 +1084,13 @@ export const getNodeUsedVars = (node: Node): ValueSelector[] => { res = [...(mixVars as ValueSelector[]), ...(vars as any)] break } + case BlockEnum.DataSource: { + const payload = data as DataSourceNodeType + const mixVars = matchNotSystemVars(Object.keys(payload.datasource_parameters)?.filter(key => payload.datasource_parameters[key].type === ToolVarType.mixed).map(key => payload.datasource_parameters[key].value) as string[]) + const vars = Object.keys(payload.datasource_parameters).filter(key => payload.datasource_parameters[key].type === ToolVarType.variable).map(key => payload.datasource_parameters[key].value as string) || [] + res = [...(mixVars as ValueSelector[]), ...(vars as any)] + break + } case BlockEnum.VariableAssigner: { res = (data as VariableAssignerNodeType)?.variables @@ -1357,6 +1364,30 @@ export const updateNodeVars = (oldNode: Node, oldVarSelector: ValueSelector, new } break } + case BlockEnum.DataSource: { + const payload = data as DataSourceNodeType + const hasShouldRenameVar = Object.keys(payload.datasource_parameters)?.filter(key => payload.datasource_parameters[key].type !== ToolVarType.constant) + if (hasShouldRenameVar) { + Object.keys(payload.datasource_parameters).forEach((key) => { + const value = payload.datasource_parameters[key] + const { type } = value + if (type === ToolVarType.variable) { + payload.datasource_parameters[key] = { + ...value, + value: newVarSelector, + } + } + + if (type === ToolVarType.mixed) { + payload.datasource_parameters[key] = { + ...value, + value: replaceOldVarInText(payload.datasource_parameters[key].value as string, oldVarSelector, newVarSelector), + } + } + }) + } + break + } case BlockEnum.VariableAssigner: { const payload = data as VariableAssignerNodeType if (payload.variables) {