diff --git a/client/src/data-provider/queries.ts b/client/src/data-provider/queries.ts index 447775067e..24d1d6c9e5 100644 --- a/client/src/data-provider/queries.ts +++ b/client/src/data-provider/queries.ts @@ -32,6 +32,20 @@ import type { } from 'librechat-data-provider'; import type { ConversationCursorData } from '~/utils/convos'; import { findConversationInInfinite } from '~/utils'; +type TConversationCosts = { + conversationId: string; + totals: { + prompt: { usd: number; tokenCount: number }; + completion: { usd: number; tokenCount: number }; + total: { usd: number; tokenCount: number }; + }; + perMessage: Array<{ + messageId: string; + tokenType: 'prompt' | 'completion'; + tokenCount: number; + usd: number; + }>; +}; export const useGetPresetsQuery = ( config?: UseQueryOptions, @@ -76,6 +90,25 @@ export const useGetConvoIdQuery = ( ); }; +export const useGetConversationCosts = ( + conversationId: string, + config?: UseQueryOptions, +): QueryObserverResult => { + return useQuery( + [QueryKeys.conversation, conversationId, 'costs'], + () => { + return dataService.getConversationCosts(conversationId) as unknown as TConversationCosts; + }, + { + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, + enabled: !!conversationId && conversationId !== 'new', + ...config, + }, + ); +}; + export const useConversationsInfiniteQuery = ( params: ConversationListParams, config?: UseInfiniteQueryOptions, diff --git a/client/src/hooks/SSE/useEventHandlers.ts b/client/src/hooks/SSE/useEventHandlers.ts index 1d860fbb7a..9b95453863 100644 --- a/client/src/hooks/SSE/useEventHandlers.ts +++ b/client/src/hooks/SSE/useEventHandlers.ts @@ -232,8 +232,14 @@ export default function useEventHandlers({ }, ]); } + + if (userMessage?.conversationId) { + queryClient.invalidateQueries({ + queryKey: [QueryKeys.conversation, userMessage.conversationId, 'costs'], + }); + } }, - [setMessages, announcePolite, setIsSubmitting], + [setMessages, announcePolite, setIsSubmitting, queryClient], ); const cancelHandler = useCallback( @@ -275,6 +281,12 @@ export default function useEventHandlers({ }); } + if (convoUpdate?.conversationId) { + queryClient.invalidateQueries({ + queryKey: [QueryKeys.conversation, convoUpdate.conversationId, 'costs'], + }); + } + setIsSubmitting(false); }, [setMessages, setConversation, genTitle, isAddedRequest, queryClient, setIsSubmitting], @@ -341,6 +353,12 @@ export default function useEventHandlers({ if (resetLatestMessage) { resetLatestMessage(); } + + if (conversationId) { + queryClient.invalidateQueries({ + queryKey: [QueryKeys.conversation, conversationId, 'costs'], + }); + } }, [ queryClient, @@ -527,6 +545,12 @@ export default function useEventHandlers({ ); } + if (conversation.conversationId) { + queryClient.invalidateQueries({ + queryKey: [QueryKeys.conversation, conversation.conversationId, 'costs'], + }); + } + if (isNewConvo && submissionConvo.conversationId) { removeConvoFromAllQueries(queryClient, submissionConvo.conversationId); }