feat: Parallel Execution of Nodes in Workflows (#8192)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>
Co-authored-by: Yi <yxiaoisme@gmail.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
takatost
2024-09-10 15:23:16 +08:00
committed by GitHub
parent 5da0182800
commit dabfd74622
156 changed files with 11158 additions and 5605 deletions

View File

@@ -1,5 +1,6 @@
import { useCallback } from 'react'
import {
getIncomers,
useReactFlow,
useStoreApi,
} from 'reactflow'
@@ -8,6 +9,7 @@ import { v4 as uuidV4 } from 'uuid'
import { usePathname } from 'next/navigation'
import { useWorkflowStore } from '../store'
import { useNodesSyncDraft } from '../hooks'
import type { Node } from '../types'
import {
NodeRunningStatus,
WorkflowRunningStatus,
@@ -140,9 +142,6 @@ export const useWorkflowRun = () => {
resultText: '',
})
let isInIteration = false
let iterationLength = 0
let ttsUrl = ''
let ttsIsPublic = false
if (params.token) {
@@ -186,7 +185,7 @@ export const useWorkflowRun = () => {
draft.forEach((edge) => {
edge.data = {
...edge.data,
_run: false,
_runned: false,
}
})
})
@@ -249,19 +248,20 @@ export const useWorkflowRun = () => {
setEdges,
transform,
} = store.getState()
if (isInIteration) {
const nodes = getNodes()
const node = nodes.find(node => node.id === data.node_id)
if (node?.parentId) {
setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
const tracing = draft.tracing!
const iterations = tracing[tracing.length - 1]
const currIteration = iterations.details![iterations.details!.length - 1]
currIteration.push({
const iterations = tracing.find(trace => trace.node_id === node?.parentId)
const currIteration = iterations?.details![node.data.iteration_index] || iterations?.details![iterations.details!.length - 1]
currIteration?.push({
...data,
status: NodeRunningStatus.Running,
} as any)
}))
}
else {
const nodes = getNodes()
setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
draft.tracing!.push({
...data,
@@ -288,11 +288,12 @@ export const useWorkflowRun = () => {
draft[currentNodeIndex].data._runningStatus = NodeRunningStatus.Running
})
setNodes(newNodes)
const incomeNodesId = getIncomers({ id: data.node_id } as Node, newNodes, edges).filter(node => node.data._runningStatus === NodeRunningStatus.Succeeded).map(node => node.id)
const newEdges = produce(edges, (draft) => {
const edge = draft.find(edge => edge.target === data.node_id && edge.source === prevNodeId)
if (edge)
edge.data = { ...edge.data, _run: true } as any
draft.forEach((edge) => {
if (edge.target === data.node_id && incomeNodesId.includes(edge.source))
edge.data = { ...edge.data, _runned: true } as any
})
})
setEdges(newEdges)
}
@@ -309,25 +310,46 @@ export const useWorkflowRun = () => {
getNodes,
setNodes,
} = store.getState()
if (isInIteration) {
const nodes = getNodes()
const nodeParentId = nodes.find(node => node.id === data.node_id)!.parentId
if (nodeParentId) {
setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
const tracing = draft.tracing!
const iterations = tracing[tracing.length - 1]
const currIteration = iterations.details![iterations.details!.length - 1]
const nodeInfo = currIteration[currIteration.length - 1]
const iterations = tracing.find(trace => trace.node_id === nodeParentId) // the iteration node
currIteration[currIteration.length - 1] = {
...nodeInfo,
...data,
status: NodeRunningStatus.Succeeded,
} as any
if (iterations && iterations.details) {
const iterationIndex = data.execution_metadata?.iteration_index || 0
if (!iterations.details[iterationIndex])
iterations.details[iterationIndex] = []
const currIteration = iterations.details[iterationIndex]
const nodeIndex = currIteration.findIndex(node =>
node.node_id === data.node_id && (
node.execution_metadata?.parallel_id === data.execution_metadata?.parallel_id || node.parallel_id === data.execution_metadata?.parallel_id),
)
if (data.status === NodeRunningStatus.Succeeded) {
if (nodeIndex !== -1) {
currIteration[nodeIndex] = {
...currIteration[nodeIndex],
...data,
} as any
}
else {
currIteration.push({
...data,
} as any)
}
}
}
}))
}
else {
const nodes = getNodes()
setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
const currentIndex = draft.tracing!.findIndex(trace => trace.node_id === data.node_id)
const currentIndex = draft.tracing!.findIndex((trace) => {
if (!trace.execution_metadata?.parallel_id)
return trace.node_id === data.node_id
return trace.node_id === data.node_id && trace.execution_metadata?.parallel_id === data.execution_metadata?.parallel_id
})
if (currentIndex > -1 && draft.tracing) {
draft.tracing[currentIndex] = {
...(draft.tracing[currentIndex].extras
@@ -337,16 +359,14 @@ export const useWorkflowRun = () => {
} as any
}
}))
const newNodes = produce(nodes, (draft) => {
const currentNode = draft.find(node => node.id === data.node_id)!
currentNode.data._runningStatus = data.status as any
})
setNodes(newNodes)
prevNodeId = data.node_id
}
if (onNodeFinished)
onNodeFinished(params)
},
@@ -371,8 +391,6 @@ export const useWorkflowRun = () => {
details: [],
} as any)
}))
isInIteration = true
iterationLength = data.metadata.iterator_length
const {
setViewport,
@@ -398,7 +416,7 @@ export const useWorkflowRun = () => {
const edge = draft.find(edge => edge.target === data.node_id && edge.source === prevNodeId)
if (edge)
edge.data = { ...edge.data, _run: true } as any
edge.data = { ...edge.data, _runned: true } as any
})
setEdges(newEdges)
@@ -418,13 +436,13 @@ export const useWorkflowRun = () => {
} = store.getState()
setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
const iteration = draft.tracing![draft.tracing!.length - 1]
if (iteration.details!.length >= iterationLength)
return
iteration.details!.push([])
const iteration = draft.tracing!.find(trace => trace.node_id === data.node_id)
if (iteration) {
if (iteration.details!.length >= iteration.metadata.iterator_length!)
return
}
iteration?.details!.push([])
}))
const nodes = getNodes()
const newNodes = produce(nodes, (draft) => {
const currentNode = draft.find(node => node.id === data.node_id)!
@@ -450,13 +468,14 @@ export const useWorkflowRun = () => {
const nodes = getNodes()
setWorkflowRunningData(produce(workflowRunningData!, (draft) => {
const tracing = draft.tracing!
tracing[tracing.length - 1] = {
...tracing[tracing.length - 1],
...data,
status: NodeRunningStatus.Succeeded,
} as any
const currIterationNode = tracing.find(trace => trace.node_id === data.node_id)
if (currIterationNode) {
Object.assign(currIterationNode, {
...data,
status: NodeRunningStatus.Succeeded,
})
}
}))
isInIteration = false
const newNodes = produce(nodes, (draft) => {
const currentNode = draft.find(node => node.id === data.node_id)!
@@ -470,6 +489,12 @@ export const useWorkflowRun = () => {
if (onIterationFinish)
onIterationFinish(params)
},
onParallelBranchStarted: (params) => {
// console.log(params, 'parallel start')
},
onParallelBranchFinished: (params) => {
// console.log(params, 'finished')
},
onTextChunk: (params) => {
const { data: { text } } = params
const {