feat: add datasets detail context and provider for improved data vali… (#16451)

web/app/components/workflow/datasets-detail-store/provider.tsx

@@ -0,0 +1,53 @@
+import type { FC } from 'react'
+import { createContext, useCallback, useEffect, useRef } from 'react'
+import { createDatasetsDetailStore } from './store'
+import type { CommonNodeType, Node } from '../types'
+import { BlockEnum } from '../types'
+import type { KnowledgeRetrievalNodeType } from '../nodes/knowledge-retrieval/types'
+import { fetchDatasets } from '@/service/datasets'
+
+type DatasetsDetailStoreApi = ReturnType<typeof createDatasetsDetailStore>
+
+type DatasetsDetailContextType = DatasetsDetailStoreApi | undefined
+
+export const DatasetsDetailContext = createContext<DatasetsDetailContextType>(undefined)
+
+type DatasetsDetailProviderProps = {
+  nodes: Node[]
+  children: React.ReactNode
+}
+
+const DatasetsDetailProvider: FC<DatasetsDetailProviderProps> = ({
+  nodes,
+  children,
+}) => {
+  const storeRef = useRef<DatasetsDetailStoreApi>()
+
+  if (!storeRef.current)
+    storeRef.current = createDatasetsDetailStore()
+
+  const updateDatasetsDetail = useCallback(async (datasetIds: string[]) => {
+    const { data: datasetsDetail } = await fetchDatasets({ url: '/datasets', params: { page: 1, ids: datasetIds } })
+    if (datasetsDetail && datasetsDetail.length > 0)
+      storeRef.current!.getState().updateDatasetsDetail(datasetsDetail)
+  }, [])
+
+  useEffect(() => {
+    if (!storeRef.current) return
+    const knowledgeRetrievalNodes = nodes.filter(node => node.data.type === BlockEnum.KnowledgeRetrieval)
+    const allDatasetIds = knowledgeRetrievalNodes.reduce<string[]>((acc, node) => {
+      return Array.from(new Set([...acc, ...(node.data as CommonNodeType<KnowledgeRetrievalNodeType>).dataset_ids]))
+    }, [])
+    if (allDatasetIds.length === 0) return
+    updateDatasetsDetail(allDatasetIds)
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  return (
+    <DatasetsDetailContext.Provider value={storeRef.current!}>
+      {children}
+    </DatasetsDetailContext.Provider>
+  )
+}
+
+export default DatasetsDetailProvider

web/app/components/workflow/datasets-detail-store/store.ts

@@ -0,0 +1,38 @@
+import { useContext } from 'react'
+import { createStore, useStore } from 'zustand'
+import type { DataSet } from '@/models/datasets'
+import { DatasetsDetailContext } from './provider'
+import produce from 'immer'
+
+type DatasetsDetailStore = {
+  datasetsDetail: Record<string, DataSet>
+  updateDatasetsDetail: (datasetsDetail: DataSet[]) => void
+}
+
+export const createDatasetsDetailStore = () => {
+  return createStore<DatasetsDetailStore>((set, get) => ({
+    datasetsDetail: {},
+    updateDatasetsDetail: (datasets: DataSet[]) => {
+      const oldDatasetsDetail = get().datasetsDetail
+      const datasetsDetail = datasets.reduce<Record<string, DataSet>>((acc, dataset) => {
+        acc[dataset.id] = dataset
+        return acc
+      }, {})
+      // Merge new datasets detail into old one
+      const newDatasetsDetail = produce(oldDatasetsDetail, (draft) => {
+        Object.entries(datasetsDetail).forEach(([key, value]) => {
+          draft[key] = value
+        })
+      })
+      set({ datasetsDetail: newDatasetsDetail })
+    },
+  }))
+}
+
+export const useDatasetsDetailStore = <T>(selector: (state: DatasetsDetailStore) => T): T => {
+  const store = useContext(DatasetsDetailContext)
+  if (!store)
+    throw new Error('Missing DatasetsDetailContext.Provider in the tree')
+
+  return useStore(store, selector)
+}