refactor: optimize state selection in data source components using useShallow for improved performance

web/app/components/app-sidebar/navLink.tsx

@@ -1,5 +1,5 @@
 'use client'
-
+import React from 'react'
 import { useSelectedLayoutSegment } from 'next/navigation'
 import Link from 'next/link'
 import classNames from '@/utils/classnames'
@@ -22,13 +22,13 @@ export type NavLinkProps = {
   disabled?: boolean
 }
 
-export default function NavLink({
+const NavLink = ({
   name,
   href,
   iconMap,
   mode = 'expand',
   disabled = false,
-}: NavLinkProps) {
+}: NavLinkProps) => {
   const segment = useSelectedLayoutSegment()
   const formattedSegment = (() => {
     let res = segment?.toLowerCase()
@@ -90,3 +90,5 @@ export default function NavLink({
     </Link>
   )
 }
+
+export default React.memo(NavLink)

web/app/components/datasets/documents/create-from-pipeline/data-source/online-documents/index.tsx

@@ -11,6 +11,7 @@ import Toast from '@/app/components/base/toast'
 import type { DataSourceNodeCompletedResponse, DataSourceNodeErrorResponse } from '@/types/pipeline'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import { useDataSourceStore, useDataSourceStoreWithSelector } from '../store'
+import { useShallow } from 'zustand/react/shallow'
 
 type OnlineDocumentsProps = {
   isInPipeline?: boolean
@@ -24,11 +25,17 @@ const OnlineDocuments = ({
   nodeData,
 }: OnlineDocumentsProps) => {
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const documentsData = useDataSourceStoreWithSelector(state => state.documentsData)
-  const searchValue = useDataSourceStoreWithSelector(state => state.searchValue)
-  const selectedPagesId = useDataSourceStoreWithSelector(state => state.selectedPagesId)
-  const currentWorkspaceId = useDataSourceStoreWithSelector(state => state.currentWorkspaceId)
-  const currentNodeIdRef = useDataSourceStoreWithSelector(state => state.currentNodeIdRef)
+  const {
+    documentsData,
+    searchValue,
+    selectedPagesId,
+    currentWorkspaceId,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    documentsData: state.documentsData,
+    searchValue: state.searchValue,
+    selectedPagesId: state.selectedPagesId,
+    currentWorkspaceId: state.currentWorkspaceId,
+  })))
   const dataSourceStore = useDataSourceStore()
 
   const PagesMapAndSelectedPagesId: DataSourceNotionPageMap = useMemo(() => {
@@ -75,15 +82,16 @@ const OnlineDocuments = ({
   }, [dataSourceStore, datasourceNodeRunURL])
 
   useEffect(() => {
+    const {
+      setDocumentsData,
+      setCurrentWorkspaceId,
+      setSearchValue,
+      setSelectedPagesId,
+      setOnlineDocuments,
+      setCurrentDocument,
+      currentNodeIdRef,
+    } = dataSourceStore.getState()
     if (nodeId !== currentNodeIdRef.current) {
-      const {
-        setDocumentsData,
-        setCurrentWorkspaceId,
-        setSearchValue,
-        setSelectedPagesId,
-        setOnlineDocuments,
-        setCurrentDocument,
-      } = dataSourceStore.getState()
       setDocumentsData([])
       setCurrentWorkspaceId('')
       setSearchValue('')

web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx

@@ -11,6 +11,7 @@ import Toast from '@/app/components/base/toast'
 import { useDataSourceStore, useDataSourceStoreWithSelector } from '../store'
 import { convertOnlineDriveData } from './utils'
 import produce from 'immer'
+import { useShallow } from 'zustand/react/shallow'
 
 type OnlineDriveProps = {
   nodeId: string
@@ -24,12 +25,19 @@ const OnlineDrive = ({
   isInPipeline = false,
 }: OnlineDriveProps) => {
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const prefix = useDataSourceStoreWithSelector(state => state.prefix)
-  const keywords = useDataSourceStoreWithSelector(state => state.keywords)
-  const bucket = useDataSourceStoreWithSelector(state => state.bucket)
-  const selectedFileKeys = useDataSourceStoreWithSelector(state => state.selectedFileKeys)
-  const fileList = useDataSourceStoreWithSelector(state => state.fileList)
-  const currentNodeIdRef = useDataSourceStoreWithSelector(state => state.currentNodeIdRef)
+  const {
+    prefix,
+    keywords,
+    bucket,
+    selectedFileKeys,
+    fileList,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    prefix: state.prefix,
+    keywords: state.keywords,
+    bucket: state.bucket,
+    selectedFileKeys: state.selectedFileKeys,
+    fileList: state.fileList,
+  })))
   const dataSourceStore = useDataSourceStore()
   const [isLoading, setIsLoading] = useState(false)
 
@@ -83,14 +91,15 @@ const OnlineDrive = ({
   }, [datasourceNodeRunURL, dataSourceStore])
 
   useEffect(() => {
+    const {
+      setFileList,
+      setBucket,
+      setPrefix,
+      setKeywords,
+      setSelectedFileKeys,
+      currentNodeIdRef,
+    } = dataSourceStore.getState()
     if (nodeId !== currentNodeIdRef.current) {
-      const {
-        setFileList,
-        setBucket,
-        setPrefix,
-        setKeywords,
-        setSelectedFileKeys,
-      } = dataSourceStore.getState()
       setFileList([])
       setBucket('')
       setPrefix([])

web/app/components/datasets/documents/create-from-pipeline/data-source/website-crawl/index.tsx

@@ -22,6 +22,7 @@ import type {
 } from '@/types/pipeline'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import { useDataSourceStore, useDataSourceStoreWithSelector } from '../store'
+import { useShallow } from 'zustand/react/shallow'
 
 const I18N_PREFIX = 'datasetCreation.stepOne.website'
 
@@ -42,10 +43,17 @@ const WebsiteCrawl = ({
   const [crawledNum, setCrawledNum] = useState(0)
   const [crawlErrorMessage, setCrawlErrorMessage] = useState('')
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const crawlResult = useDataSourceStoreWithSelector(state => state.crawlResult)
-  const step = useDataSourceStoreWithSelector(state => state.step)
-  const checkedCrawlResult = useDataSourceStoreWithSelector(state => state.websitePages)
-  const previewIndex = useDataSourceStoreWithSelector(state => state.previewIndex)
+  const {
+    crawlResult,
+    step,
+    checkedCrawlResult,
+    previewIndex,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    crawlResult: state.crawlResult,
+    step: state.step,
+    checkedCrawlResult: state.websitePages,
+    previewIndex: state.previewIndex,
+  })))
   const dataSourceStore = useDataSourceStore()
 
   const usePreProcessingParams = useRef(!isInPipeline ? usePublishedPipelinePreProcessingParams : useDraftPipelinePreProcessingParams)

web/app/components/datasets/documents/create-from-pipeline/hooks.ts

@@ -6,6 +6,7 @@ import { BlockEnum, type Node } from '@/app/components/workflow/types'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import { useDataSourceStore, useDataSourceStoreWithSelector } from './data-source/store'
 import type { DataSourceNotionPageMap, DataSourceNotionWorkspace } from '@/models/common'
+import { useShallow } from 'zustand/react/shallow'
 
 export const useAddDocumentsSteps = () => {
   const { t } = useTranslation()
@@ -62,8 +63,13 @@ export const useDatasourceOptions = (pipelineNodes: Node<DataSourceNodeType>[])
 }
 
 export const useLocalFile = () => {
-  const fileList = useDataSourceStoreWithSelector(state => state.localFileList)
-  const currentLocalFile = useDataSourceStoreWithSelector(state => state.currentLocalFile)
+  const {
+    localFileList: fileList,
+    currentLocalFile,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    localFileList: state.localFileList,
+    currentLocalFile: state.currentLocalFile,
+  })))
   const dataSourceStore = useDataSourceStore()
 
   const allFileLoaded = useMemo(() => (fileList.length > 0 && fileList.every(file => file.file.id)), [fileList])
@@ -82,10 +88,17 @@ export const useLocalFile = () => {
 }
 
 export const useOnlineDocuments = () => {
-  const documentsData = useDataSourceStoreWithSelector(state => state.documentsData)
-  const currentWorkspaceId = useDataSourceStoreWithSelector(state => state.currentWorkspaceId)
-  const onlineDocuments = useDataSourceStoreWithSelector(state => state.onlineDocuments)
-  const currentDocument = useDataSourceStoreWithSelector(state => state.currentDocument)
+  const {
+    documentsData,
+    currentWorkspaceId,
+    onlineDocuments,
+    currentDocument,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    documentsData: state.documentsData,
+    currentWorkspaceId: state.currentWorkspaceId,
+    onlineDocuments: state.onlineDocuments,
+    currentDocument: state.currentDocument,
+  })))
   const dataSourceStore = useDataSourceStore()
 
   const currentWorkspace = documentsData.find(workspace => workspace.workspace_id === currentWorkspaceId)
@@ -119,8 +132,13 @@ export const useOnlineDocuments = () => {
 }
 
 export const useWebsiteCrawl = () => {
-  const websitePages = useDataSourceStoreWithSelector(state => state.websitePages)
-  const currentWebsite = useDataSourceStoreWithSelector(state => state.currentWebsite)
+  const {
+    websitePages,
+    currentWebsite,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    websitePages: state.websitePages,
+    currentWebsite: state.currentWebsite,
+  })))
   const dataSourceStore = useDataSourceStore()
 
   const hideWebsitePreview = useCallback(() => {
@@ -137,8 +155,13 @@ export const useWebsiteCrawl = () => {
 }
 
 export const useOnlineDrive = () => {
-  const fileList = useDataSourceStoreWithSelector(state => state.fileList)
-  const selectedFileKeys = useDataSourceStoreWithSelector(state => state.selectedFileKeys)
+  const {
+    fileList,
+    selectedFileKeys,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    fileList: state.fileList,
+    selectedFileKeys: state.selectedFileKeys,
+  })))
 
   const selectedOnlineDriveFileList = useMemo(() => {
     return selectedFileKeys.map(key => fileList.find(item => item.key === key)!)

web/app/components/rag-pipeline/components/panel/test-run/index.tsx

@@ -17,13 +17,21 @@ import Header from './header'
 import FooterTips from './footer-tips'
 import DataSourceProvider from '@/app/components/datasets/documents/create-from-pipeline/data-source/store/provider'
 import { useDataSourceStore, useDataSourceStoreWithSelector } from '@/app/components/datasets/documents/create-from-pipeline/data-source/store'
+import { useShallow } from 'zustand/react/shallow'
 
 const TestRunPanel = () => {
   const setShowDebugAndPreviewPanel = useWorkflowStoreWithSelector(state => state.setShowDebugAndPreviewPanel)
-  const fileList = useDataSourceStoreWithSelector(state => state.localFileList)
-  const onlineDocuments = useDataSourceStoreWithSelector(state => state.onlineDocuments)
-  const websitePages = useDataSourceStoreWithSelector(state => state.websitePages)
-  const selectedFileKeys = useDataSourceStoreWithSelector(state => state.selectedFileKeys)
+  const {
+    localFileList: fileList,
+    onlineDocuments,
+    websitePages,
+    selectedFileKeys,
+  } = useDataSourceStoreWithSelector(useShallow(state => ({
+    localFileList: state.localFileList,
+    onlineDocuments: state.onlineDocuments,
+    websitePages: state.websitePages,
+    selectedFileKeys: state.selectedFileKeys,
+  })))
   const dataSourceStore = useDataSourceStore()
   const [datasource, setDatasource] = useState<Datasource>()