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 BlockIcon from '../block-icon'
import { useWorkflowContext } from '../context'
import { import {
BLOCK_CLASSIFICATIONS, BLOCK_CLASSIFICATIONS,
BLOCK_GROUP_BY_CLASSIFICATION, BLOCK_GROUP_BY_CLASSIFICATION,
...@@ -7,7 +12,13 @@ import { ...@@ -7,7 +12,13 @@ import {
} from './constants' } from './constants'
const Tabs = () => { const Tabs = () => {
const {
nodes,
handleAddNextNode,
} = useWorkflowContext()
const [activeTab, setActiveTab] = useState(TABS[0].key) const [activeTab, setActiveTab] = useState(TABS[0].key)
const nodeId = useNodeId()
const currentNode = nodes.find(node => node.id === nodeId)
return ( return (
<div> <div>
...@@ -46,6 +57,10 @@ const Tabs = () => { ...@@ -46,6 +57,10 @@ const Tabs = () => {
<div <div
key={block.type} key={block.type}
className='flex items-center px-3 h-8 rounded-lg hover:bg-gray-50 cursor-pointer' 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 <BlockIcon
className='mr-2' className='mr-2'
...@@ -63,4 +78,4 @@ const Tabs = () => { ...@@ -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' 'use client'
import { createContext, useContext } from 'use-context-selector' import { createContext, useContext } from 'use-context-selector'
import type { Edge } from 'reactflow' import type {
import type { Node } from './types' Edge,
ReactFlowInstance,
} from 'reactflow'
import type {
BlockEnum,
Node,
} from './types'
export type WorkflowContextValue = { export type WorkflowContextValue = {
reactFlow: ReactFlowInstance
nodes: Node[] nodes: Node[]
edges: Edge[] edges: Edge[]
selectedNodeId?: string selectedNodeId?: string
handleSelectedNodeIdChange: (nodeId: string) => void handleSelectedNodeIdChange: (nodeId: string) => void
selectedNode?: Node selectedNode?: Node
handleAddNextNode: (prevNode: Node, nextNodeType: BlockEnum) => void
} }
export const WorkflowContext = createContext<WorkflowContextValue>({ export const WorkflowContext = createContext<WorkflowContextValue>({
reactFlow: null as any,
nodes: [], nodes: [],
edges: [], edges: [],
handleSelectedNodeIdChange: () => {}, handleSelectedNodeIdChange: () => {},
handleAddNextNode: () => {},
}) })
export const useWorkflowContext = () => useContext(WorkflowContext) export const useWorkflowContext = () => useContext(WorkflowContext)
import type {
Dispatch,
SetStateAction,
} from 'react'
import { import {
useCallback, useCallback,
useMemo, useMemo,
useState, useState,
} from 'react' } 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 [selectedNodeId, setSelectedNodeId] = useState(initialSelectedNodeId)
const handleSelectedNodeIdChange = useCallback((nodeId: string) => setSelectedNodeId(nodeId), []) const handleSelectedNodeIdChange = useCallback((nodeId: string) => setSelectedNodeId(nodeId), [])
...@@ -14,9 +30,40 @@ export const useWorkflow = (nodes: Node[], initialSelectedNodeId?: string) => { ...@@ -14,9 +30,40 @@ export const useWorkflow = (nodes: Node[], initialSelectedNodeId?: string) => {
return nodes.find(node => node.id === selectedNodeId) return nodes.find(node => node.id === selectedNodeId)
}, [nodes, 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 { return {
selectedNodeId, selectedNodeId,
selectedNode, selectedNode,
handleSelectedNodeIdChange, handleSelectedNodeIdChange,
handleAddNextNode,
} }
} }
...@@ -5,6 +5,7 @@ import ReactFlow, { ...@@ -5,6 +5,7 @@ import ReactFlow, {
ReactFlowProvider, ReactFlowProvider,
useEdgesState, useEdgesState,
useNodesState, useNodesState,
useReactFlow,
} from 'reactflow' } from 'reactflow'
import 'reactflow/dist/style.css' import 'reactflow/dist/style.css'
import { import {
...@@ -60,21 +61,31 @@ const WorkflowWrap: FC<WorkflowWrapProps> = ({ ...@@ -60,21 +61,31 @@ const WorkflowWrap: FC<WorkflowWrapProps> = ({
edges: initialEdges, edges: initialEdges,
selectedNodeId: initialSelectedNodeId, selectedNodeId: initialSelectedNodeId,
}) => { }) => {
const [nodes] = useNodesState(initialNodes) const reactFlow = useReactFlow()
const [edges] = useEdgesState(initialEdges) const [nodes, setNodes] = useNodesState(initialNodes)
const [edges, setEdges] = useEdgesState(initialEdges)
const { const {
selectedNodeId, selectedNodeId,
handleSelectedNodeIdChange, handleSelectedNodeIdChange,
selectedNode, selectedNode,
} = useWorkflow(nodes, initialSelectedNodeId) handleAddNextNode,
} = useWorkflow(
nodes,
edges,
setNodes,
setEdges,
initialSelectedNodeId,
)
return ( return (
<WorkflowContext.Provider value={{ <WorkflowContext.Provider value={{
reactFlow,
selectedNodeId, selectedNodeId,
handleSelectedNodeIdChange, handleSelectedNodeIdChange,
selectedNode, selectedNode,
nodes, nodes,
edges, edges,
handleAddNextNode,
}}> }}>
<Workflow /> <Workflow />
</WorkflowContext.Provider> </WorkflowContext.Provider>
......
import type { ComponentType } from 'react' import type { ComponentType } from 'react'
import { BlockEnum } from '../types'
import StartNode from './start/node' import StartNode from './start/node'
import StartPanel from './start/panel' import StartPanel from './start/panel'
import EndNode from './end/node' import EndNode from './end/node'
...@@ -23,29 +24,29 @@ import ToolNode from './tool/node' ...@@ -23,29 +24,29 @@ import ToolNode from './tool/node'
import ToolPanel from './tool/panel' import ToolPanel from './tool/panel'
export const NodeMap: Record<string, ComponentType> = { export const NodeMap: Record<string, ComponentType> = {
start: StartNode, [BlockEnum.Start]: StartNode,
end: EndNode, [BlockEnum.End]: EndNode,
directAnswer: DirectAnswerNode, [BlockEnum.DirectAnswer]: DirectAnswerNode,
llm: LLMNode, [BlockEnum.LLM]: LLMNode,
knowledgeRetrieval: KnowledgeRetrievalNode, [BlockEnum.KnowledgeRetrieval]: KnowledgeRetrievalNode,
questionClassifier: QuestionClassifierNode, [BlockEnum.QuestionClassifier]: QuestionClassifierNode,
ifElse: IfElseNode, [BlockEnum.IfElse]: IfElseNode,
code: CodeNode, [BlockEnum.Code]: CodeNode,
templateTransform: TemplateTransformNode, [BlockEnum.TemplateTransform]: TemplateTransformNode,
http: HttpNode, [BlockEnum.HttpRequest]: HttpNode,
tool: ToolNode, [BlockEnum.Tool]: ToolNode,
} }
export const PanelMap: Record<string, ComponentType> = { export const PanelMap: Record<string, ComponentType> = {
start: StartPanel, [BlockEnum.Start]: StartPanel,
end: EndPanel, [BlockEnum.End]: EndPanel,
directAnswer: DirectAnswerPanel, [BlockEnum.DirectAnswer]: DirectAnswerPanel,
llm: LLMPanel, [BlockEnum.LLM]: LLMPanel,
knowledgeRetrieval: KnowledgeRetrievalPanel, [BlockEnum.KnowledgeRetrieval]: KnowledgeRetrievalPanel,
questionClassifier: QuestionClassifierPanel, [BlockEnum.QuestionClassifier]: QuestionClassifierPanel,
ifElse: IfElsePanel, [BlockEnum.IfElse]: IfElsePanel,
code: CodePanel, [BlockEnum.Code]: CodePanel,
templateTransform: TemplateTransformPanel, [BlockEnum.TemplateTransform]: TemplateTransformPanel,
http: HttpPanel, [BlockEnum.HttpRequest]: HttpPanel,
tool: ToolPanel, [BlockEnum.Tool]: ToolPanel,
} }
...@@ -13,6 +13,7 @@ export enum BlockEnum { ...@@ -13,6 +13,7 @@ export enum BlockEnum {
TemplateTransform = 'template-transform', TemplateTransform = 'template-transform',
HttpRequest = 'http-request', HttpRequest = 'http-request',
VariableAssigner = 'variable-assigner', VariableAssigner = 'variable-assigner',
Tool = 'tool',
} }
export type NodeData = { 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