import type { MutationOptions } from '@tanstack/react-query' import { useInfiniteQuery, useMutation, useQuery } from '@tanstack/react-query' import type { DataSet, DataSetListResponse, DatasetListRequest, IndexingStatusBatchRequest, IndexingStatusBatchResponse, ProcessRuleResponse, RelatedAppResponse, } from '@/models/datasets' import { get } from '../base' import { useReset } from '../use-base' import qs from 'qs' const NAME_SPACE = 'dataset' const DatasetListKey = [NAME_SPACE, 'list'] export const useDatasetList = (params: DatasetListRequest) => { const { initialPage, tag_ids, limit, include_all, keyword } = params return useInfiniteQuery({ queryKey: [...DatasetListKey, initialPage, tag_ids, limit, include_all, keyword], queryFn: ({ pageParam = 1 }) => { const urlParams = qs.stringify({ tag_ids, limit, include_all, keyword, page: pageParam, }, { indices: false }) return get(`/datasets?${urlParams}`) }, getNextPageParam: lastPage => lastPage.has_more ? lastPage.page + 1 : null, initialPageParam: initialPage, }) } export const useResetDatasetList = () => { return useReset([...DatasetListKey]) } export const useDatasetDetail = (datasetId: string) => { return useQuery({ queryKey: [NAME_SPACE, 'detail', datasetId], queryFn: () => get(`/datasets/${datasetId}`), enabled: !!datasetId, }) } export const useDatasetRelatedApps = (datasetId: string) => { return useQuery({ queryKey: [NAME_SPACE, 'related-apps', datasetId], queryFn: () => get(`/datasets/${datasetId}/related-apps`), }) } export const useIndexingStatusBatch = ( params: IndexingStatusBatchRequest, mutationOptions: MutationOptions = {}, ) => { const { datasetId, batchId } = params return useMutation({ mutationKey: [NAME_SPACE, 'indexing-status-batch', datasetId, batchId], mutationFn: () => get(`/datasets/${datasetId}/batch/${batchId}/indexing-status`), ...mutationOptions, }) } export const useProcessRule = (documentId: string) => { return useQuery({ queryKey: [NAME_SPACE, 'process-rule', documentId], queryFn: () => get('/datasets/process-rule', { params: { document_id: documentId } }), }) }