import { IKnowledge, IKnowledgeResult, INextTestingResult, } from '@/interfaces/database/knowledge'; import { ITestRetrievalRequestBody } from '@/interfaces/request/knowledge'; import i18n from '@/locales/config'; import kbService, { listDataset } from '@/services/knowledge-service'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { useDebounce } from 'ahooks'; import { message } from 'antd'; import { useCallback, useMemo, useState } from 'react'; import { useParams } from 'umi'; import { useGetPaginationWithRouter, useHandleSearchChange, } from './logic-hooks'; import { useSetPaginationParams } from './route-hook'; export const enum KnowledgeApiAction { TestRetrieval = 'testRetrieval', FetchKnowledgeListByPage = 'fetchKnowledgeListByPage', CreateKnowledge = 'createKnowledge', DeleteKnowledge = 'deleteKnowledge', SaveKnowledge = 'saveKnowledge', FetchKnowledgeDetail = 'fetchKnowledgeDetail', } export const useKnowledgeBaseId = () => { const { id } = useParams(); return id; }; export const useTestRetrieval = () => { const knowledgeBaseId = useKnowledgeBaseId(); const { page, size: pageSize } = useSetPaginationParams(); const [values, setValues] = useState(); const queryParams = useMemo(() => { return { ...values, kb_id: values?.kb_id || knowledgeBaseId, page, size: pageSize, }; }, [knowledgeBaseId, page, pageSize, values]); const { data, isFetching: loading, refetch, } = useQuery({ queryKey: [KnowledgeApiAction.TestRetrieval, queryParams], initialData: { chunks: [], doc_aggs: [], total: 0, }, enabled: false, gcTime: 0, queryFn: async () => { const { data } = await kbService.retrieval_test(queryParams); console.log('🚀 ~ queryFn: ~ data:', data); return data?.data ?? {}; }, }); return { data, loading, setValues, refetch }; }; export const useFetchNextKnowledgeListByPage = () => { const { searchString, handleInputChange } = useHandleSearchChange(); const { pagination, setPagination } = useGetPaginationWithRouter(); const [ownerIds, setOwnerIds] = useState([]); const debouncedSearchString = useDebounce(searchString, { wait: 500 }); const { data, isFetching: loading } = useQuery({ queryKey: [ KnowledgeApiAction.FetchKnowledgeListByPage, { debouncedSearchString, ...pagination, ownerIds, }, ], initialData: { kbs: [], total: 0, }, gcTime: 0, queryFn: async () => { const { data } = await listDataset( { keywords: debouncedSearchString, page_size: pagination.pageSize, page: pagination.current, }, { owner_ids: ownerIds, }, ); return data?.data; }, }); const onInputChange: React.ChangeEventHandler = useCallback( (e) => { // setPagination({ page: 1 }); // TODO: 这里导致重复请求 handleInputChange(e); }, [handleInputChange], ); const handleOwnerIdsChange = useCallback((ids: string[]) => { // setPagination({ page: 1 }); // TODO: 这里导致重复请求 setOwnerIds(ids); }, []); return { ...data, searchString, handleInputChange: onInputChange, pagination: { ...pagination, total: data?.total }, setPagination, loading, setOwnerIds: handleOwnerIdsChange, ownerIds, }; }; export const useCreateKnowledge = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.CreateKnowledge], mutationFn: async (params: { id?: string; name: string }) => { 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({ kb_id: id }); if (data.code === 0) { message.success(i18n.t(`message.deleted`)); queryClient.invalidateQueries({ queryKey: [KnowledgeApiAction.FetchKnowledgeListByPage], }); } return data?.data ?? []; }, }); return { data, loading, deleteKnowledge: mutateAsync }; }; export const useUpdateKnowledge = (shouldFetchList = false) => { const knowledgeBaseId = useKnowledgeBaseId(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: [KnowledgeApiAction.SaveKnowledge], mutationFn: async (params: Record) => { const { data = {} } = await kbService.updateKb({ kb_id: params?.kb_id ? params?.kb_id : knowledgeBaseId, ...params, }); 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 = () => { const { id } = useParams(); const { data, isFetching: loading } = useQuery({ queryKey: [KnowledgeApiAction.FetchKnowledgeDetail], initialData: {} as IKnowledge, gcTime: 0, queryFn: async () => { const { data } = await kbService.get_kb_detail({ kb_id: id, }); return data?.data ?? {}; }, }); return { data, loading }; };