dify/web/app/components/workflow/hooks/use-nodes-available-var-list.ts

import type { Node, NodeOutPutVar, ValueSelector, Var } from '@/app/components/workflow/types'
import { useCallback } from 'react'
import { useTranslation } from 'react-i18next'
import { useSnippetDraftStore } from '@/app/components/snippets/draft-store'
import { useIsChatMode, useWorkflow, useWorkflowVariables } from '@/app/components/workflow/hooks'
import { useHooksStore } from '@/app/components/workflow/hooks-store/store'
import {
  appendSnippetInputFieldVars,
  filterSnippetSystemVars,
  isSnippetCanvas,
} from '@/app/components/workflow/nodes/_base/hooks/snippet-input-field-vars'
import { BlockEnum } from '@/app/components/workflow/types'
import { FlowType } from '@/types/common'

type Params = {
  onlyLeafNodeVar?: boolean
  hideEnv?: boolean
  hideChatVar?: boolean
  filterVar: (payload: Var, selector: ValueSelector) => boolean
  passedInAvailableNodes?: Node[]
}

const getNodeInfo = (nodeId: string, nodes: Node[]) => {
  const allNodes = nodes
  const node = allNodes.find((n) => n.id === nodeId)
  const isInIteration = !!node?.data.isInIteration
  const isInLoop = !!node?.data.isInLoop
  const parentNodeId = node?.parentId
  const parentNode = allNodes.find((n) => n.id === parentNodeId)
  return {
    node,
    isInIteration,
    isInLoop,
    parentNode,
  }
}

// TODO: loop type?
const useNodesAvailableVarList = (
  nodes: Node[],
  {
    onlyLeafNodeVar,
    filterVar,
    hideEnv = false,
    hideChatVar = false,
    passedInAvailableNodes,
  }: Params = {
    onlyLeafNodeVar: false,
    filterVar: () => true,
  },
) => {
  const { t } = useTranslation()
  const snippetInputFields = useSnippetDraftStore((s) => s.inputFields)
  const { getTreeLeafNodes, getBeforeNodesInSameBranchIncludeParent } = useWorkflow()
  const { getNodeAvailableVars } = useWorkflowVariables()
  const isChatMode = useIsChatMode()
  const isSnippetFlow =
    useHooksStore((s) => s.configsMap?.flowType) === FlowType.snippet || isSnippetCanvas()

  const nodeAvailabilityMap: {
    [key: string]: { availableVars: NodeOutPutVar[]; availableNodes: Node[] }
  } = {}

  nodes.forEach((node) => {
    const nodeId = node.id
    const availableNodes =
      passedInAvailableNodes ||
      (onlyLeafNodeVar ? getTreeLeafNodes(nodeId) : getBeforeNodesInSameBranchIncludeParent(nodeId))
    if (node.data.type === BlockEnum.Loop) availableNodes.push(node)
    const snippetInputFieldAvailability = appendSnippetInputFieldVars({
      availableNodes,
      fields: snippetInputFields,
      title: t(($) => $.panelTitle, { ns: 'snippet' }),
    })

    const { parentNode: iterationNode } = getNodeInfo(nodeId, nodes)

    const availableVars = filterSnippetSystemVars(
      [
        ...snippetInputFieldAvailability.availableVars,
        ...getNodeAvailableVars({
          parentNode: iterationNode,
          beforeNodes: availableNodes,
          isChatMode,
          filterVar,
          hideEnv,
          hideChatVar,
        }),
      ],
      isSnippetFlow,
    )
    const result = {
      node,
      availableVars,
      availableNodes: snippetInputFieldAvailability.availableNodes,
    }
    nodeAvailabilityMap[nodeId] = result
  })
  return nodeAvailabilityMap
}

export const useGetNodesAvailableVarList = () => {
  const { t } = useTranslation()
  const snippetInputFields = useSnippetDraftStore((s) => s.inputFields)
  const { getTreeLeafNodes, getBeforeNodesInSameBranchIncludeParent } = useWorkflow()
  const { getNodeAvailableVars } = useWorkflowVariables()
  const isChatMode = useIsChatMode()
  const isSnippetFlow =
    useHooksStore((s) => s.configsMap?.flowType) === FlowType.snippet || isSnippetCanvas()
  const getNodesAvailableVarList = useCallback(
    (
      nodes: Node[],
      { onlyLeafNodeVar, filterVar, hideEnv, hideChatVar, passedInAvailableNodes }: Params = {
        onlyLeafNodeVar: false,
        filterVar: () => true,
      },
    ) => {
      const nodeAvailabilityMap: {
        [key: string]: { availableVars: NodeOutPutVar[]; availableNodes: Node[] }
      } = {}

      nodes.forEach((node) => {
        const nodeId = node.id
        const availableNodes =
          passedInAvailableNodes ||
          (onlyLeafNodeVar
            ? getTreeLeafNodes(nodeId)
            : getBeforeNodesInSameBranchIncludeParent(nodeId))
        if (node.data.type === BlockEnum.Loop) availableNodes.push(node)
        const snippetInputFieldAvailability = appendSnippetInputFieldVars({
          availableNodes,
          fields: snippetInputFields,
          title: t(($) => $.panelTitle, { ns: 'snippet' }),
        })

        const { parentNode: iterationNode } = getNodeInfo(nodeId, nodes)

        const availableVars = filterSnippetSystemVars(
          [
            ...snippetInputFieldAvailability.availableVars,
            ...getNodeAvailableVars({
              parentNode: iterationNode,
              beforeNodes: availableNodes,
              isChatMode,
              filterVar,
              hideEnv,
              hideChatVar,
            }),
          ],
          isSnippetFlow,
        )
        const result = {
          node,
          availableVars,
          availableNodes: snippetInputFieldAvailability.availableNodes,
        }
        nodeAvailabilityMap[nodeId] = result
      })
      return nodeAvailabilityMap
    },
    [
      getTreeLeafNodes,
      getBeforeNodesInSameBranchIncludeParent,
      getNodeAvailableVars,
      isChatMode,
      isSnippetFlow,
      snippetInputFields,
      t,
    ],
  )
  return {
    getNodesAvailableVarList,
  }
}

export default useNodesAvailableVarList