Refactor datasets service toward TanStack Query (#29008)

Co-authored-by: Wu Tianwei <30284043+WTW0313@users.noreply.github.com>
2026-05-05 01:48:04 +08:00 · 2025-12-09 13:44:45 +08:00
parent 57d244de69
commit 18601d8b38
12 changed files with 270 additions and 152 deletions
--- a/web/service/datasets.ts
+++ b/web/service/datasets.ts
@ -1,4 +1,3 @@
-import type { Fetcher } from 'swr'
 import qs from 'qs'
 import { del, get, patch, post, put } from './base'
 import type {
@ -50,140 +49,143 @@ export type SortType = 'created_at' | 'hit_count' | '-created_at' | '-hit_count'

 export type MetadataType = 'all' | 'only' | 'without'

-export const fetchDatasetDetail: Fetcher<DataSet, string> = (datasetId: string) => {
+export const fetchDatasetDetail = (datasetId: string): Promise<DataSet> => {
  return get<DataSet>(`/datasets/${datasetId}`)
 }

-export const updateDatasetSetting: Fetcher<DataSet, {
+export const updateDatasetSetting = ({
+  datasetId,
+  body,
+}: {
  datasetId: string
  body: Partial<Pick<DataSet,
    'name' | 'description' | 'permission' | 'partial_member_list' | 'indexing_technique' | 'retrieval_model' | 'embedding_model' | 'embedding_model_provider' | 'icon_info' | 'doc_form'
  >>
-}> = ({ datasetId, body }) => {
+}): Promise<DataSet> => {
  return patch<DataSet>(`/datasets/${datasetId}`, { body })
 }

-export const fetchDatasetRelatedApps: Fetcher<RelatedAppResponse, string> = (datasetId: string) => {
+export const fetchDatasetRelatedApps = (datasetId: string): Promise<RelatedAppResponse> => {
  return get<RelatedAppResponse>(`/datasets/${datasetId}/related-apps`)
 }

-export const fetchDatasets: Fetcher<DataSetListResponse, FetchDatasetsParams> = ({ url, params }) => {
+export const fetchDatasets = ({ url, params }: FetchDatasetsParams): Promise<DataSetListResponse> => {
  const urlParams = qs.stringify(params, { indices: false })
  return get<DataSetListResponse>(`${url}?${urlParams}`)
 }

-export const createEmptyDataset: Fetcher<DataSet, { name: string }> = ({ name }) => {
+export const createEmptyDataset = ({ name }: { name: string }): Promise<DataSet> => {
  return post<DataSet>('/datasets', { body: { name } })
 }

-export const checkIsUsedInApp: Fetcher<{ is_using: boolean }, string> = (id) => {
+export const checkIsUsedInApp = (id: string): Promise<{ is_using: boolean }> => {
  return get<{ is_using: boolean }>(`/datasets/${id}/use-check`, {}, {
    silent: true,
  })
 }

-export const deleteDataset: Fetcher<DataSet, string> = (datasetID) => {
+export const deleteDataset = (datasetID: string): Promise<DataSet> => {
  return del<DataSet>(`/datasets/${datasetID}`)
 }

-export const fetchExternalAPIList: Fetcher<ExternalAPIListResponse, { url: string }> = ({ url }) => {
+export const fetchExternalAPIList = ({ url }: { url: string }): Promise<ExternalAPIListResponse> => {
  return get<ExternalAPIListResponse>(url)
 }

-export const fetchExternalAPI: Fetcher<ExternalAPIItem, { apiTemplateId: string }> = ({ apiTemplateId }) => {
+export const fetchExternalAPI = ({ apiTemplateId }: { apiTemplateId: string }): Promise<ExternalAPIItem> => {
  return get<ExternalAPIItem>(`/datasets/external-knowledge-api/${apiTemplateId}`)
 }

-export const updateExternalAPI: Fetcher<ExternalAPIItem, { apiTemplateId: string; body: ExternalAPIItem }> = ({ apiTemplateId, body }) => {
+export const updateExternalAPI = ({ apiTemplateId, body }: { apiTemplateId: string; body: ExternalAPIItem }): Promise<ExternalAPIItem> => {
  return patch<ExternalAPIItem>(`/datasets/external-knowledge-api/${apiTemplateId}`, { body })
 }

-export const deleteExternalAPI: Fetcher<ExternalAPIDeleteResponse, { apiTemplateId: string }> = ({ apiTemplateId }) => {
+export const deleteExternalAPI = ({ apiTemplateId }: { apiTemplateId: string }): Promise<ExternalAPIDeleteResponse> => {
  return del<ExternalAPIDeleteResponse>(`/datasets/external-knowledge-api/${apiTemplateId}`)
 }

-export const checkUsageExternalAPI: Fetcher<ExternalAPIUsage, { apiTemplateId: string }> = ({ apiTemplateId }) => {
+export const checkUsageExternalAPI = ({ apiTemplateId }: { apiTemplateId: string }): Promise<ExternalAPIUsage> => {
  return get<ExternalAPIUsage>(`/datasets/external-knowledge-api/${apiTemplateId}/use-check`)
 }

-export const createExternalAPI: Fetcher<ExternalAPIItem, { body: CreateExternalAPIReq }> = ({ body }) => {
+export const createExternalAPI = ({ body }: { body: CreateExternalAPIReq }): Promise<ExternalAPIItem> => {
  return post<ExternalAPIItem>('/datasets/external-knowledge-api', { body })
 }

-export const createExternalKnowledgeBase: Fetcher<ExternalKnowledgeItem, { body: CreateKnowledgeBaseReq }> = ({ body }) => {
+export const createExternalKnowledgeBase = ({ body }: { body: CreateKnowledgeBaseReq }): Promise<ExternalKnowledgeItem> => {
  return post<ExternalKnowledgeItem>('/datasets/external', { body })
 }

-export const fetchDefaultProcessRule: Fetcher<ProcessRuleResponse, { url: string }> = ({ url }) => {
+export const fetchDefaultProcessRule = ({ url }: { url: string }): Promise<ProcessRuleResponse> => {
  return get<ProcessRuleResponse>(url)
 }
-export const fetchProcessRule: Fetcher<ProcessRuleResponse, { params: { documentId: string } }> = ({ params: { documentId } }) => {
+export const fetchProcessRule = ({ params: { documentId } }: { params: { documentId: string } }): Promise<ProcessRuleResponse> => {
  return get<ProcessRuleResponse>('/datasets/process-rule', { params: { document_id: documentId } })
 }

-export const createFirstDocument: Fetcher<createDocumentResponse, { body: CreateDocumentReq }> = ({ body }) => {
+export const createFirstDocument = ({ body }: { body: CreateDocumentReq }): Promise<createDocumentResponse> => {
  return post<createDocumentResponse>('/datasets/init', { body })
 }

-export const createDocument: Fetcher<createDocumentResponse, { datasetId: string; body: CreateDocumentReq }> = ({ datasetId, body }) => {
+export const createDocument = ({ datasetId, body }: { datasetId: string; body: CreateDocumentReq }): Promise<createDocumentResponse> => {
  return post<createDocumentResponse>(`/datasets/${datasetId}/documents`, { body })
 }

-export const fetchIndexingEstimate: Fetcher<IndexingEstimateResponse, CommonDocReq> = ({ datasetId, documentId }) => {
+export const fetchIndexingEstimate = ({ datasetId, documentId }: CommonDocReq): Promise<IndexingEstimateResponse> => {
  return get<IndexingEstimateResponse>(`/datasets/${datasetId}/documents/${documentId}/indexing-estimate`, {})
 }
-export const fetchIndexingEstimateBatch: Fetcher<IndexingEstimateResponse, BatchReq> = ({ datasetId, batchId }) => {
+export const fetchIndexingEstimateBatch = ({ datasetId, batchId }: BatchReq): Promise<IndexingEstimateResponse> => {
  return get<IndexingEstimateResponse>(`/datasets/${datasetId}/batch/${batchId}/indexing-estimate`, {})
 }

-export const fetchIndexingStatus: Fetcher<IndexingStatusResponse, CommonDocReq> = ({ datasetId, documentId }) => {
+export const fetchIndexingStatus = ({ datasetId, documentId }: CommonDocReq): Promise<IndexingStatusResponse> => {
  return get<IndexingStatusResponse>(`/datasets/${datasetId}/documents/${documentId}/indexing-status`, {})
 }

-export const fetchIndexingStatusBatch: Fetcher<IndexingStatusBatchResponse, BatchReq> = ({ datasetId, batchId }) => {
+export const fetchIndexingStatusBatch = ({ datasetId, batchId }: BatchReq): Promise<IndexingStatusBatchResponse> => {
  return get<IndexingStatusBatchResponse>(`/datasets/${datasetId}/batch/${batchId}/indexing-status`, {})
 }

-export const renameDocumentName: Fetcher<CommonResponse, CommonDocReq & { name: string }> = ({ datasetId, documentId, name }) => {
+export const renameDocumentName = ({ datasetId, documentId, name }: CommonDocReq & { name: string }): Promise<CommonResponse> => {
  return post<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/rename`, {
    body: { name },
  })
 }

-export const pauseDocIndexing: Fetcher<CommonResponse, CommonDocReq> = ({ datasetId, documentId }) => {
+export const pauseDocIndexing = ({ datasetId, documentId }: CommonDocReq): Promise<CommonResponse> => {
  return patch<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/processing/pause`)
 }

-export const resumeDocIndexing: Fetcher<CommonResponse, CommonDocReq> = ({ datasetId, documentId }) => {
+export const resumeDocIndexing = ({ datasetId, documentId }: CommonDocReq): Promise<CommonResponse> => {
  return patch<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/processing/resume`)
 }

-export const preImportNotionPages: Fetcher<{ notion_info: DataSourceNotionWorkspace[] }, { url: string; datasetId?: string }> = ({ url, datasetId }) => {
+export const preImportNotionPages = ({ url, datasetId }: { url: string; datasetId?: string }): Promise<{ notion_info: DataSourceNotionWorkspace[] }> => {
  return get<{ notion_info: DataSourceNotionWorkspace[] }>(url, { params: { dataset_id: datasetId } })
 }

-export const modifyDocMetadata: Fetcher<CommonResponse, CommonDocReq & { body: { doc_type: string; doc_metadata: Record<string, any> } }> = ({ datasetId, documentId, body }) => {
+export const modifyDocMetadata = ({ datasetId, documentId, body }: CommonDocReq & { body: { doc_type: string; doc_metadata: Record<string, any> } }): Promise<CommonResponse> => {
  return put<CommonResponse>(`/datasets/${datasetId}/documents/${documentId}/metadata`, { body })
 }

 // hit testing
-export const hitTesting: Fetcher<HitTestingResponse, { datasetId: string; queryText: string; retrieval_model: RetrievalConfig }> = ({ datasetId, queryText, retrieval_model }) => {
+export const hitTesting = ({ datasetId, queryText, retrieval_model }: { datasetId: string; queryText: string; retrieval_model: RetrievalConfig }): Promise<HitTestingResponse> => {
  return post<HitTestingResponse>(`/datasets/${datasetId}/hit-testing`, { body: { query: queryText, retrieval_model } })
 }

-export const externalKnowledgeBaseHitTesting: Fetcher<ExternalKnowledgeBaseHitTestingResponse, { datasetId: string; query: string; external_retrieval_model: { top_k: number; score_threshold: number; score_threshold_enabled: boolean } }> = ({ datasetId, query, external_retrieval_model }) => {
+export const externalKnowledgeBaseHitTesting = ({ datasetId, query, external_retrieval_model }: { datasetId: string; query: string; external_retrieval_model: { top_k: number; score_threshold: number; score_threshold_enabled: boolean } }): Promise<ExternalKnowledgeBaseHitTestingResponse> => {
  return post<ExternalKnowledgeBaseHitTestingResponse>(`/datasets/${datasetId}/external-hit-testing`, { body: { query, external_retrieval_model } })
 }

-export const fetchTestingRecords: Fetcher<HitTestingRecordsResponse, { datasetId: string; params: { page: number; limit: number } }> = ({ datasetId, params }) => {
+export const fetchTestingRecords = ({ datasetId, params }: { datasetId: string; params: { page: number; limit: number } }): Promise<HitTestingRecordsResponse> => {
  return get<HitTestingRecordsResponse>(`/datasets/${datasetId}/queries`, { params })
 }

-export const fetchFileIndexingEstimate: Fetcher<FileIndexingEstimateResponse, IndexingEstimateParams> = (body: IndexingEstimateParams) => {
+export const fetchFileIndexingEstimate = (body: IndexingEstimateParams): Promise<FileIndexingEstimateResponse> => {
  return post<FileIndexingEstimateResponse>('/datasets/indexing-estimate', { body })
 }

-export const fetchNotionPagePreview: Fetcher<{ content: string }, { workspaceID: string; pageID: string; pageType: string; credentialID: string; }> = ({ workspaceID, pageID, pageType, credentialID }) => {
+export const fetchNotionPagePreview = ({ workspaceID, pageID, pageType, credentialID }: { workspaceID: string; pageID: string; pageType: string; credentialID: string }): Promise<{ content: string }> => {
  return get<{ content: string }>(`notion/workspaces/${workspaceID}/pages/${pageID}/${pageType}/preview`, {
    params: {
      credential_id: credentialID,
@ -191,31 +193,31 @@ export const fetchNotionPagePreview: Fetcher<{ content: string }, { workspaceID:
  })
 }

-export const fetchApiKeysList: Fetcher<ApiKeysListResponse, { url: string; params: Record<string, any> }> = ({ url, params }) => {
+export const fetchApiKeysList = ({ url, params }: { url: string; params: Record<string, any> }): Promise<ApiKeysListResponse> => {
  return get<ApiKeysListResponse>(url, params)
 }

-export const delApikey: Fetcher<CommonResponse, { url: string; params: Record<string, any> }> = ({ url, params }) => {
+export const delApikey = ({ url, params }: { url: string; params: Record<string, any> }): Promise<CommonResponse> => {
  return del<CommonResponse>(url, params)
 }

-export const createApikey: Fetcher<CreateApiKeyResponse, { url: string; body: Record<string, any> }> = ({ url, body }) => {
+export const createApikey = ({ url, body }: { url: string; body: Record<string, any> }): Promise<CreateApiKeyResponse> => {
  return post<CreateApiKeyResponse>(url, body)
 }

-export const fetchDataSources = () => {
+export const fetchDataSources = (): Promise<CommonResponse> => {
  return get<CommonResponse>('api-key-auth/data-source')
 }

-export const createDataSourceApiKeyBinding: Fetcher<CommonResponse, Record<string, any>> = (body) => {
+export const createDataSourceApiKeyBinding = (body: Record<string, any>): Promise<CommonResponse> => {
  return post<CommonResponse>('api-key-auth/data-source/binding', { body })
 }

-export const removeDataSourceApiKeyBinding: Fetcher<CommonResponse, string> = (id: string) => {
+export const removeDataSourceApiKeyBinding = (id: string): Promise<CommonResponse> => {
  return del<CommonResponse>(`api-key-auth/data-source/${id}`)
 }

-export const createFirecrawlTask: Fetcher<CommonResponse, Record<string, any>> = (body) => {
+export const createFirecrawlTask = (body: Record<string, any>): Promise<CommonResponse> => {
  return post<CommonResponse>('website/crawl', {
    body: {
      ...body,
@ -224,7 +226,7 @@ export const createFirecrawlTask: Fetcher<CommonResponse, Record<string, any>> =
  })
 }

-export const checkFirecrawlTaskStatus: Fetcher<CommonResponse, string> = (jobId: string) => {
+export const checkFirecrawlTaskStatus = (jobId: string): Promise<CommonResponse> => {
  return get<CommonResponse>(`website/crawl/status/${jobId}`, {
    params: {
      provider: DataSourceProvider.fireCrawl,
@ -234,7 +236,7 @@ export const checkFirecrawlTaskStatus: Fetcher<CommonResponse, string> = (jobId:
  })
 }

-export const createJinaReaderTask: Fetcher<CommonResponse, Record<string, any>> = (body) => {
+export const createJinaReaderTask = (body: Record<string, any>): Promise<CommonResponse> => {
  return post<CommonResponse>('website/crawl', {
    body: {
      ...body,
@ -243,7 +245,7 @@ export const createJinaReaderTask: Fetcher<CommonResponse, Record<string, any>>
  })
 }

-export const checkJinaReaderTaskStatus: Fetcher<CommonResponse, string> = (jobId: string) => {
+export const checkJinaReaderTaskStatus = (jobId: string): Promise<CommonResponse> => {
  return get<CommonResponse>(`website/crawl/status/${jobId}`, {
    params: {
      provider: 'jinareader',
@ -253,7 +255,7 @@ export const checkJinaReaderTaskStatus: Fetcher<CommonResponse, string> = (jobId
  })
 }

-export const createWatercrawlTask: Fetcher<CommonResponse, Record<string, any>> = (body) => {
+export const createWatercrawlTask = (body: Record<string, any>): Promise<CommonResponse> => {
  return post<CommonResponse>('website/crawl', {
    body: {
      ...body,
@ -262,7 +264,7 @@ export const createWatercrawlTask: Fetcher<CommonResponse, Record<string, any>>
  })
 }

-export const checkWatercrawlTaskStatus: Fetcher<CommonResponse, string> = (jobId: string) => {
+export const checkWatercrawlTaskStatus = (jobId: string): Promise<CommonResponse> => {
  return get<CommonResponse>(`website/crawl/status/${jobId}`, {
    params: {
      provider: DataSourceProvider.waterCrawl,
@ -276,14 +278,14 @@ export type FileTypesRes = {
  allowed_extensions: string[]
 }

-export const fetchSupportFileTypes: Fetcher<FileTypesRes, { url: string }> = ({ url }) => {
+export const fetchSupportFileTypes = ({ url }: { url: string }): Promise<FileTypesRes> => {
  return get<FileTypesRes>(url)
 }

-export const getErrorDocs: Fetcher<ErrorDocsResponse, { datasetId: string }> = ({ datasetId }) => {
+export const getErrorDocs = ({ datasetId }: { datasetId: string }): Promise<ErrorDocsResponse> => {
  return get<ErrorDocsResponse>(`/datasets/${datasetId}/error-docs`)
 }

-export const retryErrorDocs: Fetcher<CommonResponse, { datasetId: string; document_ids: string[] }> = ({ datasetId, document_ids }) => {
+export const retryErrorDocs = ({ datasetId, document_ids }: { datasetId: string; document_ids: string[] }): Promise<CommonResponse> => {
  return post<CommonResponse>(`/datasets/${datasetId}/retry`, { body: { document_ids } })
 }
--- a/web/service/knowledge/use-dataset.ts
+++ b/web/service/knowledge/use-dataset.ts
@ -1,23 +1,86 @@
 import type { MutationOptions } from '@tanstack/react-query'
-import { useInfiniteQuery, useMutation, useQuery } from '@tanstack/react-query'
+import {
+  keepPreviousData,
+  useInfiniteQuery,
+  useMutation,
+  useQuery,
+  useQueryClient,
+} from '@tanstack/react-query'
+import qs from 'qs'
 import type {
  DataSet,
  DataSetListResponse,
  DatasetListRequest,
+  ErrorDocsResponse,
+  ExternalAPIListResponse,
+  FetchDatasetsParams,
+  HitTestingRecordsResponse,
  IndexingStatusBatchRequest,
  IndexingStatusBatchResponse,
  ProcessRuleResponse,
  RelatedAppResponse,
 } from '@/models/datasets'
+import type { ApiKeysListResponse } from '@/models/app'
 import { get, post } from '../base'
 import { useInvalid } from '../use-base'
-import qs from 'qs'
 import type { CommonResponse } from '@/models/common'

 const NAME_SPACE = 'dataset'

 const DatasetListKey = [NAME_SPACE, 'list']

+const normalizeDatasetsParams = (params: Partial<FetchDatasetsParams['params']> = {}) => {
+  const {
+    page = 1,
+    limit,
+    ids,
+    tag_ids,
+    include_all,
+    keyword,
+  } = params
+
+  return {
+    page,
+    ...(limit ? { limit } : {}),
+    ...(ids?.length ? { ids } : {}),
+    ...(tag_ids?.length ? { tag_ids } : {}),
+    ...(include_all !== undefined ? { include_all } : {}),
+    ...(keyword ? { keyword } : {}),
+  }
+}
+
+type UseInfiniteDatasetsOptions = {
+  enabled?: boolean
+  refetchOnMount?: boolean | 'always'
+  staleTime?: number
+  refetchOnReconnect?: boolean
+  refetchOnWindowFocus?: boolean
+}
+
+export const useInfiniteDatasets = (
+  params: Partial<FetchDatasetsParams['params']>,
+  options?: UseInfiniteDatasetsOptions,
+) => {
+  const normalizedParams = normalizeDatasetsParams(params)
+  const buildUrl = (pageParam: number | undefined) => {
+    const queryString = qs.stringify({
+      ...normalizedParams,
+      page: pageParam ?? normalizedParams.page,
+    }, { indices: false })
+    return `/datasets?${queryString}`
+  }
+
+  return useInfiniteQuery<DataSetListResponse>({
+    queryKey: [...DatasetListKey, 'infinite', normalizedParams],
+    queryFn: ({ pageParam = normalizedParams.page }) => get<DataSetListResponse>(buildUrl(pageParam as number | undefined)),
+    getNextPageParam: lastPage => lastPage.has_more ? lastPage.page + 1 : undefined,
+    initialPageParam: normalizedParams.page,
+    staleTime: 0,
+    refetchOnMount: 'always',
+    ...options,
+  })
+}
+
 export const useDatasetList = (params: DatasetListRequest) => {
  const { initialPage, tag_ids, limit, include_all, keyword } = params
  return useInfiniteQuery({
@ -70,10 +133,12 @@ export const useIndexingStatusBatch = (
  })
 }

-export const useProcessRule = (documentId: string) => {
+export const useProcessRule = (documentId?: string) => {
  return useQuery<ProcessRuleResponse>({
    queryKey: [NAME_SPACE, 'process-rule', documentId],
    queryFn: () => get<ProcessRuleResponse>('/datasets/process-rule', { params: { document_id: documentId } }),
+    enabled: !!documentId,
+    refetchOnWindowFocus: false,
  })
 }

@ -97,3 +162,57 @@ export const useDisableDatasetServiceApi = () => {
    mutationFn: (datasetId: string) => post<CommonResponse>(`/datasets/${datasetId}/api-keys/disable`),
  })
 }
+
+export const useDatasetApiKeys = (options?: { enabled?: boolean }) => {
+  return useQuery<ApiKeysListResponse>({
+    queryKey: [NAME_SPACE, 'api-keys'],
+    queryFn: () => get<ApiKeysListResponse>('/datasets/api-keys'),
+    enabled: options?.enabled ?? true,
+  })
+}
+
+export const useInvalidateDatasetApiKeys = () => {
+  const queryClient = useQueryClient()
+  return () => {
+    queryClient.invalidateQueries({
+      queryKey: [NAME_SPACE, 'api-keys'],
+    })
+  }
+}
+
+export const useExternalKnowledgeApiList = (options?: { enabled?: boolean }) => {
+  return useQuery<ExternalAPIListResponse>({
+    queryKey: [NAME_SPACE, 'external-knowledge-api'],
+    queryFn: () => get<ExternalAPIListResponse>('/datasets/external-knowledge-api'),
+    enabled: options?.enabled ?? true,
+  })
+}
+
+export const useInvalidateExternalKnowledgeApiList = () => {
+  const queryClient = useQueryClient()
+  return () => {
+    queryClient.invalidateQueries({
+      queryKey: [NAME_SPACE, 'external-knowledge-api'],
+    })
+  }
+}
+
+export const useDatasetTestingRecords = (
+  datasetId?: string,
+  params?: { page: number; limit: number },
+) => {
+  return useQuery<HitTestingRecordsResponse>({
+    queryKey: [NAME_SPACE, 'testing-records', datasetId, params],
+    queryFn: () => get<HitTestingRecordsResponse>(`/datasets/${datasetId}/queries`, { params }),
+    enabled: !!datasetId && !!params,
+    placeholderData: keepPreviousData,
+  })
+}
+
+export const useDatasetErrorDocs = (datasetId?: string) => {
+  return useQuery<ErrorDocsResponse>({
+    queryKey: [NAME_SPACE, 'error-docs', datasetId],
+    queryFn: () => get<ErrorDocsResponse>(`/datasets/${datasetId}/error-docs`),
+    enabled: !!datasetId,
+  })
+}