This commit is contained in:
StyleZhang 2024-03-19 20:34:10 +08:00
parent 089072432e
commit 9f024835aa
2 changed files with 40 additions and 6 deletions

View File

@ -102,7 +102,17 @@ export const useWorkflowRun = () => {
} }
}, [store, handleLoadBackupDraft, handleBackupDraft, workflowStore]) }, [store, handleLoadBackupDraft, handleBackupDraft, workflowStore])
const handleRun = useCallback((params: any, callback?: IOtherOptions) => { const handleRun = useCallback((
params: any,
callback?: IOtherOptions,
) => {
const {
onWorkflowStarted,
onWorkflowFinished,
onNodeStarted,
onNodeFinished,
...restCallback
} = callback || {}
const { const {
getNodes, getNodes,
setNodes, setNodes,
@ -130,7 +140,8 @@ export const useWorkflowRun = () => {
body: params, body: params,
}, },
{ {
onWorkflowStarted: ({ task_id, workflow_run_id, data }) => { onWorkflowStarted: (params) => {
const { task_id, workflow_run_id, data } = params
workflowStore.setState({ runningStatus: WorkflowRunningStatus.Running }) workflowStore.setState({ runningStatus: WorkflowRunningStatus.Running })
workflowStore.setState({ taskId: task_id }) workflowStore.setState({ taskId: task_id })
workflowStore.setState({ currentSequenceNumber: data.sequence_number }) workflowStore.setState({ currentSequenceNumber: data.sequence_number })
@ -141,11 +152,19 @@ export const useWorkflowRun = () => {
}) })
}) })
setNodes(newNodes) setNodes(newNodes)
if (onWorkflowStarted)
onWorkflowStarted(params)
}, },
onWorkflowFinished: ({ data }) => { onWorkflowFinished: (params) => {
const { data } = params
workflowStore.setState({ runningStatus: data.status as WorkflowRunningStatus }) workflowStore.setState({ runningStatus: data.status as WorkflowRunningStatus })
if (onWorkflowFinished)
onWorkflowFinished(params)
}, },
onNodeStarted: ({ data }) => { onNodeStarted: (params) => {
const { data } = params
const nodes = getNodes() const nodes = getNodes()
const { const {
setViewport, setViewport,
@ -171,16 +190,23 @@ export const useWorkflowRun = () => {
edge.data._runned = true edge.data._runned = true
}) })
setEdges(newEdges) setEdges(newEdges)
if (onNodeStarted)
onNodeStarted(params)
}, },
onNodeFinished: ({ data }) => { onNodeFinished: (params) => {
const { data } = params
const newNodes = produce(getNodes(), (draft) => { const newNodes = produce(getNodes(), (draft) => {
const currentNode = draft.find(node => node.id === data.node_id)! const currentNode = draft.find(node => node.id === data.node_id)!
currentNode.data._runningStatus = data.status currentNode.data._runningStatus = data.status
}) })
setNodes(newNodes) setNodes(newNodes)
if (onNodeFinished)
onNodeFinished(params)
}, },
...callback, ...restCallback,
}, },
) )
}, [store, reactflow, workflowStore]) }, [store, reactflow, workflowStore])

View File

@ -261,6 +261,14 @@ export const useChat = (
}) })
handleUpdateChatList(newChatList) handleUpdateChatList(newChatList)
}, },
onWorkflowStarted: () => {
},
onWorkflowFinished: () => {
},
onNodeStarted: () => {
},
onNodeFinished: () => {
},
}, },
) )
}, [handleRun, handleResponding, handleUpdateChatList, notify, t, updateCurrentQA, config.suggested_questions_after_answer?.enabled]) }, [handleRun, handleResponding, handleUpdateChatList, notify, t, updateCurrentQA, config.suggested_questions_after_answer?.enabled])