import { EToolResources, LocalStorageKeys, defaultAssistantsVersion, } from 'librechat-data-provider'; import { useMutation, useQueryClient } from '@tanstack/react-query'; import type { UseMutationResult } from '@tanstack/react-query'; import type t from 'librechat-data-provider'; import { addConversation, updateConversation, deleteConversation, updateConvoFields, deleteSharedLink, addSharedLink, } from '~/utils'; import { dataService, MutationKeys, QueryKeys, defaultOrderQuery } from 'librechat-data-provider'; import { useSetRecoilState } from 'recoil'; import store from '~/store'; /** Conversations */ export const useGenTitleMutation = (): UseMutationResult< t.TGenTitleResponse, unknown, t.TGenTitleRequest, unknown > => { const queryClient = useQueryClient(); return useMutation((payload: t.TGenTitleRequest) => dataService.genTitle(payload), { onSuccess: (response, vars) => { queryClient.setQueryData( [QueryKeys.conversation, vars.conversationId], (convo: t.TConversation | undefined) => { if (!convo) { return convo; } return { ...convo, title: response.title }; }, ); queryClient.setQueryData([QueryKeys.allConversations], (convoData) => { if (!convoData) { return convoData; } return updateConvoFields(convoData, { conversationId: vars.conversationId, title: response.title, } as t.TConversation); }); document.title = response.title; }, }); }; export const useUpdateConversationMutation = ( id: string, ): UseMutationResult< t.TUpdateConversationResponse, unknown, t.TUpdateConversationRequest, unknown > => { const queryClient = useQueryClient(); return useMutation( (payload: t.TUpdateConversationRequest) => dataService.updateConversation(payload), { onSuccess: (updatedConvo) => { queryClient.setQueryData([QueryKeys.conversation, id], updatedConvo); queryClient.setQueryData([QueryKeys.allConversations], (convoData) => { if (!convoData) { return convoData; } return updateConversation(convoData, updatedConvo); }); }, }, ); }; export const useArchiveConversationMutation = ( id: string, ): UseMutationResult< t.TArchiveConversationResponse, unknown, t.TArchiveConversationRequest, unknown > => { const queryClient = useQueryClient(); return useMutation( (payload: t.TArchiveConversationRequest) => dataService.archiveConversation(payload), { onSuccess: (_data, vars) => { if (vars.isArchived) { queryClient.setQueryData([QueryKeys.conversation, id], null); } else { queryClient.setQueryData([QueryKeys.conversation, id], _data); } queryClient.setQueryData([QueryKeys.allConversations], (convoData) => { if (!convoData) { return convoData; } if (vars.isArchived) { return deleteConversation(convoData, id as string); } else { return addConversation(convoData, _data); } }); queryClient.setQueryData( [QueryKeys.archivedConversations], (convoData) => { if (!convoData) { return convoData; } if (vars.isArchived) { return addConversation(convoData, _data); } else { return deleteConversation(convoData, id as string); } }, ); }, }, ); }; export const useCreateSharedLinkMutation = ( options?: t.CreateSharedLinkOptions, ): UseMutationResult => { const queryClient = useQueryClient(); const { onSuccess, ..._options } = options || {}; return useMutation((payload: t.TSharedLinkRequest) => dataService.createSharedLink(payload), { onSuccess: (_data, vars, context) => { if (!vars.conversationId) { return; } queryClient.setQueryData([QueryKeys.sharedLinks], (sharedLink) => { if (!sharedLink) { return sharedLink; } // If the shared link is public, add it to the shared links cache list if (vars.isPublic) { return addSharedLink(sharedLink, _data); } else { return deleteSharedLink(sharedLink, _data.shareId); } }); queryClient.setQueryData([QueryKeys.sharedLinks, _data.shareId], _data); onSuccess?.(_data, vars, context); }, ...(_options || {}), }); }; export const useUpdateSharedLinkMutation = ( options?: t.UpdateSharedLinkOptions, ): UseMutationResult => { const queryClient = useQueryClient(); const { onSuccess, ..._options } = options || {}; return useMutation((payload: t.TSharedLinkRequest) => dataService.updateSharedLink(payload), { onSuccess: (_data, vars, context) => { if (!vars.conversationId) { return; } queryClient.setQueryData([QueryKeys.sharedLinks], (sharedLink) => { if (!sharedLink) { return sharedLink; } // If the shared link is public, add it to the shared links cache list. if (vars.isPublic) { // Even if the SharedLink data exists in the database, it is not registered in the cache when isPublic is false. // Therefore, when isPublic is true, use addSharedLink instead of updateSharedLink. return addSharedLink(sharedLink, _data); } else { return deleteSharedLink(sharedLink, _data.shareId); } }); queryClient.setQueryData([QueryKeys.sharedLinks, _data.shareId], _data); onSuccess?.(_data, vars, context); }, ...(_options || {}), }); }; export const useDeleteSharedLinkMutation = ( options?: t.DeleteSharedLinkOptions, ): UseMutationResult => { const queryClient = useQueryClient(); const { onSuccess, ..._options } = options || {}; return useMutation(({ shareId }) => dataService.deleteSharedLink(shareId), { onSuccess: (_data, vars, context) => { if (!vars.shareId) { return; } queryClient.setQueryData([QueryKeys.sharedMessages, vars.shareId], null); queryClient.setQueryData([QueryKeys.sharedLinks], (data) => { if (!data) { return data; } return deleteSharedLink(data, vars.shareId); }); onSuccess?.(_data, vars, context); }, ...(_options || {}), }); }; export const useDeleteConversationMutation = ( options?: t.DeleteConversationOptions, ): UseMutationResult< t.TDeleteConversationResponse, unknown, t.TDeleteConversationRequest, unknown > => { const queryClient = useQueryClient(); const { onSuccess, ..._options } = options || {}; return useMutation( (payload: t.TDeleteConversationRequest) => dataService.deleteConversation(payload), { onSuccess: (_data, vars, context) => { if (!vars.conversationId) { return; } const handleDelete = (convoData) => { if (!convoData) { return convoData; } return deleteConversation(convoData, vars.conversationId as string); }; queryClient.setQueryData([QueryKeys.conversation, vars.conversationId], null); queryClient.setQueryData([QueryKeys.allConversations], handleDelete); queryClient.setQueryData( [QueryKeys.archivedConversations], handleDelete, ); onSuccess?.(_data, vars, context); }, ...(_options || {}), }, ); }; export const useForkConvoMutation = ( options?: t.ForkConvoOptions, ): UseMutationResult => { const queryClient = useQueryClient(); const { onSuccess, ..._options } = options || {}; return useMutation((payload: t.TForkConvoRequest) => dataService.forkConversation(payload), { onSuccess: (data, vars, context) => { if (!vars.conversationId) { return; } queryClient.setQueryData( [QueryKeys.conversation, data.conversation.conversationId], data.conversation, ); queryClient.setQueryData([QueryKeys.allConversations], (convoData) => { if (!convoData) { return convoData; } return addConversation(convoData, data.conversation); }); queryClient.setQueryData( [QueryKeys.messages, data.conversation.conversationId], data.messages, ); onSuccess?.(data, vars, context); }, ...(_options || {}), }); }; export const useUploadConversationsMutation = ( _options?: t.MutationOptions, ) => { const queryClient = useQueryClient(); const { onSuccess, onError, onMutate } = _options || {}; // returns the job status or reason of failure const checkJobStatus = async (jobId) => { try { const response = await dataService.queryImportConversationJobStatus(jobId); return response; } catch (error) { throw new Error('Failed to check job status'); } }; // Polls the job status until it is completed, failed, or timed out const pollJobStatus = (jobId, onSuccess, onError) => { let timeElapsed = 0; const timeout = 60000; // Timeout after a minute const pollInterval = 500; // Poll every 500ms const intervalId = setInterval(async () => { try { const statusResponse = await checkJobStatus(jobId); console.log('Polling job status', statusResponse); if (statusResponse.status === 'completed' || statusResponse.status === 'failed') { clearInterval(intervalId); if (statusResponse.status === 'completed') { onSuccess && onSuccess(statusResponse); } else { onError && onError( new Error( statusResponse.failReason ? statusResponse.failReason : 'Failed to import conversations', ), ); } } timeElapsed += pollInterval; // Increment time elapsed by polling interval if (timeElapsed >= timeout) { clearInterval(intervalId); onError && onError(new Error('Polling timed out')); } } catch (error) { clearInterval(intervalId); onError && onError(error); } }, pollInterval); }; return useMutation({ mutationFn: (formData: FormData) => dataService.importConversationsFile(formData), onSuccess: (data, variables, context) => { queryClient.invalidateQueries([QueryKeys.allConversations]); // Assuming the job ID is in the response data const jobId = data.jobId; if (jobId) { // Start polling for job status pollJobStatus( jobId, (statusResponse) => { // This is the final success callback when the job is completed queryClient.invalidateQueries([QueryKeys.allConversations]); // Optionally refresh conversations query if (onSuccess) { onSuccess(statusResponse, variables, context); } }, (error) => { // This is the error callback for job failure or polling errors if (onError) { onError(error, variables, context); } }, ); } }, onError: (err, variables, context) => { if (onError) { onError(err, variables, context); } }, onMutate, }); }; export const useUploadFileMutation = ( _options?: t.UploadMutationOptions, ): UseMutationResult< t.TFileUpload, // response data unknown, // error FormData, // request unknown // context > => { const queryClient = useQueryClient(); const { onSuccess, ...options } = _options || {}; return useMutation([MutationKeys.fileUpload], { mutationFn: (body: FormData) => { const width = body.get('width'); const height = body.get('height'); const version = body.get('version') as number | string; if (height && width && (!version || version != 2)) { return dataService.uploadImage(body); } return dataService.uploadFile(body); }, ...(options || {}), onSuccess: (data, formData, context) => { queryClient.setQueryData([QueryKeys.files], (_files) => [ data, ...(_files ?? []), ]); const endpoint = formData.get('endpoint'); const assistant_id = formData.get('assistant_id'); const message_file = formData.get('message_file'); const tool_resource = formData.get('tool_resource'); if (!assistant_id || message_file === 'true') { onSuccess?.(data, formData, context); return; } queryClient.setQueryData( [QueryKeys.assistants, endpoint, defaultOrderQuery], (prev) => { if (!prev) { return prev; } return { ...prev, data: prev?.data.map((assistant) => { if (assistant.id !== assistant_id) { return assistant; } const update = {}; if (!tool_resource) { update['file_ids'] = [...assistant.file_ids, data.file_id]; } if (tool_resource === EToolResources.code_interpreter) { const prevResources = assistant.tool_resources ?? {}; const prevResource = assistant.tool_resources?.[tool_resource as string] ?? { file_ids: [], }; prevResource.file_ids.push(data.file_id); update['tool_resources'] = { ...prevResources, [tool_resource as string]: prevResource, }; } return { ...assistant, ...update, }; }), }; }, ); onSuccess?.(data, formData, context); }, }); }; export const useDeleteFilesMutation = ( _options?: t.DeleteMutationOptions, ): UseMutationResult< t.DeleteFilesResponse, // response data unknown, // error t.DeleteFilesBody, // request unknown // context > => { const queryClient = useQueryClient(); const { onSuccess, ...options } = _options || {}; return useMutation([MutationKeys.fileDelete], { mutationFn: (body: t.DeleteFilesBody) => dataService.deleteFiles(body.files, body.assistant_id, body.tool_resource), ...(options || {}), onSuccess: (data, ...args) => { queryClient.setQueryData([QueryKeys.files], (cachefiles) => { const { files: filesDeleted } = args[0]; const fileMap = filesDeleted.reduce((acc, file) => { acc.set(file.file_id, file); return acc; }, new Map()); return (cachefiles ?? []).filter((file) => !fileMap.has(file.file_id)); }); onSuccess?.(data, ...args); }, }); }; export const useUpdatePresetMutation = ( options?: t.UpdatePresetOptions, ): UseMutationResult< t.TPreset, // response data unknown, t.TPreset, unknown > => { return useMutation([MutationKeys.updatePreset], { mutationFn: (preset: t.TPreset) => dataService.updatePreset(preset), ...(options || {}), }); }; export const useDeletePresetMutation = ( options?: t.DeletePresetOptions, ): UseMutationResult< t.PresetDeleteResponse, // response data unknown, t.TPreset | undefined, unknown > => { return useMutation([MutationKeys.deletePreset], { mutationFn: (preset: t.TPreset | undefined) => dataService.deletePreset(preset), ...(options || {}), }); }; /* login/logout */ export const useLogoutUserMutation = ( options?: t.LogoutOptions, ): UseMutationResult => { const queryClient = useQueryClient(); const setDefaultPreset = useSetRecoilState(store.defaultPreset); return useMutation([MutationKeys.logoutUser], { mutationFn: () => dataService.logout(), ...(options || {}), onSuccess: (...args) => { options?.onSuccess?.(...args); }, onMutate: (...args) => { setDefaultPreset(null); queryClient.removeQueries(); localStorage.removeItem(LocalStorageKeys.LAST_CONVO_SETUP); localStorage.removeItem(LocalStorageKeys.LAST_MODEL); localStorage.removeItem(LocalStorageKeys.LAST_TOOLS); localStorage.removeItem(LocalStorageKeys.FILES_TO_DELETE); // localStorage.removeItem('lastAssistant'); options?.onMutate?.(...args); }, }); }; /* Avatar upload */ export const useUploadAvatarMutation = ( options?: t.UploadAvatarOptions, ): UseMutationResult< t.AvatarUploadResponse, // response data unknown, // error FormData, // request unknown // context > => { return useMutation([MutationKeys.avatarUpload], { mutationFn: (variables: FormData) => dataService.uploadAvatar(variables), ...(options || {}), }); }; /* Speech to text */ export const useSpeechToTextMutation = ( options?: t.SpeechToTextOptions, ): UseMutationResult< t.SpeechToTextResponse, // response data unknown, // error FormData, // request unknown // context > => { return useMutation([MutationKeys.speechToText], { mutationFn: (variables: FormData) => dataService.speechToText(variables), ...(options || {}), }); }; /* Text to speech */ export const useTextToSpeechMutation = ( options?: t.TextToSpeechOptions, ): UseMutationResult< ArrayBuffer, // response data unknown, // error FormData, // request unknown // context > => { return useMutation([MutationKeys.textToSpeech], { mutationFn: (variables: FormData) => dataService.textToSpeech(variables), ...(options || {}), }); }; /** * ASSISTANTS */ /** * Create a new assistant */ export const useCreateAssistantMutation = ( options?: t.CreateAssistantMutationOptions, ): UseMutationResult => { const queryClient = useQueryClient(); return useMutation( (newAssistantData: t.AssistantCreateParams) => dataService.createAssistant(newAssistantData), { onMutate: (variables) => options?.onMutate?.(variables), onError: (error, variables, context) => options?.onError?.(error, variables, context), onSuccess: (newAssistant, variables, context) => { const listRes = queryClient.getQueryData([ QueryKeys.assistants, variables.endpoint, defaultOrderQuery, ]); if (!listRes) { return options?.onSuccess?.(newAssistant, variables, context); } const currentAssistants = [newAssistant, ...JSON.parse(JSON.stringify(listRes.data))]; queryClient.setQueryData( [QueryKeys.assistants, variables.endpoint, defaultOrderQuery], { ...listRes, data: currentAssistants, }, ); return options?.onSuccess?.(newAssistant, variables, context); }, }, ); }; /** * Hook for updating an assistant */ export const useUpdateAssistantMutation = ( options?: t.UpdateAssistantMutationOptions, ): UseMutationResult< t.Assistant, Error, { assistant_id: string; data: t.AssistantUpdateParams } > => { const queryClient = useQueryClient(); return useMutation( ({ assistant_id, data }: { assistant_id: string; data: t.AssistantUpdateParams }) => { const { endpoint } = data; const endpointsConfig = queryClient.getQueryData([QueryKeys.endpoints]); const version = endpointsConfig?.[endpoint]?.version ?? defaultAssistantsVersion[endpoint]; return dataService.updateAssistant({ data, version, assistant_id, }); }, { onMutate: (variables) => options?.onMutate?.(variables), onError: (error, variables, context) => options?.onError?.(error, variables, context), onSuccess: (updatedAssistant, variables, context) => { const listRes = queryClient.getQueryData([ QueryKeys.assistants, variables.data.endpoint, defaultOrderQuery, ]); if (!listRes) { return options?.onSuccess?.(updatedAssistant, variables, context); } queryClient.setQueryData( [QueryKeys.assistants, variables.data.endpoint, defaultOrderQuery], { ...listRes, data: listRes.data.map((assistant) => { if (assistant.id === variables.assistant_id) { return updatedAssistant; } return assistant; }), }, ); return options?.onSuccess?.(updatedAssistant, variables, context); }, }, ); }; /** * Hook for deleting an assistant */ export const useDeleteAssistantMutation = ( options?: t.DeleteAssistantMutationOptions, ): UseMutationResult => { const queryClient = useQueryClient(); return useMutation( ({ assistant_id, model, endpoint }: t.DeleteAssistantBody) => { const endpointsConfig = queryClient.getQueryData([QueryKeys.endpoints]); const version = endpointsConfig?.[endpoint]?.version ?? defaultAssistantsVersion[endpoint]; return dataService.deleteAssistant({ assistant_id, model, version, endpoint }); }, { onMutate: (variables) => options?.onMutate?.(variables), onError: (error, variables, context) => options?.onError?.(error, variables, context), onSuccess: (_data, variables, context) => { const listRes = queryClient.getQueryData([ QueryKeys.assistants, variables.endpoint, defaultOrderQuery, ]); if (!listRes) { return options?.onSuccess?.(_data, variables, context); } const data = listRes.data.filter((assistant) => assistant.id !== variables.assistant_id); queryClient.setQueryData( [QueryKeys.assistants, variables.endpoint, defaultOrderQuery], { ...listRes, data, }, ); return options?.onSuccess?.(_data, variables, data); }, }, ); }; /** * Hook for uploading an assistant avatar */ export const useUploadAssistantAvatarMutation = ( options?: t.UploadAssistantAvatarOptions, ): UseMutationResult< t.Assistant, // response data unknown, // error t.AssistantAvatarVariables, // request unknown // context > => { return useMutation([MutationKeys.assistantAvatarUpload], { // eslint-disable-next-line @typescript-eslint/no-unused-vars mutationFn: ({ postCreation, ...variables }: t.AssistantAvatarVariables) => dataService.uploadAssistantAvatar(variables), ...(options || {}), }); }; /** * Hook for updating Assistant Actions */ export const useUpdateAction = ( options?: t.UpdateActionOptions, ): UseMutationResult< t.UpdateActionResponse, // response data unknown, // error t.UpdateActionVariables, // request unknown // context > => { const queryClient = useQueryClient(); return useMutation([MutationKeys.updateAction], { mutationFn: (variables: t.UpdateActionVariables) => dataService.updateAction(variables), onMutate: (variables) => options?.onMutate?.(variables), onError: (error, variables, context) => options?.onError?.(error, variables, context), onSuccess: (updateActionResponse, variables, context) => { const listRes = queryClient.getQueryData([ QueryKeys.assistants, variables.endpoint, defaultOrderQuery, ]); if (!listRes) { return options?.onSuccess?.(updateActionResponse, variables, context); } const updatedAssistant = updateActionResponse[1]; queryClient.setQueryData( [QueryKeys.assistants, variables.endpoint, defaultOrderQuery], { ...listRes, data: listRes.data.map((assistant) => { if (assistant.id === variables.assistant_id) { return updatedAssistant; } return assistant; }), }, ); queryClient.setQueryData([QueryKeys.actions], (prev) => { return prev ?.map((action) => { if (action.action_id === variables.action_id) { return updateActionResponse[2]; } return action; }) .concat(variables.action_id ? [] : [updateActionResponse[2]]); }); return options?.onSuccess?.(updateActionResponse, variables, context); }, }); }; /** * Hook for deleting an Assistant Action */ export const useDeleteAction = ( options?: t.DeleteActionOptions, ): UseMutationResult< void, // response data for a delete operation is typically void Error, // error type t.DeleteActionVariables, // request variables unknown // context > => { const queryClient = useQueryClient(); return useMutation([MutationKeys.deleteAction], { mutationFn: (variables: t.DeleteActionVariables) => { const { endpoint } = variables; const endpointsConfig = queryClient.getQueryData([QueryKeys.endpoints]); const version = endpointsConfig?.[endpoint]?.version ?? defaultAssistantsVersion[endpoint]; return dataService.deleteAction({ ...variables, version, }); }, onMutate: (variables) => options?.onMutate?.(variables), onError: (error, variables, context) => options?.onError?.(error, variables, context), onSuccess: (_data, variables, context) => { let domain: string | undefined = ''; queryClient.setQueryData([QueryKeys.actions], (prev) => { return prev?.filter((action) => { domain = action.metadata.domain; return action.action_id !== variables.action_id; }); }); queryClient.setQueryData( [QueryKeys.assistants, variables.endpoint, defaultOrderQuery], (prev) => { if (!prev) { return prev; } return { ...prev, data: prev?.data.map((assistant) => { if (assistant.id === variables.assistant_id) { return { ...assistant, tools: assistant.tools.filter( (tool) => !tool.function?.name.includes(domain ?? ''), ), }; } return assistant; }), }; }, ); return options?.onSuccess?.(_data, variables, context); }, }); };