import { useHandleFilterSubmit } from '@/components/list-filter-bar/use-handle-filter-submit'; import message from '@/components/ui/message'; import { ParseType } from '@/constants/knowledge'; import { ResponsePostType, ResponseType } from '@/interfaces/database/base'; import { IArtifact, IArtifactGraph, IArtifactPage, IArtifactTopic, IDataset, IDatasetListResult, IKnowledgeGraph, INextTestingResult, IRenameTag, ITestingResult, IWikiCommit, IWikiCommitDetail, IWikiCommitListResponse, } from '@/interfaces/database/dataset'; import { IFetchArtifactGraphRequestParams, ITestRetrievalRequestBody, IUpdateArtifactPageRequestParams, } from '@/interfaces/request/knowledge'; import i18n from '@/locales/config'; import kbService, { clearWiki, deleteKnowledgeGraph, getArtifactGraph, getArtifactPage, getKbDetail, getKnowledgeGraph, getWikiCommit, listArtifactTopics, listArtifacts, listDataset, listTag, listWikiCommits, removeTag, renameTag, updateArtifactPage, updateKb, } from '@/services/knowledge-service'; import { useInfiniteQuery, useIsMutating, useMutation, useMutationState, useQuery, useQueryClient, } from '@tanstack/react-query'; import { useDebounce } from 'ahooks'; import { omit } from 'lodash'; import { useCallback, useMemo, useState } from 'react'; import { useParams, useSearchParams } from 'react-router'; import { useGetPaginationWithRouter, useHandleSearchChange, } from './logic-hooks'; import { extractParserConfigExt } from './parser-config-utils'; import { useSetPaginationParams } from './route-hook'; export const enum KnowledgeApiAction { FetchKnowledgeListByPage = 'fetchKnowledgeListByPage', CreateKnowledge = 'createKnowledge', DeleteKnowledge = 'deleteKnowledge', SaveKnowledge = 'saveKnowledge', FetchKnowledgeDetail = 'fetchKnowledgeDetail', FetchDatasetPipelineConfiguration = 'fetchDatasetPipelineConfiguration', FetchKnowledgeGraph = 'fetchKnowledgeGraph', FetchArtifactList = 'fetchArtifactList', FetchArtifactTopicList = 'fetchArtifactTopicList', FetchArtifactPage = 'fetchArtifactPage', FetchArtifactGraph = 'fetchArtifactGraph', UpdateArtifactPage = 'updateArtifactPage', FetchWikiCommits = 'fetchWikiCommits', FetchWikiCommit = 'fetchWikiCommit', FetchMetadata = 'fetchMetadata', FetchMetadataKeys = 'fetchMetadataKeys', FetchKnowledgeList = 'fetchKnowledgeList', RemoveKnowledgeGraph = 'removeKnowledgeGraph', ClearWiki = 'clearWiki', } export const useKnowledgeBaseId = (): string => { const { id } = useParams(); return (id as string) || ''; }; export const useTestRetrieval = () => { const knowledgeBaseId = useKnowledgeBaseId(); const [values, setValues] = useState(); const { filterValue, setFilterValue } = useHandleFilterSubmit(); const queryParams = useMemo(() => { return { ...values, kb_id: values?.kb_id || knowledgeBaseId, page: 1, doc_ids: filterValue.doc_ids, highlight: true, }; }, [filterValue, knowledgeBaseId, values]); const mutation = useMutation({ mutationFn: async (params) => { const { data } = await kbService.retrievalTest(params); const result = data?.data ?? {}; return { ...result, isRuned: true }; }, }); const refetch = useCallback(() => { if (queryParams.question) { const newParams = { ...queryParams, page: 1 }; mutation.mutate(newParams); } }, [mutation, queryParams]); const handleFilterSubmit = useCallback( (value: { doc_ids?: string[] }) => { setFilterValue(value); if (mutation.data && queryParams.question) { const newParams = { ...queryParams, doc_ids: value.doc_ids ?? [], page: 1, }; mutation.mutate(newParams); } }, [mutation, queryParams, setFilterValue], ); const data = useMemo( () => mutation.data ?? { chunks: [], doc_aggs: [], total: 0, isRuned: false, }, [mutation.data], ); return { data, loading: mutation.isPending, setValues, refetch, handleFilterSubmit, filterValue, }; }; export const useFetchNextKnowledgeListByPage = () => { const { searchString, handleInputChange } = useHandleSearchChange(); const { pagination, setPagination } = useGetPaginationWithRouter(); const debouncedSearchString = useDebounce(searchString, { wait: 500 }); const { filterValue, handleFilterSubmit } = useHandleFilterSubmit(); const { data, isFetching: loading } = useQuery({ queryKey: [ KnowledgeApiAction.FetchKnowledgeListByPage, { debouncedSearchString, ...pagination, filterValue, }, ], initialData: { kbs: [], total_datasets: 0, }, gcTime: 0, queryFn: async () => { const { data } = await listDataset({ page_size: pagination.pageSize, page: pagination.current, ext: { keywords: debouncedSearchString, owner_ids: filterValue.owner as string[], }, }); return { kbs: data?.data, total_datasets: data?.total_datasets }; }, }); const onInputChange: React.ChangeEventHandler = useCallback( (e) => { // setPagination({ page: 1 }); // TODO: This results in repeated requests handleInputChange(e); }, [handleInputChange], ); return { ...data, searchString, handleInputChange: onInputChange, pagination: { ...pagination, total: data?.total_datasets }, setPagination, loading, filterValue, handleFilterSubmit, }; }; export const useCreateKnowledge = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.CreateKnowledge], mutationFn: async (params: { id?: string; name: string; embedding_model?: string; chunk_method?: string; parseType?: ParseType; pipeline_id?: string | null; ext?: { language?: string; [key: string]: any; }; }) => { const { data = {} } = await kbService.createKb(params); if (data.code === 0) { message.success( i18n.t(`message.${params?.id ? 'modified' : 'created'}`), ); queryClient.invalidateQueries({ queryKey: ['fetchKnowledgeList'] }); } return data; }, }); return { data, loading, createKnowledge: mutateAsync }; }; export const useDeleteKnowledge = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.DeleteKnowledge], mutationFn: async (id: string) => { const { data } = await kbService.rmKb({ ids: [id] }); if (data.code === 0) { message.success(i18n.t(`message.deleted`)); queryClient.invalidateQueries({ queryKey: [KnowledgeApiAction.FetchKnowledgeListByPage], }); } return data?.data ?? []; }, }); return { data, loading, deleteKnowledge: mutateAsync }; }; function isPipelineParserConfig( parserConfig: Record | undefined, ): boolean { if (!parserConfig || typeof parserConfig !== 'object') { return false; } return Object.keys(parserConfig).some((key) => key.includes(':')); } export const useUpdateKnowledge = (shouldFetchList = false) => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.SaveKnowledge], mutationFn: async (params: { kb_id?: string; name?: string; embedding_model?: string; chunk_method?: string; pipeline_id?: string | null; avatar?: string | null; description?: string; permission?: string; pagerank?: number; parser_config?: Record; [key: string]: any; }) => { const kbId = params?.kb_id || knowledgeBaseId; const { embedding_model, chunk_method, pipeline_id, avatar, description, permission, pagerank, parser_config, ...ext } = params; const requestBody: Record = { name, embedding_model, chunk_method, pipeline_id, avatar, description, permission, pagerank, parser_config: isPipelineParserConfig(parser_config) ? parser_config : extractParserConfigExt(parser_config), ...omit(ext, ['kb_id']), }; const { data = {} } = await updateKb(kbId, requestBody); if (data.code === 0) { message.success(i18n.t(`message.updated`)); if (shouldFetchList) { queryClient.invalidateQueries({ queryKey: [KnowledgeApiAction.FetchKnowledgeListByPage], }); } else { queryClient.invalidateQueries({ queryKey: ['fetchKnowledgeDetail'] }); } } return data; }, }); return { data, loading, saveKnowledgeConfiguration: mutateAsync }; }; export const useFetchKnowledgeBaseConfiguration = (props?: { isEdit?: boolean; }) => { const { isEdit = true } = props || { isEdit: true }; const { id } = useParams(); const [searchParams] = useSearchParams(); const knowledgeBaseId = searchParams.get('id') || id; const { data, isFetching: loading } = useQuery({ queryKey: [KnowledgeApiAction.FetchKnowledgeDetail, knowledgeBaseId], initialData: {} as IDataset, gcTime: 0, enabled: !!knowledgeBaseId && isEdit, queryFn: async () => { const { data } = await getKbDetail(knowledgeBaseId || ''); return data?.data ?? {}; }, }); return { data, loading }; }; export const DatasetPipelineConfigurationKeys = { detail: (knowledgeBaseId: string | null | undefined) => [ KnowledgeApiAction.FetchDatasetPipelineConfiguration, knowledgeBaseId, ] as const, }; export const useFetchDatasetPipelineConfiguration = (props?: { isEdit?: boolean; }) => { const { isEdit = true } = props || { isEdit: true }; const { id } = useParams(); const [searchParams] = useSearchParams(); const knowledgeBaseId = searchParams.get('id') || id; const { data, isFetching: loading } = useQuery({ queryKey: DatasetPipelineConfigurationKeys.detail(knowledgeBaseId), initialData: {} as IDataset, gcTime: 0, enabled: !!knowledgeBaseId && isEdit, queryFn: async () => { const { data } = await getKbDetail(knowledgeBaseId || ''); return data?.data ?? {}; }, }); return { data, loading }; }; export const ArtifactKeys = { list: ( datasetId: string, keywords: string, topic?: string, pageType?: string, ) => [ KnowledgeApiAction.FetchArtifactList, datasetId, keywords, topic, pageType, ] as const, listByDataset: (datasetId: string) => [KnowledgeApiAction.FetchArtifactList, datasetId] as const, detail: (datasetId: string, pageType: string, slug: string) => [KnowledgeApiAction.FetchArtifactPage, datasetId, pageType, slug] as const, }; export const ArtifactTopicKeys = { list: (datasetId: string, keywords: string) => [KnowledgeApiAction.FetchArtifactTopicList, datasetId, keywords] as const, listByDataset: (datasetId: string) => [KnowledgeApiAction.FetchArtifactTopicList, datasetId] as const, }; const wikiCommitKeys = { list: (datasetId: string, pageType: string, slug: string) => [KnowledgeApiAction.FetchWikiCommits, datasetId, pageType, slug] as const, detail: (datasetId: string, commitId: string) => [KnowledgeApiAction.FetchWikiCommit, datasetId, commitId] as const, }; export const useFetchWikiCommits = ( artifact: IArtifact | null, enabled = true, ) => { const knowledgeBaseId = useKnowledgeBaseId(); const pageType = artifact?.page_type ?? ''; const slug = artifact?.slug ?? ''; const { data, isFetching: loading } = useQuery({ queryKey: wikiCommitKeys.list(knowledgeBaseId, pageType, slug), enabled: !!knowledgeBaseId && !!artifact && !!pageType && !!slug && enabled, gcTime: 0, queryFn: async () => { const { data } = await listWikiCommits(knowledgeBaseId, pageType, slug); // The merged file-commit endpoint returns {total, page, page_size, commits}, // while the existing components expect {total, items}. Normalize here. const raw = (data?.data ?? {}) as { total?: number; items?: IWikiCommit[]; commits?: IWikiCommit[]; }; return { total: raw.total ?? 0, items: raw.items ?? raw.commits ?? [], }; }, }); return { commits: data?.items ?? [], total: data?.total ?? 0, loading, }; }; export function useFetchWikiCommit(commitId: string | null) { const knowledgeBaseId = useKnowledgeBaseId(); const { data, isFetching: loading } = useQuery({ queryKey: wikiCommitKeys.detail(knowledgeBaseId, commitId ?? ''), enabled: !!knowledgeBaseId && !!commitId, gcTime: 0, queryFn: async () => { const { data } = await getWikiCommit(knowledgeBaseId, commitId!); return data?.data ?? null; }, }); return { data, loading }; } type UseFetchArtifactListOptions = { keywords?: string; topic?: string; pageType?: string; enabled?: boolean; }; export const useFetchArtifactList = ( options: UseFetchArtifactListOptions = {}, ) => { const { keywords = '', topic, pageType, enabled = true } = options; const knowledgeBaseId = useKnowledgeBaseId(); const { data, fetchNextPage, hasNextPage, isFetching, isFetchingNextPage } = useInfiniteQuery<{ artifacts: IArtifact[]; total: number; }>({ queryKey: ArtifactKeys.list(knowledgeBaseId, keywords, topic, pageType), enabled: !!knowledgeBaseId && enabled && !!topic, gcTime: 0, initialPageParam: 1, queryFn: async ({ pageParam }) => { const page = pageParam as number; const { data } = await listArtifacts(knowledgeBaseId, { page, page_size: 30, keywords, topic, page_type: pageType, }); const responseData = data?.data; return { artifacts: responseData?.items ?? [], total: responseData?.total ?? 0, }; }, getNextPageParam: (lastPage, allPages) => { const loadedCount = allPages.reduce( (sum, page) => sum + page.artifacts.length, 0, ); return loadedCount < lastPage.total ? allPages.length + 1 : undefined; }, }); const artifacts = useMemo( () => data?.pages.flatMap((page) => page.artifacts) ?? [], [data], ); const loading = isFetching || isFetchingNextPage; const handleScroll = useCallback( (e: React.UIEvent) => { const { scrollTop, scrollHeight, clientHeight } = e.currentTarget; const threshold = 50; if ( scrollHeight - scrollTop - clientHeight <= threshold && hasNextPage && !isFetchingNextPage ) { fetchNextPage(); } }, [fetchNextPage, hasNextPage, isFetchingNextPage], ); return { artifacts, loading, handleScroll, hasMore: !!hasNextPage, }; }; type UseFetchArtifactTopicListOptions = { keywords?: string; enabled?: boolean; }; export const useFetchArtifactTopicList = ( options: UseFetchArtifactTopicListOptions = {}, ) => { const { keywords = '', enabled = true } = options; const knowledgeBaseId = useKnowledgeBaseId(); const { data, fetchNextPage, hasNextPage, isFetching, isFetchingNextPage } = useInfiniteQuery<{ topics: IArtifactTopic[]; total: number; }>({ queryKey: ArtifactTopicKeys.list(knowledgeBaseId, keywords), enabled: !!knowledgeBaseId && enabled, gcTime: 0, initialPageParam: 1, queryFn: async ({ pageParam }) => { const page = pageParam as number; const { data } = await listArtifactTopics(knowledgeBaseId, { page, page_size: 30, keywords, }); const responseData = data?.data; return { topics: responseData?.items ?? [], total: responseData?.total ?? 0, }; }, getNextPageParam: (lastPage, allPages) => { const loadedCount = allPages.reduce( (sum, page) => sum + page.topics.length, 0, ); return loadedCount < lastPage.total ? allPages.length + 1 : undefined; }, }); const topics = useMemo( () => data?.pages.flatMap((page) => page.topics) ?? [], [data], ); const loading = isFetching || isFetchingNextPage; const handleScroll = useCallback( (e: React.UIEvent) => { const { scrollTop, scrollHeight, clientHeight } = e.currentTarget; const threshold = 50; if ( scrollHeight - scrollTop - clientHeight <= threshold && hasNextPage && !isFetchingNextPage ) { fetchNextPage(); } }, [fetchNextPage, hasNextPage, isFetchingNextPage], ); return { topics, loading, handleScroll, hasMore: !!hasNextPage, }; }; export function useFetchArtifactPage( artifact: IArtifact | null, enabled = true, ) { const knowledgeBaseId = useKnowledgeBaseId(); const pageType = artifact?.page_type ?? ''; const slug = artifact?.slug ?? ''; const { data, isFetching: loading } = useQuery({ queryKey: ArtifactKeys.detail(knowledgeBaseId, pageType, slug), enabled: !!knowledgeBaseId && !!artifact && !!pageType && !!slug && enabled, queryFn: async () => { const { data } = await getArtifactPage(knowledgeBaseId, pageType, slug); return data?.data ?? null; }, }); return { data, loading }; } export const useUpdateArtifactPage = () => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation< ResponseType, Error, IUpdateArtifactPageRequestParams >({ mutationKey: [KnowledgeApiAction.UpdateArtifactPage], mutationFn: async (params) => { const { data = {} } = await updateArtifactPage( knowledgeBaseId, params.pageType, params.slug, params.body, ); if (data.code === 0) { message.success(i18n.t(`message.updated`)); queryClient.invalidateQueries({ queryKey: ArtifactKeys.detail( knowledgeBaseId, params.pageType, params.slug, ), }); } return data; }, }); return { data, loading, updateArtifactPage: mutateAsync }; }; export function useFetchKnowledgeGraph() { const knowledgeBaseId = useKnowledgeBaseId(); const { data, isFetching: loading } = useQuery({ queryKey: [KnowledgeApiAction.FetchKnowledgeGraph, knowledgeBaseId], initialData: { graph: {}, mind_map: {} } as IKnowledgeGraph, enabled: !!knowledgeBaseId, gcTime: 0, queryFn: async () => { const { data } = await getKnowledgeGraph(knowledgeBaseId); return data?.data; }, }); return { data, loading }; } export const artifactGraphKeys = { graph: (datasetId: string, params?: IFetchArtifactGraphRequestParams) => [KnowledgeApiAction.FetchArtifactGraph, datasetId, params?.node] as const, }; export function useFetchArtifactGraph( params?: IFetchArtifactGraphRequestParams, options?: { enabled?: boolean }, ) { const knowledgeBaseId = useKnowledgeBaseId(); const { data, isFetching: loading } = useQuery({ queryKey: artifactGraphKeys.graph(knowledgeBaseId, params), initialData: { entities: [], relations: [] } as IArtifactGraph, enabled: !!knowledgeBaseId && (options?.enabled ?? true), gcTime: 0, queryFn: async () => { const { data } = await getArtifactGraph(knowledgeBaseId, params); return data?.data ?? { entities: [], relations: [] }; }, }); return { data, loading }; } export function useFetchKnowledgeMetadata(kbIds: string[] = []) { const { data, isFetching: loading } = useQuery< Record> >({ queryKey: [KnowledgeApiAction.FetchMetadata, kbIds], initialData: {}, enabled: kbIds.length > 0, gcTime: 0, queryFn: async () => { const { data } = await kbService.getMeta({ dataset_ids: kbIds.join(','), }); return data?.data ?? {}; }, }); return { data, loading }; } export function useFetchKnowledgeMetadataKeys(kbIds: string[] = []) { const sortedKbIds = useMemo(() => [...kbIds].sort(), [kbIds]); const { data, isFetching: loading } = useQuery({ queryKey: [KnowledgeApiAction.FetchMetadataKeys, sortedKbIds], initialData: [], enabled: sortedKbIds.length > 0, gcTime: 0, queryFn: async () => { const { data } = await kbService.getMetaKeys({ kb_ids: sortedKbIds.join(','), }); return data?.data ?? []; }, }); return { data, loading }; } export const useRemoveKnowledgeGraph = () => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.RemoveKnowledgeGraph], mutationFn: async () => { const { data } = await deleteKnowledgeGraph(knowledgeBaseId); if (data.code === 0) { message.success(i18n.t(`message.deleted`)); queryClient.invalidateQueries({ queryKey: [KnowledgeApiAction.FetchKnowledgeGraph], }); } return data?.code; }, }); return { data, loading, removeKnowledgeGraph: mutateAsync }; }; export const useClearWiki = () => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.ClearWiki], mutationFn: async () => { const { data } = await clearWiki(knowledgeBaseId); if (data?.code === 0) { message.success(i18n.t('message.deleted')); queryClient.invalidateQueries({ queryKey: ArtifactKeys.listByDataset(knowledgeBaseId), }); queryClient.invalidateQueries({ queryKey: ArtifactTopicKeys.listByDataset(knowledgeBaseId), }); queryClient.invalidateQueries({ queryKey: artifactGraphKeys.graph(knowledgeBaseId), }); } return data; }, }); return { data, loading, clearWiki: mutateAsync }; }; export const useFetchKnowledgeList = ( shouldFilterListWithoutDocument: boolean = false, keywords = '', ): { list: IDataset[]; loading: boolean; } => { const { data, isFetching: loading } = useQuery({ queryKey: [ KnowledgeApiAction.FetchKnowledgeList, shouldFilterListWithoutDocument, keywords, ], initialData: [], gcTime: 0, // https://tanstack.com/query/latest/docs/framework/react/guides/caching?from=reactQueryV3 queryFn: async () => { const { data } = await listDataset( keywords ? { ext: { keywords, }, } : undefined, ); const list = data?.data ?? []; return shouldFilterListWithoutDocument ? list.filter((x: IDataset) => x.chunk_count > 0) : list; }, }); return { list: data, loading }; }; export const useSelectKnowledgeOptions = () => { const { list } = useFetchKnowledgeList(); const options = list?.map((item) => ({ label: item.name, value: item.id, })); return options; }; //#region tags export const useRenameTag = () => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['renameTag'], mutationFn: async (params: IRenameTag) => { const { data } = await renameTag(knowledgeBaseId, params); if (data.code === 0) { message.success(i18n.t(`message.modified`)); queryClient.invalidateQueries({ queryKey: ['fetchTagList'], }); } return data?.data ?? []; }, }); return { data, loading, renameTag: mutateAsync }; }; export const useTagIsRenaming = () => { return useIsMutating({ mutationKey: ['renameTag'] }) > 0; }; export const useFetchTagListByKnowledgeIds = () => { const [knowledgeIds, setKnowledgeIds] = useState([]); const { data, isFetching: loading } = useQuery>({ queryKey: ['fetchTagListByKnowledgeIds', knowledgeIds], enabled: knowledgeIds.length > 0, initialData: [], gcTime: 0, // https://tanstack.com/query/latest/docs/framework/react/guides/caching?from=reactQueryV3 queryFn: async () => { const { data } = await kbService.listTagByKnowledgeIds({ dataset_ids: knowledgeIds.join(','), }); const list = (data?.data || []) as Array< [string, number] | { value?: string; count?: number } >; return list.flatMap((tag): Array<[string, number]> => { if (Array.isArray(tag)) { return [tag]; } return tag.value ? [[tag.value, tag.count ?? 0]] : []; }); }, }); return { list: data, loading, setKnowledgeIds }; }; export const useFetchTagList = () => { const knowledgeBaseId = useKnowledgeBaseId(); const { data, isFetching: loading } = useQuery>({ queryKey: ['fetchTagList'], initialData: [], gcTime: 0, // https://tanstack.com/query/latest/docs/framework/react/guides/caching?from=reactQueryV3 queryFn: async () => { const { data } = await listTag(knowledgeBaseId); const list = data?.data || []; return list; }, }); return { list: data, loading }; }; export const useDeleteTag = () => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['deleteTag'], mutationFn: async (tags: string[]) => { const { data } = await removeTag(knowledgeBaseId, tags); if (data.code === 0) { message.success(i18n.t(`message.deleted`)); queryClient.invalidateQueries({ queryKey: ['fetchTagList'], }); } return data?.data ?? []; }, }); return { data, loading, deleteTag: mutateAsync }; }; // #endregion //#region Retrieval testing export const useTestChunkRetrieval = (): ResponsePostType & { testChunk: (...params: any[]) => void; } => { const knowledgeBaseId = useKnowledgeBaseId(); const { page, size: pageSize } = useSetPaginationParams(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['testChunk'], // This method is invalid gcTime: 0, mutationFn: async (values: any) => { const { data } = await kbService.retrievalTest({ ...values, kb_id: values.kb_id ?? knowledgeBaseId, page, size: pageSize, }); if (data.code === 0) { const res = data.data; return { ...res, documents: res.doc_aggs, }; } return ( data?.data ?? { chunks: [], documents: [], total: 0, } ); }, }); return { data: data ?? { chunks: [], documents: [], total: 0 }, loading, testChunk: mutateAsync, }; }; export const useTestChunkAllRetrieval = (): ResponsePostType & { testChunkAll: (...params: any[]) => void; } => { const knowledgeBaseId = useKnowledgeBaseId(); const { page, size: pageSize } = useSetPaginationParams(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['testChunkAll'], // This method is invalid gcTime: 0, mutationFn: async (values: any) => { const { data } = await kbService.retrievalTest({ ...values, kb_id: values.kb_id ?? knowledgeBaseId, doc_ids: [], page, size: pageSize, }); if (data.code === 0) { const res = data.data; return { ...res, documents: res.doc_aggs, }; } return ( data?.data ?? { chunks: [], documents: [], total: 0, } ); }, }); return { data: data ?? { chunks: [], documents: [], total: 0 }, loading, testChunkAll: mutateAsync, }; }; export const useChunkIsTesting = () => { return useIsMutating({ mutationKey: ['testChunk'] }) > 0; }; export const useSelectTestingResult = (): ITestingResult => { const data = useMutationState({ filters: { mutationKey: ['testChunk'] }, select: (mutation) => { return mutation.state.data; }, }); return (data.at(-1) ?? { chunks: [], documents: [], total: 0, }) as ITestingResult; }; export const useSelectIsTestingSuccess = () => { const status = useMutationState({ filters: { mutationKey: ['testChunk'] }, select: (mutation) => { return mutation.state.status; }, }); return status.at(-1) === 'success'; }; export const useAllTestingSuccess = () => { const status = useMutationState({ filters: { mutationKey: ['testChunkAll'] }, select: (mutation) => { return mutation.state.status; }, }); return status.at(-1) === 'success'; }; export const useAllTestingResult = (): ITestingResult => { const data = useMutationState({ filters: { mutationKey: ['testChunkAll'] }, select: (mutation) => { return mutation.state.data; }, }); return (data.at(-1) ?? { chunks: [], documents: [], total: 0, }) as ITestingResult; }; //#endregion