import { IDocumentInfo } from '@/interfaces/database/document'; import { IChunk } from '@/interfaces/database/knowledge'; import { IChangeParserConfigRequestBody } from '@/interfaces/request/document'; import i18n from '@/locales/config'; import chatService from '@/services/chat-service'; import kbService from '@/services/knowledge-service'; import { api_host } from '@/utils/api'; import { buildChunkHighlights } from '@/utils/document-util'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { UploadFile, message } from 'antd'; import { get } from 'lodash'; import { useCallback, useMemo, useState } from 'react'; import { IHighlight } from 'react-pdf-highlighter'; import { useGetPaginationWithRouter, useHandleSearchChange, } from './logic-hooks'; import { useGetKnowledgeSearchParams, useSetPaginationParams, } from './route-hook'; export const useGetDocumentUrl = (documentId?: string) => { const getDocumentUrl = useCallback( (id?: string) => { return `${api_host}/document/get/${documentId || id}`; }, [documentId], ); return getDocumentUrl; }; export const useGetChunkHighlights = (selectedChunk: IChunk) => { const [size, setSize] = useState({ width: 849, height: 1200 }); const highlights: IHighlight[] = useMemo(() => { return buildChunkHighlights(selectedChunk, size); }, [selectedChunk, size]); const setWidthAndHeight = (width: number, height: number) => { setSize((pre) => { if (pre.height !== height || pre.width !== width) { return { height, width }; } return pre; }); }; return { highlights, setWidthAndHeight }; }; export const useFetchNextDocumentList = () => { const { knowledgeId } = useGetKnowledgeSearchParams(); const { searchString, handleInputChange } = useHandleSearchChange(); const { pagination, setPagination } = useGetPaginationWithRouter(); const { data, isFetching: loading } = useQuery<{ docs: IDocumentInfo[]; total: number; }>({ queryKey: ['fetchDocumentList', searchString, pagination], initialData: { docs: [], total: 0 }, refetchInterval: 15000, queryFn: async () => { const ret = await kbService.get_document_list({ kb_id: knowledgeId, keywords: searchString, page_size: pagination.pageSize, page: pagination.current, }); if (ret.data.retcode === 0) { return ret.data.data; } return { docs: [], total: 0, }; }, }); const onInputChange: React.ChangeEventHandler = useCallback( (e) => { setPagination({ page: 1 }); handleInputChange(e); }, [handleInputChange, setPagination], ); return { loading, searchString, documents: data.docs, pagination: { ...pagination, total: data?.total }, handleInputChange: onInputChange, setPagination, }; }; export const useSetNextDocumentStatus = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['updateDocumentStatus'], mutationFn: async ({ status, documentId, }: { status: boolean; documentId: string; }) => { const { data } = await kbService.document_change_status({ doc_id: documentId, status: Number(status), }); if (data.retcode === 0) { message.success(i18n.t('message.modified')); queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); } return data; }, }); return { setDocumentStatus: mutateAsync, data, loading }; }; export const useSaveNextDocumentName = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['saveDocumentName'], mutationFn: async ({ name, documentId, }: { name: string; documentId: string; }) => { const { data } = await kbService.document_rename({ doc_id: documentId, name: name, }); if (data.retcode === 0) { message.success(i18n.t('message.renamed')); queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); } return data.retcode; }, }); return { loading, saveName: mutateAsync, data }; }; export const useCreateNextDocument = () => { const { knowledgeId } = useGetKnowledgeSearchParams(); const { setPaginationParams, page } = useSetPaginationParams(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['createDocument'], mutationFn: async (name: string) => { const { data } = await kbService.document_create({ name, kb_id: knowledgeId, }); if (data.retcode === 0) { if (page === 1) { queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); } else { setPaginationParams(); // fetch document list } message.success(i18n.t('message.created')); } return data.retcode; }, }); return { createDocument: mutateAsync, loading, data }; }; export const useSetNextDocumentParser = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['setDocumentParser'], mutationFn: async ({ parserId, documentId, parserConfig, }: { parserId: string; documentId: string; parserConfig: IChangeParserConfigRequestBody; }) => { const { data } = await kbService.document_change_parser({ parser_id: parserId, doc_id: documentId, parser_config: parserConfig, }); if (data.retcode === 0) { queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); message.success(i18n.t('message.modified')); } return data.retcode; }, }); return { setDocumentParser: mutateAsync, data, loading }; }; export const useUploadNextDocument = () => { const queryClient = useQueryClient(); const { knowledgeId } = useGetKnowledgeSearchParams(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['uploadDocument'], mutationFn: async (fileList: UploadFile[]) => { const formData = new FormData(); formData.append('kb_id', knowledgeId); fileList.forEach((file: any) => { formData.append('file', file); }); try { const ret = await kbService.document_upload(formData); const retcode = get(ret, 'data.retcode'); if (retcode === 0) { message.success(i18n.t('message.uploaded')); } if (retcode === 0 || retcode === 500) { queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); } return ret?.data; } catch (error) { console.warn(error); return {}; } }, }); return { uploadDocument: mutateAsync, loading, data }; }; export const useNextWebCrawl = () => { const { knowledgeId } = useGetKnowledgeSearchParams(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['webCrawl'], mutationFn: async ({ name, url }: { name: string; url: string }) => { const formData = new FormData(); formData.append('name', name); formData.append('url', url); formData.append('kb_id', knowledgeId); const ret = await kbService.web_crawl(formData); const retcode = get(ret, 'data.retcode'); if (retcode === 0) { message.success(i18n.t('message.uploaded')); } return retcode; }, }); return { data, loading, webCrawl: mutateAsync, }; }; export const useRunNextDocument = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['runDocumentByIds'], mutationFn: async ({ documentIds, run, }: { documentIds: string[]; run: number; }) => { const ret = await kbService.document_run({ doc_ids: documentIds, run, }); const retcode = get(ret, 'data.retcode'); if (retcode === 0) { queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); message.success(i18n.t('message.operated')); } return retcode; }, }); return { runDocumentByIds: mutateAsync, loading, data }; }; export const useFetchDocumentInfosByIds = () => { const [ids, setDocumentIds] = useState([]); const { data } = useQuery({ queryKey: ['fetchDocumentInfos', ids], enabled: ids.length > 0, initialData: [], queryFn: async () => { const { data } = await kbService.document_infos({ doc_ids: ids }); if (data.retcode === 0) { return data.data; } return []; }, }); return { data, setDocumentIds }; }; export const useFetchDocumentThumbnailsByIds = () => { const [ids, setDocumentIds] = useState([]); const { data } = useQuery>({ queryKey: ['fetchDocumentThumbnails', ids], enabled: ids.length > 0, initialData: {}, queryFn: async () => { const { data } = await kbService.document_thumbnails({ doc_ids: ids }); if (data.retcode === 0) { return data.data; } return {}; }, }); return { data, setDocumentIds }; }; export const useRemoveNextDocument = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['removeDocument'], mutationFn: async (documentIds: string | string[]) => { const { data } = await kbService.document_rm({ doc_id: documentIds }); if (data.retcode === 0) { message.success(i18n.t('message.deleted')); queryClient.invalidateQueries({ queryKey: ['fetchDocumentList'] }); } return data.retcode; }, }); return { data, loading, removeDocument: mutateAsync }; }; export const useDeleteDocument = () => { // const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['deleteDocument'], mutationFn: async (documentIds: string[]) => { const data = await kbService.document_delete({ doc_ids: documentIds }); // if (data.retcode === 0) { // queryClient.invalidateQueries({ queryKey: ['fetchFlowList'] }); // } return data; }, }); return { data, loading, deleteDocument: mutateAsync }; }; export const useUploadAndParseDocument = (uploadMethod: string) => { const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['uploadAndParseDocument'], mutationFn: async ({ conversationId, fileList, }: { conversationId: string; fileList: UploadFile[]; }) => { try { const formData = new FormData(); formData.append('conversation_id', conversationId); fileList.forEach((file: UploadFile) => { formData.append('file', file as any); }); if (uploadMethod === 'upload_and_parse') { const data = await kbService.upload_and_parse(formData); return data?.data; } const data = await chatService.uploadAndParseExternal(formData); return data?.data; } catch (error) { console.log('🚀 ~ useUploadAndParseDocument ~ error:', error); } }, }); return { data, loading, uploadAndParseDocument: mutateAsync }; };