feat: add current vars store

This commit is contained in:
Joel 2025-04-16 15:50:20 +08:00
parent 9d7357058a
commit 379d382914
4 changed files with 132 additions and 6 deletions

View File

@ -0,0 +1,30 @@
import type { FC } from 'react'
import { createContext, useRef } from 'react'
import { createCurrentVarsStore } from './store'
type CurrentVarsStoreApi = ReturnType<typeof createCurrentVarsStore>
type CurrentVarsContextType = CurrentVarsStoreApi | undefined
export const CurrentVarsContext = createContext<CurrentVarsContextType>(undefined)
type CurrentVarsProviderProps = {
children: React.ReactNode
}
const CurrentVarsProvider: FC<CurrentVarsProviderProps> = ({
children,
}) => {
const storeRef = useRef<CurrentVarsStoreApi>()
if (!storeRef.current)
storeRef.current = createCurrentVarsStore()
return (
<CurrentVarsContext.Provider value={storeRef.current!}>
{children}
</CurrentVarsContext.Provider>
)
}
export default CurrentVarsProvider

View File

@ -0,0 +1,91 @@
import { useContext } from 'react'
import { createStore, useStore } from 'zustand'
import { CurrentVarsContext } from './provider'
import produce from 'immer'
type NodeVars = {
id: string
name: string
type: string
vars: {
key: string
type: string
value: any
}[]
}
type CurrentVarsState = {
nodes: NodeVars[]
}
type CurrentVarsActions = {
setVars: (vars: NodeVars[]) => void
getVars: () => NodeVars[]
clearVars: () => void
setNodeVars: (nodeId: string, payload: NodeVars) => void
clearNodeVars: (nodeId: string) => void
getNodeVars: (nodeId: string) => NodeVars | undefined
}
type CurrentVarsStore = CurrentVarsState & CurrentVarsActions
export const createCurrentVarsStore = () => {
return createStore<CurrentVarsStore>((set, get) => ({
nodes: [{
id: '',
name: '',
type: '',
vars: [],
}],
setVars: (vars) => {
set(() => ({
nodes: vars,
}))
},
getVars: () => {
return get().nodes
},
clearVars: () => {
set(() => ({
nodes: [],
}))
},
setNodeVars: (nodeId, vars) => {
set((state) => {
// eslint-disable-next-line sonarjs/no-nested-functions
const nodes = state.nodes.map((node) => {
if (node.id === nodeId) {
return produce(node, (draft) => {
draft.vars = vars.vars
})
}
return node
})
return {
nodes,
}
})
},
clearNodeVars: (nodeId) => {
set(produce((state: CurrentVarsStore) => {
// eslint-disable-next-line sonarjs/no-nested-functions
const nodes = state.nodes.filter(node => node.id !== nodeId)
state.nodes = nodes
},
))
},
getNodeVars: (nodeId) => {
const nodes = get().nodes
return nodes.find(node => node.id === nodeId)
},
}))
}
export const useCurrentVarsStore = <T>(selector: (state: CurrentVarsStore) => T): T => {
const store = useContext(CurrentVarsContext)
if (!store)
throw new Error('Missing CurrentVarsContext.Provider in the tree')
return useStore(store, selector)
}

View File

@ -102,6 +102,7 @@ import Confirm from '@/app/components/base/confirm'
import { FILE_EXTS } from '@/app/components/base/prompt-editor/constants'
import { fetchFileUploadConfig } from '@/service/common'
import DatasetsDetailProvider from './datasets-detail-store/provider'
import CurrentVarsProvider from './current-vars-store/provider'
const nodeTypes = {
[CUSTOM_NODE]: CustomNode,
@ -453,11 +454,13 @@ const WorkflowWrap = memo(() => {
edges={edgesData} >
<FeaturesProvider features={initialFeatures}>
<DatasetsDetailProvider nodes={nodesData}>
<Workflow
nodes={nodesData}
edges={edgesData}
viewport={data?.graph.viewport}
/>
<CurrentVarsProvider>
<Workflow
nodes={nodesData}
edges={edgesData}
viewport={data?.graph.viewport}
/>
</CurrentVarsProvider>
</DatasetsDetailProvider>
</FeaturesProvider>
</WorkflowHistoryProvider>

View File

@ -20,6 +20,7 @@ import type { Props as FormProps } from '@/app/components/workflow/nodes/_base/c
import ResultPanel from '@/app/components/workflow/run/result-panel'
import Tooltip from '@/app/components/base/tooltip'
import Editor from '@/app/components/workflow/nodes/_base/components/prompt/editor'
import { useCurrentVarsStore } from '../../current-vars-store/store'
const i18nPrefix = 'workflow.nodes.llm'
@ -28,7 +29,8 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
data,
}) => {
const { t } = useTranslation()
const currentVars = useCurrentVarsStore(state => state.getVars())
console.log(currentVars)
const {
readOnly,
inputs,