Refactor datasets service toward TanStack Query (#29008)

Co-authored-by: Wu Tianwei <30284043+WTW0313@users.noreply.github.com>

web/app/components/app/configuration/dataset-config/select-dataset/index.tsx

@@ -1,18 +1,18 @@
 'use client'
 import type { FC } from 'react'
-import React, { useRef, useState } from 'react'
-import { useGetState, useInfiniteScroll } from 'ahooks'
+import React, { useEffect, useMemo, useRef, useState } from 'react'
+import { useInfiniteScroll } from 'ahooks'
 import { useTranslation } from 'react-i18next'
 import Link from 'next/link'
 import Modal from '@/app/components/base/modal'
 import type { DataSet } from '@/models/datasets'
 import Button from '@/app/components/base/button'
-import { fetchDatasets } from '@/service/datasets'
 import Loading from '@/app/components/base/loading'
 import Badge from '@/app/components/base/badge'
 import { useKnowledge } from '@/hooks/use-knowledge'
 import cn from '@/utils/classnames'
 import AppIcon from '@/app/components/base/app-icon'
+import { useInfiniteDatasets } from '@/service/knowledge/use-dataset'
 import { ModelFeatureEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
 import FeatureIcon from '@/app/components/header/account-setting/model-provider-page/model-selector/feature-icon'
 
@@ -30,51 +30,70 @@ const SelectDataSet: FC<ISelectDataSetProps> = ({
   onSelect,
 }) => {
   const { t } = useTranslation()
-  const [selected, setSelected] = React.useState<DataSet[]>([])
-  const [loaded, setLoaded] = React.useState(false)
-  const [datasets, setDataSets] = React.useState<DataSet[] | null>(null)
-  const [hasInitialized, setHasInitialized] = React.useState(false)
-  const hasNoData = !datasets || datasets?.length === 0
+  const [selected, setSelected] = useState<DataSet[]>([])
   const canSelectMulti = true
+  const { formatIndexingTechniqueAndMethod } = useKnowledge()
+  const { data, isLoading, isFetchingNextPage, fetchNextPage, hasNextPage } = useInfiniteDatasets(
+    { page: 1 },
+    { enabled: isShow, staleTime: 0, refetchOnMount: 'always' },
+  )
+  const pages = data?.pages || []
+  const datasets = useMemo(() => {
+    return pages.flatMap(page => page.data.filter(item => item.indexing_technique || item.provider === 'external'))
+  }, [pages])
+  const hasNoData = !isLoading && datasets.length === 0
 
   const listRef = useRef<HTMLDivElement>(null)
-  const [page, setPage, getPage] = useGetState(1)
-  const [isNoMore, setIsNoMore] = useState(false)
-  const { formatIndexingTechniqueAndMethod } = useKnowledge()
+  const isNoMore = hasNextPage === false
 
   useInfiniteScroll(
     async () => {
-      if (!isNoMore) {
-        const { data, has_more } = await fetchDatasets({ url: '/datasets', params: { page } })
-        setPage(getPage() + 1)
-        setIsNoMore(!has_more)
-        const newList = [...(datasets || []), ...data.filter(item => item.indexing_technique || item.provider === 'external')]
-        setDataSets(newList)
-        setLoaded(true)
-
-        // Initialize selected datasets based on selectedIds and available datasets
-        if (!hasInitialized) {
-          if (selectedIds.length > 0) {
-            const validSelectedDatasets = selectedIds
-              .map(id => newList.find(item => item.id === id))
-              .filter(Boolean) as DataSet[]
-            setSelected(validSelectedDatasets)
-          }
-          setHasInitialized(true)
-        }
-      }
+      if (!hasNextPage || isFetchingNextPage)
+        return { list: [] }
+      await fetchNextPage()
       return { list: [] }
     },
     {
       target: listRef,
-      isNoMore: () => {
-        return isNoMore
-      },
-      reloadDeps: [isNoMore],
+      isNoMore: () => isNoMore,
+      reloadDeps: [isNoMore, isFetchingNextPage],
     },
   )
 
+  const prevSelectedIdsRef = useRef<string[]>([])
+  const hasUserModifiedSelectionRef = useRef(false)
+  useEffect(() => {
+    if (isShow)
+      hasUserModifiedSelectionRef.current = false
+  }, [isShow])
+  useEffect(() => {
+    const prevSelectedIds = prevSelectedIdsRef.current
+    const idsChanged = selectedIds.length !== prevSelectedIds.length
+      || selectedIds.some((id, idx) => id !== prevSelectedIds[idx])
+
+    if (!selectedIds.length && (!hasUserModifiedSelectionRef.current || idsChanged)) {
+      setSelected([])
+      prevSelectedIdsRef.current = selectedIds
+      hasUserModifiedSelectionRef.current = false
+      return
+    }
+
+    if (!idsChanged && hasUserModifiedSelectionRef.current)
+      return
+
+    setSelected((prev) => {
+      const prevMap = new Map(prev.map(item => [item.id, item]))
+      const nextSelected = selectedIds
+        .map(id => datasets.find(item => item.id === id) || prevMap.get(id))
+        .filter(Boolean) as DataSet[]
+      return nextSelected
+    })
+    prevSelectedIdsRef.current = selectedIds
+    hasUserModifiedSelectionRef.current = false
+  }, [datasets, selectedIds])
+
   const toggleSelect = (dataSet: DataSet) => {
+    hasUserModifiedSelectionRef.current = true
     const isSelected = selected.some(item => item.id === dataSet.id)
     if (isSelected) {
       setSelected(selected.filter(item => item.id !== dataSet.id))
@@ -98,13 +117,13 @@ const SelectDataSet: FC<ISelectDataSetProps> = ({
       className='w-[400px]'
       title={t('appDebug.feature.dataSet.selectTitle')}
     >
-      {!loaded && (
+      {(isLoading && datasets.length === 0) && (
         <div className='flex h-[200px]'>
           <Loading type='area' />
         </div>
       )}
 
-      {(loaded && hasNoData) && (
+      {hasNoData && (
         <div className='mt-6 flex h-[128px] items-center justify-center space-x-1  rounded-lg border text-[13px]'
           style={{
             background: 'rgba(0, 0, 0, 0.02)',
@@ -116,7 +135,7 @@ const SelectDataSet: FC<ISelectDataSetProps> = ({
         </div>
       )}
 
-      {datasets && datasets?.length > 0 && (
+      {datasets.length > 0 && (
         <>
           <div ref={listRef} className='mt-7 max-h-[286px] space-y-1 overflow-y-auto'>
             {datasets.map(item => (
@@ -171,7 +190,7 @@ const SelectDataSet: FC<ISelectDataSetProps> = ({
           </div>
         </>
       )}
-      {loaded && (
+      {!isLoading && (
         <div className='mt-8 flex items-center justify-between'>
           <div className='text-sm  font-medium text-text-secondary'>
             {selected.length > 0 && `${selected.length} ${t('appDebug.feature.dataSet.selected')}`}

web/app/components/datasets/common/document-status-with-action/index-failed.tsx

@@ -2,11 +2,11 @@
 import type { FC } from 'react'
 import React, { useEffect, useReducer } from 'react'
 import { useTranslation } from 'react-i18next'
-import useSWR from 'swr'
 import StatusWithAction from './status-with-action'
-import { getErrorDocs, retryErrorDocs } from '@/service/datasets'
+import { retryErrorDocs } from '@/service/datasets'
 import type { IndexingStatusResponse } from '@/models/datasets'
 import { noop } from 'lodash-es'
+import { useDatasetErrorDocs } from '@/service/knowledge/use-dataset'
 
 type Props = {
   datasetId: string
@@ -35,16 +35,19 @@ const indexStateReducer = (state: IIndexState, action: IAction) => {
 const RetryButton: FC<Props> = ({ datasetId }) => {
   const { t } = useTranslation()
   const [indexState, dispatch] = useReducer(indexStateReducer, { value: 'success' })
-  const { data: errorDocs, isLoading } = useSWR({ datasetId }, getErrorDocs)
+  const { data: errorDocs, isLoading, refetch: refetchErrorDocs } = useDatasetErrorDocs(datasetId)
 
   const onRetryErrorDocs = async () => {
     dispatch({ type: 'retry' })
     const document_ids = errorDocs?.data.map((doc: IndexingStatusResponse) => doc.id) || []
     const res = await retryErrorDocs({ datasetId, document_ids })
-    if (res.result === 'success')
+    if (res.result === 'success') {
+      refetchErrorDocs()
       dispatch({ type: 'success' })
-    else
+    }
+    else {
       dispatch({ type: 'error' })
+    }
   }
 
   useEffect(() => {

web/app/components/datasets/create/embedding-process/index.tsx

@@ -1,9 +1,7 @@
 import type { FC } from 'react'
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'
-import useSWR from 'swr'
 import { useRouter } from 'next/navigation'
 import { useTranslation } from 'react-i18next'
-import { omit } from 'lodash-es'
 import {
   RiArrowRightLine,
   RiCheckboxCircleFill,
@@ -25,7 +23,7 @@ import type {
   LegacyDataSourceInfo,
   ProcessRuleResponse,
 } from '@/models/datasets'
-import { fetchIndexingStatusBatch as doFetchIndexingStatus, fetchProcessRule } from '@/service/datasets'
+import { fetchIndexingStatusBatch as doFetchIndexingStatus } from '@/service/datasets'
 import { DataSourceType, ProcessMode } from '@/models/datasets'
 import NotionIcon from '@/app/components/base/notion-icon'
 import PriorityLabel from '@/app/components/billing/priority-label'
@@ -40,6 +38,7 @@ import { useInvalidDocumentList } from '@/service/knowledge/use-document'
 import Divider from '@/app/components/base/divider'
 import { useDatasetApiAccessUrl } from '@/hooks/use-api-access-url'
 import Link from 'next/link'
+import { useProcessRule } from '@/service/knowledge/use-dataset'
 
 type Props = {
   datasetId: string
@@ -207,12 +206,7 @@ const EmbeddingProcess: FC<Props> = ({ datasetId, batchId, documents = [], index
   }, [])
 
   // get rule
-  const { data: ruleDetail } = useSWR({
-    action: 'fetchProcessRule',
-    params: { documentId: getFirstDocument.id },
-  }, apiParams => fetchProcessRule(omit(apiParams, 'action')), {
-    revalidateOnFocus: false,
-  })
+  const { data: ruleDetail } = useProcessRule(getFirstDocument?.id)
 
   const router = useRouter()
   const invalidDocumentList = useInvalidDocumentList()

web/app/components/datasets/create/file-uploader/index.tsx

@@ -2,7 +2,6 @@
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { useContext } from 'use-context-selector'
-import useSWR from 'swr'
 import { RiDeleteBinLine, RiUploadCloud2Line } from '@remixicon/react'
 import DocumentFileIcon from '../../common/document-file-icon'
 import cn from '@/utils/classnames'
@@ -11,8 +10,7 @@ import { ToastContext } from '@/app/components/base/toast'
 import SimplePieChart from '@/app/components/base/simple-pie-chart'
 
 import { upload } from '@/service/base'
-import { fetchFileUploadConfig } from '@/service/common'
-import { fetchSupportFileTypes } from '@/service/datasets'
+import { useFileSupportTypes, useFileUploadConfig } from '@/service/use-common'
 import I18n from '@/context/i18n'
 import { LanguagesSupported } from '@/i18n-config/language'
 import { IS_CE_EDITION } from '@/config'
@@ -48,8 +46,8 @@ const FileUploader = ({
   const fileUploader = useRef<HTMLInputElement>(null)
   const hideUpload = notSupportBatchUpload && fileList.length > 0
 
-  const { data: fileUploadConfigResponse } = useSWR({ url: '/files/upload' }, fetchFileUploadConfig)
-  const { data: supportFileTypesResponse } = useSWR({ url: '/files/support-type' }, fetchSupportFileTypes)
+  const { data: fileUploadConfigResponse } = useFileUploadConfig()
+  const { data: supportFileTypesResponse } = useFileSupportTypes()
   const supportTypes = supportFileTypesResponse?.allowed_extensions || []
   const supportTypesShowNames = (() => {
     const extensionMap: { [key: string]: string } = {

web/app/components/datasets/documents/detail/batch-modal/csv-uploader.tsx

@@ -13,11 +13,10 @@ import Button from '@/app/components/base/button'
 import type { FileItem } from '@/models/datasets'
 import { upload } from '@/service/base'
 import { getFileUploadErrorMessage } from '@/app/components/base/file-uploader/utils'
-import useSWR from 'swr'
-import { fetchFileUploadConfig } from '@/service/common'
 import SimplePieChart from '@/app/components/base/simple-pie-chart'
 import { Theme } from '@/types/app'
 import useTheme from '@/hooks/use-theme'
+import { useFileUploadConfig } from '@/service/use-common'
 
 export type Props = {
   file: FileItem | undefined
@@ -34,7 +33,7 @@ const CSVUploader: FC<Props> = ({
   const dropRef = useRef<HTMLDivElement>(null)
   const dragRef = useRef<HTMLDivElement>(null)
   const fileUploader = useRef<HTMLInputElement>(null)
-  const { data: fileUploadConfigResponse } = useSWR({ url: '/files/upload' }, fetchFileUploadConfig)
+  const { data: fileUploadConfigResponse } = useFileUploadConfig()
   const fileUploadConfig = useMemo(() => fileUploadConfigResponse ?? {
     file_size_limit: 15,
   }, [fileUploadConfigResponse])

web/app/components/datasets/documents/detail/embedding/index.tsx

@@ -1,9 +1,7 @@
 import type { FC } from 'react'
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'
-import useSWR from 'swr'
 import { useContext } from 'use-context-selector'
 import { useTranslation } from 'react-i18next'
-import { omit } from 'lodash-es'
 import { RiLoader2Line, RiPauseCircleLine, RiPlayCircleLine } from '@remixicon/react'
 import Image from 'next/image'
 import { FieldInfo } from '../metadata'
@@ -21,10 +19,10 @@ import type { CommonResponse } from '@/models/common'
 import { asyncRunSafe, sleep } from '@/utils'
 import {
   fetchIndexingStatus as doFetchIndexingStatus,
-  fetchProcessRule,
   pauseDocIndexing,
   resumeDocIndexing,
 } from '@/service/datasets'
+import { useProcessRule } from '@/service/knowledge/use-dataset'
 
 type IEmbeddingDetailProps = {
   datasetId?: string
@@ -207,12 +205,7 @@ const EmbeddingDetail: FC<IEmbeddingDetailProps> = ({
     }
   }, [startQueryStatus, stopQueryStatus])
 
-  const { data: ruleDetail } = useSWR({
-    action: 'fetchProcessRule',
-    params: { documentId: localDocumentId },
-  }, apiParams => fetchProcessRule(omit(apiParams, 'action')), {
-    revalidateOnFocus: false,
-  })
+  const { data: ruleDetail } = useProcessRule(localDocumentId)
 
   const isEmbedding = useMemo(() => ['indexing', 'splitting', 'parsing', 'cleaning'].includes(indexingStatusDetail?.indexing_status || ''), [indexingStatusDetail])
   const isEmbeddingCompleted = useMemo(() => ['completed'].includes(indexingStatusDetail?.indexing_status || ''), [indexingStatusDetail])

web/app/components/datasets/hit-testing/index.tsx

@@ -32,9 +32,8 @@ import Records from './components/records'
 import {
   useExternalKnowledgeBaseHitTesting,
   useHitTesting,
-  useHitTestingRecords,
-  useInvalidateHitTestingRecords,
 } from '@/service/knowledge/use-hit-testing'
+import { useDatasetTestingRecords } from '@/service/knowledge/use-dataset'
 
 const limit = 10
 
@@ -48,14 +47,13 @@ const HitTestingPage: FC<Props> = ({ datasetId }: Props) => {
   const media = useBreakpoints()
   const isMobile = media === MediaType.mobile
 
-  const [hitResult, setHitResult] = useState<HitTestingResponse | undefined>() // 初始化记录为空数组
+  const [hitResult, setHitResult] = useState<HitTestingResponse | undefined>()
   const [externalHitResult, setExternalHitResult] = useState<ExternalKnowledgeBaseHitTestingResponse | undefined>()
   const [queries, setQueries] = useState<Query[]>([])
   const [queryInputKey, setQueryInputKey] = useState(Date.now())
 
   const [currPage, setCurrPage] = useState<number>(0)
-  const { data: recordsRes, isLoading: isRecordsLoading } = useHitTestingRecords({ datasetId, page: currPage + 1, limit })
-  const invalidateHitTestingRecords = useInvalidateHitTestingRecords(datasetId)
+  const { data: recordsRes, refetch: recordsRefetch, isLoading: isRecordsLoading } = useDatasetTestingRecords(datasetId, { limit, page: currPage + 1 })
 
   const total = recordsRes?.total || 0
 
@@ -107,8 +105,7 @@ const HitTestingPage: FC<Props> = ({ datasetId }: Props) => {
   )
 
   const handleClickRecord = useCallback((record: HitTestingRecord) => {
-    const { queries } = record
-    setQueries(queries)
+    setQueries(record.queries)
     setQueryInputKey(Date.now())
   }, [])
 
@@ -128,7 +125,7 @@ const HitTestingPage: FC<Props> = ({ datasetId }: Props) => {
           setHitResult={setHitResult}
           setExternalHitResult={setExternalHitResult}
           onSubmit={showRightPanel}
-          onUpdateList={invalidateHitTestingRecords}
+          onUpdateList={recordsRefetch}
           loading={isRetrievalLoading}
           queries={queries}
           setQueries={setQueries}
@@ -140,11 +137,9 @@ const HitTestingPage: FC<Props> = ({ datasetId }: Props) => {
           externalKnowledgeBaseHitTestingMutation={externalKnowledgeBaseHitTestingMutation}
         />
         <div className='mb-3 mt-6 text-base font-semibold text-text-primary'>{t('datasetHitTesting.records')}</div>
-        {isRecordsLoading
-          && (
-            <div className='flex-1'><Loading type='app' /></div>
-          )
-        }
+        {isRecordsLoading && (
+          <div className='flex-1'><Loading type='app' /></div>
+        )}
         {!isRecordsLoading && recordsRes?.data && recordsRes.data.length > 0 && (
           <>
             <Records records={recordsRes?.data} onClickRecord={handleClickRecord}/>

web/app/components/develop/secret-key/secret-key-modal.tsx

@@ -5,7 +5,6 @@ import {
 import { useTranslation } from 'react-i18next'
 import { RiDeleteBinLine } from '@remixicon/react'
 import { PlusIcon, XMarkIcon } from '@heroicons/react/20/solid'
-import useSWR from 'swr'
 import SecretKeyGenerateModal from './secret-key-generate'
 import s from './style.module.css'
 import ActionButton from '@/app/components/base/action-button'
@@ -19,7 +18,6 @@ import {
 import {
   createApikey as createDatasetApikey,
   delApikey as delDatasetApikey,
-  fetchApiKeysList as fetchDatasetApiKeysList,
 } from '@/service/datasets'
 import type { CreateApiKeyResponse } from '@/models/app'
 import Loading from '@/app/components/base/loading'
@@ -27,6 +25,7 @@ import Confirm from '@/app/components/base/confirm'
 import useTimestamp from '@/hooks/use-timestamp'
 import { useAppContext } from '@/context/app-context'
 import { useAppApiKeys, useInvalidateAppApiKeys } from '@/service/use-apps'
+import { useDatasetApiKeys, useInvalidateDatasetApiKeys } from '@/service/knowledge/use-dataset'
 
 type ISecretKeyModalProps = {
   isShow: boolean
@@ -46,11 +45,9 @@ const SecretKeyModal = ({
   const [isVisible, setVisible] = useState(false)
   const [newKey, setNewKey] = useState<CreateApiKeyResponse | undefined>(undefined)
   const invalidateAppApiKeys = useInvalidateAppApiKeys()
+  const invalidateDatasetApiKeys = useInvalidateDatasetApiKeys()
   const { data: appApiKeys, isLoading: isAppApiKeysLoading } = useAppApiKeys(appId, { enabled: !!appId && isShow })
-  const { data: datasetApiKeys, isLoading: isDatasetApiKeysLoading, mutate: mutateDatasetApiKeys } = useSWR(
-    !appId && isShow ? { url: '/datasets/api-keys', params: {} } : null,
-    fetchDatasetApiKeysList,
-  )
+  const { data: datasetApiKeys, isLoading: isDatasetApiKeysLoading } = useDatasetApiKeys({ enabled: !appId && isShow })
   const apiKeysList = appId ? appApiKeys : datasetApiKeys
   const isApiKeysLoading = appId ? isAppApiKeysLoading : isDatasetApiKeysLoading
 
@@ -69,7 +66,7 @@ const SecretKeyModal = ({
     if (appId)
       invalidateAppApiKeys(appId)
     else
-      mutateDatasetApiKeys()
+      invalidateDatasetApiKeys()
   }
 
   const onCreate = async () => {
@@ -83,7 +80,7 @@ const SecretKeyModal = ({
     if (appId)
       invalidateAppApiKeys(appId)
     else
-      mutateDatasetApiKeys()
+      invalidateDatasetApiKeys()
   }
 
   const generateToken = (token: string) => {

web/app/components/workflow/nodes/document-extractor/panel.tsx

@@ -1,6 +1,5 @@
 import type { FC } from 'react'
 import React from 'react'
-import useSWR from 'swr'
 import { useTranslation } from 'react-i18next'
 import { useContext } from 'use-context-selector'
 import VarReferencePicker from '../_base/components/variable/var-reference-picker'
@@ -9,11 +8,11 @@ import Split from '../_base/components/split'
 import { useNodeHelpLink } from '../_base/hooks/use-node-help-link'
 import useConfig from './use-config'
 import type { DocExtractorNodeType } from './types'
-import { fetchSupportFileTypes } from '@/service/datasets'
 import Field from '@/app/components/workflow/nodes/_base/components/field'
 import { BlockEnum, type NodePanelProps } from '@/app/components/workflow/types'
 import I18n from '@/context/i18n'
 import { LanguagesSupported } from '@/i18n-config/language'
+import { useFileSupportTypes } from '@/service/use-common'
 
 const i18nPrefix = 'workflow.nodes.docExtractor'
 
@@ -24,7 +23,7 @@ const Panel: FC<NodePanelProps<DocExtractorNodeType>> = ({
   const { t } = useTranslation()
   const { locale } = useContext(I18n)
   const link = useNodeHelpLink(BlockEnum.DocExtractor)
-  const { data: supportFileTypesResponse } = useSWR({ url: '/files/support-type' }, fetchSupportFileTypes)
+  const { data: supportFileTypesResponse } = useFileSupportTypes()
   const supportTypes = supportFileTypesResponse?.allowed_extensions || []
   const supportTypesShowNames = (() => {
     const extensionMap: { [key: string]: string } = {