dify/web/app/components/workflow/utils.ts

364 lines
10 KiB
TypeScript

import {
Position,
getConnectedEdges,
getOutgoers,
} from 'reactflow'
import dagre from 'dagre'
import { v4 as uuid4 } from 'uuid'
import {
cloneDeep,
uniqBy,
} from 'lodash-es'
import type {
Edge,
InputVar,
Node,
ToolWithProvider,
} from './types'
import { BlockEnum } from './types'
import {
NODE_WIDTH_X_OFFSET,
START_INITIAL_POSITION,
} from './constants'
import type { QuestionClassifierNodeType } from './nodes/question-classifier/types'
import type { ToolNodeType } from './nodes/tool/types'
import { CollectionType } from '@/app/components/tools/types'
import { toolParametersToFormSchemas } from '@/app/components/tools/utils/to-form-schema'
const WHITE = 'WHITE'
const GRAY = 'GRAY'
const BLACK = 'BLACK'
const isCyclicUtil = (nodeId: string, color: Record<string, string>, adjaList: Record<string, string[]>, stack: string[]) => {
color[nodeId] = GRAY
stack.push(nodeId)
for (let i = 0; i < adjaList[nodeId].length; ++i) {
const childId = adjaList[nodeId][i]
if (color[childId] === GRAY) {
stack.push(childId)
return true
}
if (color[childId] === WHITE && isCyclicUtil(childId, color, adjaList, stack))
return true
}
color[nodeId] = BLACK
if (stack.length > 0 && stack[stack.length - 1] === nodeId)
stack.pop()
return false
}
const getCycleEdges = (nodes: Node[], edges: Edge[]) => {
const adjaList: Record<string, string[]> = {}
const color: Record<string, string> = {}
const stack: string[] = []
for (const node of nodes) {
color[node.id] = WHITE
adjaList[node.id] = []
}
for (const edge of edges)
adjaList[edge.source]?.push(edge.target)
for (let i = 0; i < nodes.length; i++) {
if (color[nodes[i].id] === WHITE)
isCyclicUtil(nodes[i].id, color, adjaList, stack)
}
const cycleEdges = []
if (stack.length > 0) {
const cycleNodes = new Set(stack)
for (const edge of edges) {
if (cycleNodes.has(edge.source) && cycleNodes.has(edge.target))
cycleEdges.push(edge)
}
}
return cycleEdges
}
export const initialNodes = (originNodes: Node[], originEdges: Edge[]) => {
const nodes = cloneDeep(originNodes)
const edges = cloneDeep(originEdges)
const firstNode = nodes[0]
if (!firstNode?.position) {
nodes.forEach((node, index) => {
node.position = {
x: START_INITIAL_POSITION.x + index * NODE_WIDTH_X_OFFSET,
y: START_INITIAL_POSITION.y,
}
})
}
return nodes.map((node) => {
node.type = 'custom'
const connectedEdges = getConnectedEdges([node], edges)
node.data._connectedSourceHandleIds = connectedEdges.filter(edge => edge.source === node.id).map(edge => edge.sourceHandle || 'source')
node.data._connectedTargetHandleIds = connectedEdges.filter(edge => edge.target === node.id).map(edge => edge.targetHandle || 'target')
if (node.data.type === BlockEnum.IfElse) {
node.data._targetBranches = [
{
id: 'true',
name: 'IS TRUE',
},
{
id: 'false',
name: 'IS FALSE',
},
]
}
if (node.data.type === BlockEnum.QuestionClassifier) {
node.data._targetBranches = (node.data as QuestionClassifierNodeType).classes.map((topic) => {
return topic
})
}
return node
})
}
export const initialEdges = (originEdges: Edge[], originNodes: Node[]) => {
const nodes = cloneDeep(originNodes)
const edges = cloneDeep(originEdges)
let selectedNode: Node | null = null
const nodesMap = nodes.reduce((acc, node) => {
acc[node.id] = node
if (node.data?.selected)
selectedNode = node
return acc
}, {} as Record<string, Node>)
const cycleEdges = getCycleEdges(nodes, edges)
return edges.filter((edge) => {
return !cycleEdges.find(cycEdge => cycEdge.source === edge.source && cycEdge.target === edge.target)
}).map((edge) => {
edge.type = 'custom'
if (!edge.sourceHandle)
edge.sourceHandle = 'source'
if (!edge.targetHandle)
edge.targetHandle = 'target'
if (!edge.data?.sourceType && edge.source) {
edge.data = {
...edge.data,
sourceType: nodesMap[edge.source].data.type!,
} as any
}
if (!edge.data?.targetType && edge.target) {
edge.data = {
...edge.data,
targetType: nodesMap[edge.target].data.type!,
} as any
}
if (selectedNode) {
edge.data = {
...edge.data,
_connectedNodeIsSelected: edge.source === selectedNode.id || edge.target === selectedNode.id,
} as any
}
return edge
})
}
const dagreGraph = new dagre.graphlib.Graph()
dagreGraph.setDefaultEdgeLabel(() => ({}))
export const getLayoutByDagre = (originNodes: Node[], originEdges: Edge[]) => {
const nodes = cloneDeep(originNodes)
const edges = cloneDeep(originEdges)
dagreGraph.setGraph({
rankdir: 'LR',
align: 'UL',
nodesep: 40,
ranksep: 60,
})
nodes.forEach((node) => {
dagreGraph.setNode(node.id, { width: node.width, height: node.height })
})
edges.forEach((edge) => {
dagreGraph.setEdge(edge.source, edge.target)
})
dagre.layout(dagreGraph)
return dagreGraph
}
export const canRunBySingle = (nodeType: BlockEnum) => {
return nodeType === BlockEnum.LLM
|| nodeType === BlockEnum.KnowledgeRetrieval
|| nodeType === BlockEnum.Code
|| nodeType === BlockEnum.TemplateTransform
|| nodeType === BlockEnum.QuestionClassifier
|| nodeType === BlockEnum.HttpRequest
|| nodeType === BlockEnum.Tool
}
type ConnectedSourceOrTargetNodesChange = {
type: string
edge: Edge
}[]
export const getNodesConnectedSourceOrTargetHandleIdsMap = (changes: ConnectedSourceOrTargetNodesChange, nodes: Node[]) => {
const nodesConnectedSourceOrTargetHandleIdsMap = {} as Record<string, any>
changes.forEach((change) => {
const {
edge,
type,
} = change
const sourceNode = nodes.find(node => node.id === edge.source)!
if (sourceNode) {
nodesConnectedSourceOrTargetHandleIdsMap[sourceNode.id] = nodesConnectedSourceOrTargetHandleIdsMap[sourceNode.id] || {
_connectedSourceHandleIds: [...(sourceNode?.data._connectedSourceHandleIds || [])],
_connectedTargetHandleIds: [...(sourceNode?.data._connectedTargetHandleIds || [])],
}
}
const targetNode = nodes.find(node => node.id === edge.target)!
if (targetNode) {
nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id] = nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id] || {
_connectedSourceHandleIds: [...(targetNode?.data._connectedSourceHandleIds || [])],
_connectedTargetHandleIds: [...(targetNode?.data._connectedTargetHandleIds || [])],
}
}
if (sourceNode) {
if (type === 'remove')
nodesConnectedSourceOrTargetHandleIdsMap[sourceNode.id]._connectedSourceHandleIds = nodesConnectedSourceOrTargetHandleIdsMap[sourceNode.id]._connectedSourceHandleIds.filter((handleId: string) => handleId !== edge.sourceHandle)
if (type === 'add')
nodesConnectedSourceOrTargetHandleIdsMap[sourceNode.id]._connectedSourceHandleIds.push(edge.sourceHandle || 'source')
}
if (targetNode) {
if (type === 'remove')
nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id]._connectedTargetHandleIds = nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id]._connectedTargetHandleIds.filter((handleId: string) => handleId !== edge.targetHandle)
if (type === 'add')
nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id]._connectedTargetHandleIds.push(edge.targetHandle || 'target')
}
})
return nodesConnectedSourceOrTargetHandleIdsMap
}
export const generateNewNode = ({ data, position, id }: Pick<Node, 'data' | 'position'> & { id?: string }) => {
return {
id: id || `${Date.now()}`,
type: 'custom',
data,
position,
targetPosition: Position.Left,
sourcePosition: Position.Right,
} as Node
}
export const getValidTreeNodes = (nodes: Node[], edges: Edge[]) => {
const startNode = nodes.find(node => node.data.type === BlockEnum.Start)
if (!startNode) {
return {
validNodes: [],
maxDepth: 0,
}
}
const list: Node[] = [startNode]
let maxDepth = 1
const traverse = (root: Node, depth: number) => {
if (depth > maxDepth)
maxDepth = depth
const outgoers = getOutgoers(root, nodes, edges)
if (outgoers.length) {
outgoers.forEach((outgoer) => {
list.push(outgoer)
traverse(outgoer, depth + 1)
})
}
else {
list.push(root)
}
}
traverse(startNode, maxDepth)
return {
validNodes: uniqBy(list, 'id'),
maxDepth,
}
}
export const getToolCheckParams = (
toolData: ToolNodeType,
buildInTools: ToolWithProvider[],
customTools: ToolWithProvider[],
language: string,
) => {
const { provider_id, provider_type, tool_name } = toolData
const isBuiltIn = provider_type === CollectionType.builtIn
const currentTools = isBuiltIn ? buildInTools : customTools
const currCollection = currentTools.find(item => item.id === provider_id)
const currTool = currCollection?.tools.find(tool => tool.name === tool_name)
const formSchemas = currTool ? toolParametersToFormSchemas(currTool.parameters) : []
const toolInputVarSchema = formSchemas.filter((item: any) => item.form === 'llm')
const toolSettingSchema = formSchemas.filter((item: any) => item.form !== 'llm')
return {
toolInputsSchema: (() => {
const formInputs: InputVar[] = []
toolInputVarSchema.forEach((item: any) => {
formInputs.push({
label: item.label[language] || item.label.en_US,
variable: item.variable,
type: item.type,
required: item.required,
})
})
return formInputs
})(),
notAuthed: isBuiltIn && !!currCollection?.allow_delete && !currCollection?.is_team_authorization,
toolSettingSchema,
language,
}
}
export const changeNodesAndEdgesId = (nodes: Node[], edges: Edge[]) => {
const idMap = nodes.reduce((acc, node) => {
acc[node.id] = uuid4()
return acc
}, {} as Record<string, string>)
const newNodes = nodes.map((node) => {
return {
...node,
id: idMap[node.id],
}
})
const newEdges = edges.map((edge) => {
return {
...edge,
source: idMap[edge.source],
target: idMap[edge.target],
}
})
return [newNodes, newEdges] as [Node[], Edge[]]
}