dify/web/app/components/rag-pipeline/hooks/use-pipeline.tsx

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<DataSourceNodeType>[]
    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,
  }
}