Commit 3d3bc4c5 authored by StyleZhang's avatar StyleZhang

initial node data

parent 044ed624
import { useState } from 'react'
import {
memo,
useState,
} from 'react'
import { useNodeId } from 'reactflow'
import BlockIcon from '../block-icon'
import { useWorkflowContext } from '../context'
import {
BLOCK_CLASSIFICATIONS,
BLOCK_GROUP_BY_CLASSIFICATION,
......@@ -7,7 +12,13 @@ import {
} from './constants'
const Tabs = () => {
const {
nodes,
handleAddNextNode,
} = useWorkflowContext()
const [activeTab, setActiveTab] = useState(TABS[0].key)
const nodeId = useNodeId()
const currentNode = nodes.find(node => node.id === nodeId)
return (
<div>
......@@ -46,6 +57,10 @@ const Tabs = () => {
<div
key={block.type}
className='flex items-center px-3 h-8 rounded-lg hover:bg-gray-50 cursor-pointer'
onClick={(e) => {
e.stopPropagation()
handleAddNextNode(currentNode!, block.type)
}}
>
<BlockIcon
className='mr-2'
......@@ -63,4 +78,4 @@ const Tabs = () => {
)
}
export default Tabs
export default memo(Tabs)
import { BlockEnum } from './types'
export const NodeInitialData = {
[BlockEnum.Start]: {
type: BlockEnum.Start,
title: '',
desc: '',
variables: [],
},
[BlockEnum.End]: {
type: BlockEnum.End,
title: '',
desc: '',
outputs: {},
},
[BlockEnum.DirectAnswer]: {
type: BlockEnum.DirectAnswer,
title: '',
desc: '',
variables: [],
},
[BlockEnum.LLM]: {
type: BlockEnum.LLM,
title: '',
desc: '',
variables: [],
},
[BlockEnum.KnowledgeRetrieval]: {
type: BlockEnum.KnowledgeRetrieval,
title: '',
desc: '',
query_variable_selector: [],
dataset_ids: [],
retrieval_mode: 'single',
},
[BlockEnum.IfElse]: {
type: BlockEnum.IfElse,
title: '',
desc: '',
logical_operator: 'and',
conditions: [],
},
[BlockEnum.Code]: {
type: BlockEnum.Code,
title: '',
desc: '',
variables: [],
code_language: 'python3',
code: '',
outputs: [],
},
[BlockEnum.TemplateTransform]: {
type: BlockEnum.TemplateTransform,
title: '',
desc: '',
variables: [],
template: '',
},
[BlockEnum.QuestionClassifier]: {
type: BlockEnum.QuestionClassifier,
title: '',
desc: '',
query_variable_selector: [],
topics: [],
},
[BlockEnum.HttpRequest]: {
type: BlockEnum.HttpRequest,
title: '',
desc: '',
variables: [],
},
[BlockEnum.VariableAssigner]: {
type: BlockEnum.VariableAssigner,
title: '',
desc: '',
variables: [],
output_type: '',
},
[BlockEnum.Tool]: {
type: BlockEnum.Tool,
title: '',
desc: '',
},
}
'use client'
import { createContext, useContext } from 'use-context-selector'
import type { Edge } from 'reactflow'
import type { Node } from './types'
import type {
Edge,
ReactFlowInstance,
} from 'reactflow'
import type {
BlockEnum,
Node,
} from './types'
export type WorkflowContextValue = {
reactFlow: ReactFlowInstance
nodes: Node[]
edges: Edge[]
selectedNodeId?: string
handleSelectedNodeIdChange: (nodeId: string) => void
selectedNode?: Node
handleAddNextNode: (prevNode: Node, nextNodeType: BlockEnum) => void
}
export const WorkflowContext = createContext<WorkflowContextValue>({
reactFlow: null as any,
nodes: [],
edges: [],
handleSelectedNodeIdChange: () => {},
handleAddNextNode: () => {},
})
export const useWorkflowContext = () => useContext(WorkflowContext)
import type {
Dispatch,
SetStateAction,
} from 'react'
import {
useCallback,
useMemo,
useState,
} from 'react'
import type { Node } from './types'
import produce from 'immer'
import type { Edge } from 'reactflow'
import type {
BlockEnum,
Node,
} from './types'
import { NodeInitialData } from './constants'
export const useWorkflow = (nodes: Node[], initialSelectedNodeId?: string) => {
export const useWorkflow = (
nodes: Node[],
edges: Edge[],
setNodes: Dispatch<SetStateAction<Node[]>>,
setEdges: Dispatch<SetStateAction<Edge[]>>,
initialSelectedNodeId?: string,
) => {
const [selectedNodeId, setSelectedNodeId] = useState(initialSelectedNodeId)
const handleSelectedNodeIdChange = useCallback((nodeId: string) => setSelectedNodeId(nodeId), [])
......@@ -14,9 +30,40 @@ export const useWorkflow = (nodes: Node[], initialSelectedNodeId?: string) => {
return nodes.find(node => node.id === selectedNodeId)
}, [nodes, selectedNodeId])
const handleAddNextNode = useCallback((prevNode: Node, nextNodeType: BlockEnum) => {
const prevNodeDom = document.querySelector(`.react-flow__node-custom[data-id="${prevNode.id}"]`)
const prevNodeDomHeight = prevNodeDom?.getBoundingClientRect().height || 0
const nextNode = {
id: `node-${Date.now()}`,
type: 'custom',
position: {
x: prevNode.position.x,
y: prevNode.position.y + prevNodeDomHeight + 64,
},
data: NodeInitialData[nextNodeType],
}
const newEdge = {
id: `edge-${Date.now()}`,
source: prevNode.id,
target: nextNode.id,
}
setNodes((oldNodes) => {
return produce(oldNodes, (draft) => {
draft.push(nextNode)
})
})
setEdges((oldEdges) => {
return produce(oldEdges, (draft) => {
draft.push(newEdge)
})
})
}, [setNodes, setEdges])
return {
selectedNodeId,
selectedNode,
handleSelectedNodeIdChange,
handleAddNextNode,
}
}
......@@ -5,6 +5,7 @@ import ReactFlow, {
ReactFlowProvider,
useEdgesState,
useNodesState,
useReactFlow,
} from 'reactflow'
import 'reactflow/dist/style.css'
import {
......@@ -60,21 +61,31 @@ const WorkflowWrap: FC<WorkflowWrapProps> = ({
edges: initialEdges,
selectedNodeId: initialSelectedNodeId,
}) => {
const [nodes] = useNodesState(initialNodes)
const [edges] = useEdgesState(initialEdges)
const reactFlow = useReactFlow()
const [nodes, setNodes] = useNodesState(initialNodes)
const [edges, setEdges] = useEdgesState(initialEdges)
const {
selectedNodeId,
handleSelectedNodeIdChange,
selectedNode,
} = useWorkflow(nodes, initialSelectedNodeId)
handleAddNextNode,
} = useWorkflow(
nodes,
edges,
setNodes,
setEdges,
initialSelectedNodeId,
)
return (
<WorkflowContext.Provider value={{
reactFlow,
selectedNodeId,
handleSelectedNodeIdChange,
selectedNode,
nodes,
edges,
handleAddNextNode,
}}>
<Workflow />
</WorkflowContext.Provider>
......
import type { ComponentType } from 'react'
import { BlockEnum } from '../types'
import StartNode from './start/node'
import StartPanel from './start/panel'
import EndNode from './end/node'
......@@ -23,29 +24,29 @@ import ToolNode from './tool/node'
import ToolPanel from './tool/panel'
export const NodeMap: Record<string, ComponentType> = {
start: StartNode,
end: EndNode,
directAnswer: DirectAnswerNode,
llm: LLMNode,
knowledgeRetrieval: KnowledgeRetrievalNode,
questionClassifier: QuestionClassifierNode,
ifElse: IfElseNode,
code: CodeNode,
templateTransform: TemplateTransformNode,
http: HttpNode,
tool: ToolNode,
[BlockEnum.Start]: StartNode,
[BlockEnum.End]: EndNode,
[BlockEnum.DirectAnswer]: DirectAnswerNode,
[BlockEnum.LLM]: LLMNode,
[BlockEnum.KnowledgeRetrieval]: KnowledgeRetrievalNode,
[BlockEnum.QuestionClassifier]: QuestionClassifierNode,
[BlockEnum.IfElse]: IfElseNode,
[BlockEnum.Code]: CodeNode,
[BlockEnum.TemplateTransform]: TemplateTransformNode,
[BlockEnum.HttpRequest]: HttpNode,
[BlockEnum.Tool]: ToolNode,
}
export const PanelMap: Record<string, ComponentType> = {
start: StartPanel,
end: EndPanel,
directAnswer: DirectAnswerPanel,
llm: LLMPanel,
knowledgeRetrieval: KnowledgeRetrievalPanel,
questionClassifier: QuestionClassifierPanel,
ifElse: IfElsePanel,
code: CodePanel,
templateTransform: TemplateTransformPanel,
http: HttpPanel,
tool: ToolPanel,
[BlockEnum.Start]: StartPanel,
[BlockEnum.End]: EndPanel,
[BlockEnum.DirectAnswer]: DirectAnswerPanel,
[BlockEnum.LLM]: LLMPanel,
[BlockEnum.KnowledgeRetrieval]: KnowledgeRetrievalPanel,
[BlockEnum.QuestionClassifier]: QuestionClassifierPanel,
[BlockEnum.IfElse]: IfElsePanel,
[BlockEnum.Code]: CodePanel,
[BlockEnum.TemplateTransform]: TemplateTransformPanel,
[BlockEnum.HttpRequest]: HttpPanel,
[BlockEnum.Tool]: ToolPanel,
}
......@@ -13,6 +13,7 @@ export enum BlockEnum {
TemplateTransform = 'template-transform',
HttpRequest = 'http-request',
VariableAssigner = 'variable-assigner',
Tool = 'tool',
}
export type NodeData = {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment