import { LanguageTranslationMap } from '@/constants/common'; import { ResponseGetType } from '@/interfaces/database/base'; import { IToken } from '@/interfaces/database/chat'; import { ITenantInfo } from '@/interfaces/database/knowledge'; import { ISystemStatus, ITenant, ITenantUser, IUserInfo, } from '@/interfaces/database/user-setting'; import userService, { addTenantUser, agreeTenant, deleteTenantUser, listTenant, listTenantUser, } from '@/services/user-service'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { Modal, message } from 'antd'; import DOMPurify from 'dompurify'; import { isEmpty } from 'lodash'; import { useCallback, useMemo, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { history } from 'umi'; export const useFetchUserInfo = (): ResponseGetType => { const { i18n } = useTranslation(); const { data, isFetching: loading } = useQuery({ queryKey: ['userInfo'], initialData: {}, gcTime: 0, queryFn: async () => { const { data } = await userService.user_info(); if (data.code === 0) { i18n.changeLanguage( LanguageTranslationMap[ data.data.language as keyof typeof LanguageTranslationMap ], ); } return data?.data ?? {}; }, }); return { data, loading }; }; export const useFetchTenantInfo = (): ResponseGetType => { const { t } = useTranslation(); const { data, isFetching: loading } = useQuery({ queryKey: ['tenantInfo'], initialData: {}, gcTime: 0, queryFn: async () => { const { data: res } = await userService.get_tenant_info(); if (res.code === 0) { // llm_id is chat_id // asr_id is speech2txt const { data } = res; if (isEmpty(data.embd_id) || isEmpty(data.llm_id)) { Modal.warning({ title: t('common.warn'), content: (
), onOk() { history.push('/user-setting/model'); }, }); } data.chat_id = data.llm_id; data.speech2text_id = data.asr_id; return data; } return res; }, }); return { data, loading }; }; export const useSelectParserList = (): Array<{ value: string; label: string; }> => { const { data: tenantInfo } = useFetchTenantInfo(); const parserList = useMemo(() => { const parserArray: Array = tenantInfo?.parser_ids?.split(',') ?? []; return parserArray.map((x) => { const arr = x.split(':'); return { value: arr[0], label: arr[1] }; }); }, [tenantInfo]); return parserList; }; export const useSaveSetting = () => { const queryClient = useQueryClient(); const { t } = useTranslation(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['saveSetting'], mutationFn: async ( userInfo: { new_password: string } | Partial, ) => { const { data } = await userService.setting(userInfo); if (data.code === 0) { message.success(t('message.modified')); queryClient.invalidateQueries({ queryKey: ['userInfo'] }); } return data?.code; }, }); return { data, loading, saveSetting: mutateAsync }; }; export const useFetchSystemVersion = () => { const [version, setVersion] = useState(''); const [loading, setLoading] = useState(false); const fetchSystemVersion = useCallback(async () => { try { setLoading(true); const { data } = await userService.getSystemVersion(); if (data.code === 0) { setVersion(data.data); setLoading(false); } } catch (error) { setLoading(false); } }, []); return { fetchSystemVersion, version, loading }; }; export const useFetchSystemStatus = () => { const [systemStatus, setSystemStatus] = useState( {} as ISystemStatus, ); const [loading, setLoading] = useState(false); const fetchSystemStatus = useCallback(async () => { setLoading(true); const { data } = await userService.getSystemStatus(); if (data.code === 0) { setSystemStatus(data.data); setLoading(false); } }, []); return { systemStatus, fetchSystemStatus, loading, }; }; export const useFetchManualSystemTokenList = () => { const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['fetchManualSystemTokenList'], mutationFn: async () => { const { data } = await userService.listToken(); return data?.data ?? []; }, }); return { data, loading, fetchSystemTokenList: mutateAsync }; }; export const useFetchSystemTokenList = () => { const { data, isFetching: loading, refetch, } = useQuery({ queryKey: ['fetchSystemTokenList'], initialData: [], gcTime: 0, queryFn: async () => { const { data } = await userService.listToken(); return data?.data ?? []; }, }); return { data, loading, refetch }; }; export const useRemoveSystemToken = () => { const queryClient = useQueryClient(); const { t } = useTranslation(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['removeSystemToken'], mutationFn: async (token: string) => { const { data } = await userService.removeToken({}, token); if (data.code === 0) { message.success(t('message.deleted')); queryClient.invalidateQueries({ queryKey: ['fetchSystemTokenList'] }); } return data?.data ?? []; }, }); return { data, loading, removeToken: mutateAsync }; }; export const useCreateSystemToken = () => { const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['createSystemToken'], mutationFn: async (params: Record) => { const { data } = await userService.createToken(params); if (data.code === 0) { queryClient.invalidateQueries({ queryKey: ['fetchSystemTokenList'] }); } return data?.data ?? []; }, }); return { data, loading, createToken: mutateAsync }; }; export const useListTenantUser = () => { const { data: tenantInfo } = useFetchTenantInfo(); const tenantId = tenantInfo.tenant_id; const { data, isFetching: loading, refetch, } = useQuery({ queryKey: ['listTenantUser', tenantId], initialData: [], gcTime: 0, enabled: !!tenantId, queryFn: async () => { const { data } = await listTenantUser(tenantId); return data?.data ?? []; }, }); return { data, loading, refetch }; }; export const useAddTenantUser = () => { const { data: tenantInfo } = useFetchTenantInfo(); const queryClient = useQueryClient(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['addTenantUser'], mutationFn: async (email: string) => { const { data } = await addTenantUser(tenantInfo.tenant_id, email); if (data.code === 0) { queryClient.invalidateQueries({ queryKey: ['listTenantUser'] }); } return data?.code; }, }); return { data, loading, addTenantUser: mutateAsync }; }; export const useDeleteTenantUser = () => { const { data: tenantInfo } = useFetchTenantInfo(); const queryClient = useQueryClient(); const { t } = useTranslation(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['deleteTenantUser'], mutationFn: async ({ userId, tenantId, }: { userId: string; tenantId?: string; }) => { const { data } = await deleteTenantUser({ tenantId: tenantId ?? tenantInfo.tenant_id, userId, }); if (data.code === 0) { message.success(t('message.deleted')); queryClient.invalidateQueries({ queryKey: ['listTenantUser'] }); queryClient.invalidateQueries({ queryKey: ['listTenant'] }); } return data?.data ?? []; }, }); return { data, loading, deleteTenantUser: mutateAsync }; }; export const useListTenant = () => { const { data: tenantInfo } = useFetchTenantInfo(); const tenantId = tenantInfo.tenant_id; const { data, isFetching: loading, refetch, } = useQuery({ queryKey: ['listTenant', tenantId], initialData: [], gcTime: 0, enabled: !!tenantId, queryFn: async () => { const { data } = await listTenant(); return data?.data ?? []; }, }); return { data, loading, refetch }; }; export const useAgreeTenant = () => { const queryClient = useQueryClient(); const { t } = useTranslation(); const { data, isPending: loading, mutateAsync, } = useMutation({ mutationKey: ['agreeTenant'], mutationFn: async (tenantId: string) => { const { data } = await agreeTenant(tenantId); if (data.code === 0) { message.success(t('message.operated')); queryClient.invalidateQueries({ queryKey: ['listTenant'] }); } return data?.data ?? []; }, }); return { data, loading, agreeTenant: mutateAsync }; };