mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-20 18:30:15 +01:00
refactor: title generation logic
- Changed the title generation endpoint from POST to GET, allowing for more efficient retrieval of titles based on conversation ID. - Implemented exponential backoff for title fetching retries, improving responsiveness and reducing server load. - Introduced a queuing mechanism for title generation, ensuring titles are generated only after job completion. - Updated relevant components and hooks to utilize the new title generation logic, enhancing user experience and application performance.
This commit is contained in:
parent
24263c9874
commit
c7bc5548bc
9 changed files with 118 additions and 97 deletions
|
|
@ -67,16 +67,17 @@ router.get('/:conversationId', async (req, res) => {
|
|||
}
|
||||
});
|
||||
|
||||
router.post('/gen_title', async (req, res) => {
|
||||
const { conversationId } = req.body;
|
||||
router.get('/gen_title/:conversationId', async (req, res) => {
|
||||
const { conversationId } = req.params;
|
||||
const titleCache = getLogStores(CacheKeys.GEN_TITLE);
|
||||
const key = `${req.user.id}-${conversationId}`;
|
||||
let title = await titleCache.get(key);
|
||||
|
||||
if (!title) {
|
||||
// Retry every 1s for up to 20s
|
||||
for (let i = 0; i < 20; i++) {
|
||||
await sleep(1000);
|
||||
// Exponential backoff: 500ms, 1s, 2s, 4s, 8s (total ~15.5s max wait)
|
||||
const delays = [500, 1000, 2000, 4000, 8000];
|
||||
for (const delay of delays) {
|
||||
await sleep(delay);
|
||||
title = await titleCache.get(key);
|
||||
if (title) {
|
||||
break;
|
||||
|
|
|
|||
|
|
@ -3,8 +3,9 @@ import { useRecoilValue } from 'recoil';
|
|||
import { AnimatePresence, motion } from 'framer-motion';
|
||||
import { Skeleton, useMediaQuery } from '@librechat/client';
|
||||
import { PermissionTypes, Permissions } from 'librechat-data-provider';
|
||||
import type { ConversationListResponse } from 'librechat-data-provider';
|
||||
import type { InfiniteQueryObserverResult } from '@tanstack/react-query';
|
||||
import type { ConversationListResponse } from 'librechat-data-provider';
|
||||
import type { List } from 'react-virtualized';
|
||||
import {
|
||||
useLocalize,
|
||||
useHasAccess,
|
||||
|
|
@ -12,7 +13,7 @@ import {
|
|||
useLocalStorage,
|
||||
useNavScrolling,
|
||||
} from '~/hooks';
|
||||
import { useConversationsInfiniteQuery } from '~/data-provider';
|
||||
import { useConversationsInfiniteQuery, useTitleGeneration } from '~/data-provider';
|
||||
import { Conversations } from '~/components/Conversations';
|
||||
import SearchBar from './SearchBar';
|
||||
import NewChat from './NewChat';
|
||||
|
|
@ -63,6 +64,7 @@ const Nav = memo(
|
|||
}) => {
|
||||
const localize = useLocalize();
|
||||
const { isAuthenticated } = useAuthContext();
|
||||
useTitleGeneration(isAuthenticated);
|
||||
|
||||
const [navWidth, setNavWidth] = useState(NAV_WIDTH_DESKTOP);
|
||||
const isSmallScreen = useMediaQuery('(max-width: 768px)');
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import { useQuery } from '@tanstack/react-query';
|
||||
import { useEffect, useMemo, useState } from 'react';
|
||||
import { QueryKeys, request, dataService } from 'librechat-data-provider';
|
||||
import type { Agents } from 'librechat-data-provider';
|
||||
import { useQuery, useQueries, useQueryClient } from '@tanstack/react-query';
|
||||
import type { Agents, TConversation } from 'librechat-data-provider';
|
||||
import { updateConvoInAllQueries } from '~/utils';
|
||||
|
||||
export interface StreamStatusResponse {
|
||||
active: boolean;
|
||||
|
|
@ -11,67 +13,123 @@ export interface StreamStatusResponse {
|
|||
resumeState?: Agents.ResumeState;
|
||||
}
|
||||
|
||||
/**
|
||||
* Query key for stream status
|
||||
*/
|
||||
export const streamStatusQueryKey = (conversationId: string) => ['streamStatus', conversationId];
|
||||
|
||||
/**
|
||||
* Fetch stream status for a conversation
|
||||
*/
|
||||
export const fetchStreamStatus = async (conversationId: string): Promise<StreamStatusResponse> => {
|
||||
console.log('[fetchStreamStatus] Fetching status for:', conversationId);
|
||||
const result = await request.get<StreamStatusResponse>(
|
||||
`/api/agents/chat/status/${conversationId}`,
|
||||
);
|
||||
console.log('[fetchStreamStatus] Result:', result);
|
||||
return result;
|
||||
return request.get<StreamStatusResponse>(`/api/agents/chat/status/${conversationId}`);
|
||||
};
|
||||
|
||||
/**
|
||||
* React Query hook for checking if a conversation has an active generation stream.
|
||||
* Only fetches when conversationId is provided and resumable streams are enabled.
|
||||
*/
|
||||
export function useStreamStatus(conversationId: string | undefined, enabled = true) {
|
||||
return useQuery({
|
||||
queryKey: streamStatusQueryKey(conversationId || ''),
|
||||
queryFn: () => fetchStreamStatus(conversationId!),
|
||||
enabled: !!conversationId && enabled,
|
||||
staleTime: 1000, // Consider stale after 1 second
|
||||
staleTime: 1000,
|
||||
refetchOnMount: true,
|
||||
refetchOnWindowFocus: true,
|
||||
retry: false,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Query key for active jobs
|
||||
*/
|
||||
export const activeJobsQueryKey = [QueryKeys.activeJobs] as const;
|
||||
export const genTitleQueryKey = (conversationId: string) => ['genTitle', conversationId] as const;
|
||||
|
||||
// Module-level queue for title generation (survives re-renders)
|
||||
// Stores conversationIds that need title generation once their job completes
|
||||
const titleQueue = new Set<string>();
|
||||
const processedTitles = new Set<string>();
|
||||
|
||||
/** Queue a conversation for title generation (call when starting new conversation) */
|
||||
export function queueTitleGeneration(conversationId: string) {
|
||||
if (!processedTitles.has(conversationId)) {
|
||||
titleQueue.add(conversationId);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* React Query hook for getting all active job IDs for the current user.
|
||||
* Used to show generation indicators in the conversation list.
|
||||
*
|
||||
* Key behaviors:
|
||||
* - Fetches on mount to get initial state (handles page refresh)
|
||||
* - Refetches on window focus (handles multi-tab scenarios)
|
||||
* - Optimistic updates from useResumableSSE when jobs start/complete
|
||||
* - Polls every 5s while there are active jobs (catches completions when navigated away)
|
||||
* Hook to process the title generation queue.
|
||||
* Only fetches titles AFTER the job completes (not in activeJobIds).
|
||||
* Place this high in the component tree (e.g., Nav.tsx).
|
||||
*/
|
||||
export function useTitleGeneration(enabled = true) {
|
||||
const queryClient = useQueryClient();
|
||||
const [readyToFetch, setReadyToFetch] = useState<string[]>([]);
|
||||
|
||||
const { data: activeJobsData } = useActiveJobs(enabled);
|
||||
const activeJobIds = useMemo(
|
||||
() => activeJobsData?.activeJobIds ?? [],
|
||||
[activeJobsData?.activeJobIds],
|
||||
);
|
||||
|
||||
// Check queue for completed jobs and fetch titles immediately
|
||||
useEffect(() => {
|
||||
const activeSet = new Set(activeJobIds);
|
||||
const completedJobs: string[] = [];
|
||||
|
||||
for (const conversationId of titleQueue) {
|
||||
if (!activeSet.has(conversationId) && !processedTitles.has(conversationId)) {
|
||||
completedJobs.push(conversationId);
|
||||
}
|
||||
}
|
||||
|
||||
if (completedJobs.length > 0) {
|
||||
setReadyToFetch((prev) => [...new Set([...prev, ...completedJobs])]);
|
||||
}
|
||||
}, [activeJobIds]);
|
||||
|
||||
// Fetch titles for ready conversations
|
||||
const titleQueries = useQueries({
|
||||
queries: readyToFetch.map((conversationId) => ({
|
||||
queryKey: genTitleQueryKey(conversationId),
|
||||
queryFn: () => dataService.genTitle({ conversationId }),
|
||||
staleTime: Infinity,
|
||||
retry: false,
|
||||
})),
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
titleQueries.forEach((titleQuery, index) => {
|
||||
const conversationId = readyToFetch[index];
|
||||
if (!conversationId || processedTitles.has(conversationId)) return;
|
||||
|
||||
if (titleQuery.isSuccess && titleQuery.data) {
|
||||
const { title } = titleQuery.data;
|
||||
queryClient.setQueryData(
|
||||
[QueryKeys.conversation, conversationId],
|
||||
(convo: TConversation | undefined) => (convo ? { ...convo, title } : convo),
|
||||
);
|
||||
updateConvoInAllQueries(queryClient, conversationId, (c) => ({ ...c, title }));
|
||||
// Only update document title if this conversation is currently active
|
||||
if (window.location.pathname.includes(conversationId)) {
|
||||
document.title = title;
|
||||
}
|
||||
processedTitles.add(conversationId);
|
||||
titleQueue.delete(conversationId);
|
||||
setReadyToFetch((prev) => prev.filter((id) => id !== conversationId));
|
||||
} else if (titleQuery.isError) {
|
||||
// Mark as processed even on error to avoid infinite retries
|
||||
processedTitles.add(conversationId);
|
||||
titleQueue.delete(conversationId);
|
||||
setReadyToFetch((prev) => prev.filter((id) => id !== conversationId));
|
||||
}
|
||||
});
|
||||
}, [titleQueries, readyToFetch, queryClient]);
|
||||
}
|
||||
|
||||
/**
|
||||
* React Query hook for active job IDs.
|
||||
* - Polls while jobs are active
|
||||
* - Shows generation indicators in conversation list
|
||||
*/
|
||||
export function useActiveJobs(enabled = true) {
|
||||
return useQuery({
|
||||
queryKey: activeJobsQueryKey,
|
||||
queryFn: () => dataService.getActiveJobs(),
|
||||
enabled,
|
||||
staleTime: 5_000, // 5s - short to catch completions quickly
|
||||
staleTime: 5_000,
|
||||
refetchOnMount: true,
|
||||
refetchOnWindowFocus: true, // Catch up on tab switch (multi-tab scenario)
|
||||
// Poll every 5s while there are active jobs to catch completions when navigated away
|
||||
refetchInterval: (data) => {
|
||||
const hasActiveJobs = (data?.activeJobIds?.length ?? 0) > 0;
|
||||
return hasActiveJobs ? 5_000 : false;
|
||||
},
|
||||
refetchOnWindowFocus: true,
|
||||
refetchInterval: (data) => ((data?.activeJobIds?.length ?? 0) > 0 ? 5_000 : false),
|
||||
retry: false,
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,24 +25,6 @@ export type TGenTitleMutation = UseMutationResult<
|
|||
unknown
|
||||
>;
|
||||
|
||||
export const useGenTitleMutation = (): TGenTitleMutation => {
|
||||
const queryClient = useQueryClient();
|
||||
return useMutation((payload: t.TGenTitleRequest) => dataService.genTitle(payload), {
|
||||
onSuccess: (response, vars) => {
|
||||
queryClient.setQueryData(
|
||||
[QueryKeys.conversation, vars.conversationId],
|
||||
(convo: t.TConversation | undefined) =>
|
||||
convo ? { ...convo, title: response.title } : convo,
|
||||
);
|
||||
updateConvoInAllQueries(queryClient, vars.conversationId, (c) => ({
|
||||
...c,
|
||||
title: response.title,
|
||||
}));
|
||||
document.title = response.title;
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
export const useUpdateConversationMutation = (
|
||||
id: string,
|
||||
): UseMutationResult<
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ import type {
|
|||
} from 'librechat-data-provider';
|
||||
import type { TResData, TFinalResData, ConvoGenerator } from '~/common';
|
||||
import type { InfiniteData } from '@tanstack/react-query';
|
||||
import type { TGenTitleMutation } from '~/data-provider';
|
||||
import type { SetterOrUpdater, Resetter } from 'recoil';
|
||||
import type { ConversationCursorData } from '~/utils';
|
||||
import {
|
||||
|
|
@ -54,7 +53,6 @@ type TSyncData = {
|
|||
|
||||
export type EventHandlerParams = {
|
||||
isAddedRequest?: boolean;
|
||||
genTitle?: TGenTitleMutation;
|
||||
setCompleted: React.Dispatch<React.SetStateAction<Set<unknown>>>;
|
||||
setMessages: (messages: TMessage[]) => void;
|
||||
getMessages: () => TMessage[] | undefined;
|
||||
|
|
@ -167,7 +165,6 @@ export const getConvoTitle = ({
|
|||
};
|
||||
|
||||
export default function useEventHandlers({
|
||||
genTitle,
|
||||
setMessages,
|
||||
getMessages,
|
||||
setCompleted,
|
||||
|
|
@ -258,13 +255,6 @@ export default function useEventHandlers({
|
|||
removeConvoFromAllQueries(queryClient, submission.conversation.conversationId as string);
|
||||
}
|
||||
|
||||
// refresh title
|
||||
if (genTitle && isNewConvo && requestMessage.parentMessageId === Constants.NO_PARENT) {
|
||||
setTimeout(() => {
|
||||
genTitle.mutate({ conversationId: convoUpdate.conversationId as string });
|
||||
}, 2500);
|
||||
}
|
||||
|
||||
if (setConversation && !isAddedRequest) {
|
||||
setConversation((prevState) => {
|
||||
const update = { ...prevState, ...convoUpdate };
|
||||
|
|
@ -274,7 +264,7 @@ export default function useEventHandlers({
|
|||
|
||||
setIsSubmitting(false);
|
||||
},
|
||||
[setMessages, setConversation, genTitle, isAddedRequest, queryClient, setIsSubmitting],
|
||||
[setMessages, setConversation, isAddedRequest, queryClient, setIsSubmitting],
|
||||
);
|
||||
|
||||
const syncHandler = useCallback(
|
||||
|
|
@ -443,7 +433,7 @@ export default function useEventHandlers({
|
|||
messages,
|
||||
conversation: submissionConvo,
|
||||
isRegenerate = false,
|
||||
isTemporary = false,
|
||||
isTemporary: _isTemporary = false,
|
||||
} = submission;
|
||||
|
||||
try {
|
||||
|
|
@ -532,19 +522,6 @@ export default function useEventHandlers({
|
|||
removeConvoFromAllQueries(queryClient, submissionConvo.conversationId);
|
||||
}
|
||||
|
||||
/* Refresh title */
|
||||
if (
|
||||
genTitle &&
|
||||
isNewConvo &&
|
||||
!isTemporary &&
|
||||
requestMessage &&
|
||||
requestMessage.parentMessageId === Constants.NO_PARENT
|
||||
) {
|
||||
setTimeout(() => {
|
||||
genTitle.mutate({ conversationId: conversation.conversationId as string });
|
||||
}, 2500);
|
||||
}
|
||||
|
||||
if (setConversation && isAddedRequest !== true) {
|
||||
setConversation((prevState) => {
|
||||
const update = {
|
||||
|
|
@ -588,7 +565,6 @@ export default function useEventHandlers({
|
|||
},
|
||||
[
|
||||
navigate,
|
||||
genTitle,
|
||||
getMessages,
|
||||
setMessages,
|
||||
queryClient,
|
||||
|
|
|
|||
|
|
@ -12,8 +12,8 @@ import {
|
|||
} from 'librechat-data-provider';
|
||||
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
|
||||
import type { EventHandlerParams } from './useEventHandlers';
|
||||
import { useGenTitleMutation, useGetStartupConfig, useGetUserBalance } from '~/data-provider';
|
||||
import { activeJobsQueryKey } from '~/data-provider/SSE/queries';
|
||||
import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
|
||||
import { activeJobsQueryKey, queueTitleGeneration } from '~/data-provider/SSE/queries';
|
||||
import { useAuthContext } from '~/hooks/AuthContext';
|
||||
import useEventHandlers from './useEventHandlers';
|
||||
import store from '~/store';
|
||||
|
|
@ -61,7 +61,6 @@ export default function useResumableSSE(
|
|||
isAddedRequest = false,
|
||||
runIndex = 0,
|
||||
) {
|
||||
const genTitle = useGenTitleMutation();
|
||||
const queryClient = useQueryClient();
|
||||
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
|
||||
|
||||
|
|
@ -123,7 +122,6 @@ export default function useResumableSSE(
|
|||
attachmentHandler,
|
||||
resetContentHandler,
|
||||
} = useEventHandlers({
|
||||
genTitle,
|
||||
setMessages,
|
||||
getMessages,
|
||||
setCompleted,
|
||||
|
|
@ -596,6 +594,11 @@ export default function useResumableSSE(
|
|||
setStreamId(newStreamId);
|
||||
// Optimistically add to active jobs
|
||||
addActiveJob(newStreamId);
|
||||
// Queue title generation if this is a new conversation (first message)
|
||||
const isNewConvo = submission.userMessage?.parentMessageId === Constants.NO_PARENT;
|
||||
if (isNewConvo) {
|
||||
queueTitleGeneration(newStreamId);
|
||||
}
|
||||
subscribeToStream(newStreamId, submission);
|
||||
} else {
|
||||
console.error('[ResumableSSE] Failed to get streamId from startGeneration');
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import {
|
|||
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
|
||||
import type { EventHandlerParams } from './useEventHandlers';
|
||||
import type { TResData } from '~/common';
|
||||
import { useGenTitleMutation, useGetStartupConfig, useGetUserBalance } from '~/data-provider';
|
||||
import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
|
||||
import { useAuthContext } from '~/hooks/AuthContext';
|
||||
import useEventHandlers from './useEventHandlers';
|
||||
import store from '~/store';
|
||||
|
|
@ -44,7 +44,6 @@ export default function useSSE(
|
|||
isAddedRequest = false,
|
||||
runIndex = 0,
|
||||
) {
|
||||
const genTitle = useGenTitleMutation();
|
||||
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
|
||||
|
||||
const { token, isAuthenticated } = useAuthContext();
|
||||
|
|
@ -73,7 +72,6 @@ export default function useSSE(
|
|||
attachmentHandler,
|
||||
abortConversation,
|
||||
} = useEventHandlers({
|
||||
genTitle,
|
||||
setMessages,
|
||||
getMessages,
|
||||
setCompleted,
|
||||
|
|
|
|||
|
|
@ -101,7 +101,8 @@ export const conversations = (params: q.ConversationListParams) => {
|
|||
|
||||
export const conversationById = (id: string) => `${conversationsRoot}/${id}`;
|
||||
|
||||
export const genTitle = () => `${conversationsRoot}/gen_title`;
|
||||
export const genTitle = (conversationId: string) =>
|
||||
`${conversationsRoot}/gen_title/${encodeURIComponent(conversationId)}`;
|
||||
|
||||
export const updateConversation = () => `${conversationsRoot}/update`;
|
||||
|
||||
|
|
|
|||
|
|
@ -724,7 +724,7 @@ export function archiveConversation(
|
|||
}
|
||||
|
||||
export function genTitle(payload: m.TGenTitleRequest): Promise<m.TGenTitleResponse> {
|
||||
return request.post(endpoints.genTitle(), payload);
|
||||
return request.get(endpoints.genTitle(payload.conversationId));
|
||||
}
|
||||
|
||||
export const listMessages = (params?: q.MessagesListParams): Promise<q.MessagesListResponse> => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue