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:
Danny Avila 2025-12-17 17:10:14 -05:00
parent 5a278f4f4f
commit 1ea35373ae
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
9 changed files with 118 additions and 97 deletions

View file

@ -67,16 +67,17 @@ router.get('/:conversationId', async (req, res) => {
} }
}); });
router.post('/gen_title', async (req, res) => { router.get('/gen_title/:conversationId', async (req, res) => {
const { conversationId } = req.body; const { conversationId } = req.params;
const titleCache = getLogStores(CacheKeys.GEN_TITLE); const titleCache = getLogStores(CacheKeys.GEN_TITLE);
const key = `${req.user.id}-${conversationId}`; const key = `${req.user.id}-${conversationId}`;
let title = await titleCache.get(key); let title = await titleCache.get(key);
if (!title) { if (!title) {
// Retry every 1s for up to 20s // Exponential backoff: 500ms, 1s, 2s, 4s, 8s (total ~15.5s max wait)
for (let i = 0; i < 20; i++) { const delays = [500, 1000, 2000, 4000, 8000];
await sleep(1000); for (const delay of delays) {
await sleep(delay);
title = await titleCache.get(key); title = await titleCache.get(key);
if (title) { if (title) {
break; break;

View file

@ -3,8 +3,9 @@ import { useRecoilValue } from 'recoil';
import { AnimatePresence, motion } from 'framer-motion'; import { AnimatePresence, motion } from 'framer-motion';
import { Skeleton, useMediaQuery } from '@librechat/client'; import { Skeleton, useMediaQuery } from '@librechat/client';
import { PermissionTypes, Permissions } from 'librechat-data-provider'; import { PermissionTypes, Permissions } from 'librechat-data-provider';
import type { ConversationListResponse } from 'librechat-data-provider';
import type { InfiniteQueryObserverResult } from '@tanstack/react-query'; import type { InfiniteQueryObserverResult } from '@tanstack/react-query';
import type { ConversationListResponse } from 'librechat-data-provider';
import type { List } from 'react-virtualized';
import { import {
useLocalize, useLocalize,
useHasAccess, useHasAccess,
@ -12,7 +13,7 @@ import {
useLocalStorage, useLocalStorage,
useNavScrolling, useNavScrolling,
} from '~/hooks'; } from '~/hooks';
import { useConversationsInfiniteQuery } from '~/data-provider'; import { useConversationsInfiniteQuery, useTitleGeneration } from '~/data-provider';
import { Conversations } from '~/components/Conversations'; import { Conversations } from '~/components/Conversations';
import SearchBar from './SearchBar'; import SearchBar from './SearchBar';
import NewChat from './NewChat'; import NewChat from './NewChat';
@ -63,6 +64,7 @@ const Nav = memo(
}) => { }) => {
const localize = useLocalize(); const localize = useLocalize();
const { isAuthenticated } = useAuthContext(); const { isAuthenticated } = useAuthContext();
useTitleGeneration(isAuthenticated);
const [navWidth, setNavWidth] = useState(NAV_WIDTH_DESKTOP); const [navWidth, setNavWidth] = useState(NAV_WIDTH_DESKTOP);
const isSmallScreen = useMediaQuery('(max-width: 768px)'); const isSmallScreen = useMediaQuery('(max-width: 768px)');

View file

@ -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 { 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 { export interface StreamStatusResponse {
active: boolean; active: boolean;
@ -11,67 +13,123 @@ export interface StreamStatusResponse {
resumeState?: Agents.ResumeState; resumeState?: Agents.ResumeState;
} }
/**
* Query key for stream status
*/
export const streamStatusQueryKey = (conversationId: string) => ['streamStatus', conversationId]; export const streamStatusQueryKey = (conversationId: string) => ['streamStatus', conversationId];
/**
* Fetch stream status for a conversation
*/
export const fetchStreamStatus = async (conversationId: string): Promise<StreamStatusResponse> => { export const fetchStreamStatus = async (conversationId: string): Promise<StreamStatusResponse> => {
console.log('[fetchStreamStatus] Fetching status for:', conversationId); return request.get<StreamStatusResponse>(`/api/agents/chat/status/${conversationId}`);
const result = await request.get<StreamStatusResponse>(
`/api/agents/chat/status/${conversationId}`,
);
console.log('[fetchStreamStatus] Result:', result);
return result;
}; };
/**
* 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) { export function useStreamStatus(conversationId: string | undefined, enabled = true) {
return useQuery({ return useQuery({
queryKey: streamStatusQueryKey(conversationId || ''), queryKey: streamStatusQueryKey(conversationId || ''),
queryFn: () => fetchStreamStatus(conversationId!), queryFn: () => fetchStreamStatus(conversationId!),
enabled: !!conversationId && enabled, enabled: !!conversationId && enabled,
staleTime: 1000, // Consider stale after 1 second staleTime: 1000,
refetchOnMount: true, refetchOnMount: true,
refetchOnWindowFocus: true, refetchOnWindowFocus: true,
retry: false, retry: false,
}); });
} }
/**
* Query key for active jobs
*/
export const activeJobsQueryKey = [QueryKeys.activeJobs] as const; 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. * Hook to process the title generation queue.
* Used to show generation indicators in the conversation list. * Only fetches titles AFTER the job completes (not in activeJobIds).
* * Place this high in the component tree (e.g., Nav.tsx).
* Key behaviors: */
* - Fetches on mount to get initial state (handles page refresh) export function useTitleGeneration(enabled = true) {
* - Refetches on window focus (handles multi-tab scenarios) const queryClient = useQueryClient();
* - Optimistic updates from useResumableSSE when jobs start/complete const [readyToFetch, setReadyToFetch] = useState<string[]>([]);
* - Polls every 5s while there are active jobs (catches completions when navigated away)
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) { export function useActiveJobs(enabled = true) {
return useQuery({ return useQuery({
queryKey: activeJobsQueryKey, queryKey: activeJobsQueryKey,
queryFn: () => dataService.getActiveJobs(), queryFn: () => dataService.getActiveJobs(),
enabled, enabled,
staleTime: 5_000, // 5s - short to catch completions quickly staleTime: 5_000,
refetchOnMount: true, refetchOnMount: true,
refetchOnWindowFocus: true, // Catch up on tab switch (multi-tab scenario) refetchOnWindowFocus: true,
// Poll every 5s while there are active jobs to catch completions when navigated away refetchInterval: (data) => ((data?.activeJobIds?.length ?? 0) > 0 ? 5_000 : false),
refetchInterval: (data) => {
const hasActiveJobs = (data?.activeJobIds?.length ?? 0) > 0;
return hasActiveJobs ? 5_000 : false;
},
retry: false, retry: false,
}); });
} }

View file

@ -25,24 +25,6 @@ export type TGenTitleMutation = UseMutationResult<
unknown 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 = ( export const useUpdateConversationMutation = (
id: string, id: string,
): UseMutationResult< ): UseMutationResult<

View file

@ -21,7 +21,6 @@ import type {
} from 'librechat-data-provider'; } from 'librechat-data-provider';
import type { TResData, TFinalResData, ConvoGenerator } from '~/common'; import type { TResData, TFinalResData, ConvoGenerator } from '~/common';
import type { InfiniteData } from '@tanstack/react-query'; import type { InfiniteData } from '@tanstack/react-query';
import type { TGenTitleMutation } from '~/data-provider';
import type { SetterOrUpdater, Resetter } from 'recoil'; import type { SetterOrUpdater, Resetter } from 'recoil';
import type { ConversationCursorData } from '~/utils'; import type { ConversationCursorData } from '~/utils';
import { import {
@ -54,7 +53,6 @@ type TSyncData = {
export type EventHandlerParams = { export type EventHandlerParams = {
isAddedRequest?: boolean; isAddedRequest?: boolean;
genTitle?: TGenTitleMutation;
setCompleted: React.Dispatch<React.SetStateAction<Set<unknown>>>; setCompleted: React.Dispatch<React.SetStateAction<Set<unknown>>>;
setMessages: (messages: TMessage[]) => void; setMessages: (messages: TMessage[]) => void;
getMessages: () => TMessage[] | undefined; getMessages: () => TMessage[] | undefined;
@ -167,7 +165,6 @@ export const getConvoTitle = ({
}; };
export default function useEventHandlers({ export default function useEventHandlers({
genTitle,
setMessages, setMessages,
getMessages, getMessages,
setCompleted, setCompleted,
@ -258,13 +255,6 @@ export default function useEventHandlers({
removeConvoFromAllQueries(queryClient, submission.conversation.conversationId as string); 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) { if (setConversation && !isAddedRequest) {
setConversation((prevState) => { setConversation((prevState) => {
const update = { ...prevState, ...convoUpdate }; const update = { ...prevState, ...convoUpdate };
@ -274,7 +264,7 @@ export default function useEventHandlers({
setIsSubmitting(false); setIsSubmitting(false);
}, },
[setMessages, setConversation, genTitle, isAddedRequest, queryClient, setIsSubmitting], [setMessages, setConversation, isAddedRequest, queryClient, setIsSubmitting],
); );
const syncHandler = useCallback( const syncHandler = useCallback(
@ -443,7 +433,7 @@ export default function useEventHandlers({
messages, messages,
conversation: submissionConvo, conversation: submissionConvo,
isRegenerate = false, isRegenerate = false,
isTemporary = false, isTemporary: _isTemporary = false,
} = submission; } = submission;
try { try {
@ -532,19 +522,6 @@ export default function useEventHandlers({
removeConvoFromAllQueries(queryClient, submissionConvo.conversationId); 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) { if (setConversation && isAddedRequest !== true) {
setConversation((prevState) => { setConversation((prevState) => {
const update = { const update = {
@ -588,7 +565,6 @@ export default function useEventHandlers({
}, },
[ [
navigate, navigate,
genTitle,
getMessages, getMessages,
setMessages, setMessages,
queryClient, queryClient,

View file

@ -12,8 +12,8 @@ import {
} from 'librechat-data-provider'; } from 'librechat-data-provider';
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider'; import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
import type { EventHandlerParams } from './useEventHandlers'; import type { EventHandlerParams } from './useEventHandlers';
import { useGenTitleMutation, useGetStartupConfig, useGetUserBalance } from '~/data-provider'; import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
import { activeJobsQueryKey } from '~/data-provider/SSE/queries'; import { activeJobsQueryKey, queueTitleGeneration } from '~/data-provider/SSE/queries';
import { useAuthContext } from '~/hooks/AuthContext'; import { useAuthContext } from '~/hooks/AuthContext';
import useEventHandlers from './useEventHandlers'; import useEventHandlers from './useEventHandlers';
import store from '~/store'; import store from '~/store';
@ -61,7 +61,6 @@ export default function useResumableSSE(
isAddedRequest = false, isAddedRequest = false,
runIndex = 0, runIndex = 0,
) { ) {
const genTitle = useGenTitleMutation();
const queryClient = useQueryClient(); const queryClient = useQueryClient();
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex)); const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
@ -123,7 +122,6 @@ export default function useResumableSSE(
attachmentHandler, attachmentHandler,
resetContentHandler, resetContentHandler,
} = useEventHandlers({ } = useEventHandlers({
genTitle,
setMessages, setMessages,
getMessages, getMessages,
setCompleted, setCompleted,
@ -596,6 +594,11 @@ export default function useResumableSSE(
setStreamId(newStreamId); setStreamId(newStreamId);
// Optimistically add to active jobs // Optimistically add to active jobs
addActiveJob(newStreamId); 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); subscribeToStream(newStreamId, submission);
} else { } else {
console.error('[ResumableSSE] Failed to get streamId from startGeneration'); console.error('[ResumableSSE] Failed to get streamId from startGeneration');

View file

@ -13,7 +13,7 @@ import {
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider'; import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
import type { EventHandlerParams } from './useEventHandlers'; import type { EventHandlerParams } from './useEventHandlers';
import type { TResData } from '~/common'; import type { TResData } from '~/common';
import { useGenTitleMutation, useGetStartupConfig, useGetUserBalance } from '~/data-provider'; import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
import { useAuthContext } from '~/hooks/AuthContext'; import { useAuthContext } from '~/hooks/AuthContext';
import useEventHandlers from './useEventHandlers'; import useEventHandlers from './useEventHandlers';
import store from '~/store'; import store from '~/store';
@ -44,7 +44,6 @@ export default function useSSE(
isAddedRequest = false, isAddedRequest = false,
runIndex = 0, runIndex = 0,
) { ) {
const genTitle = useGenTitleMutation();
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex)); const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
const { token, isAuthenticated } = useAuthContext(); const { token, isAuthenticated } = useAuthContext();
@ -73,7 +72,6 @@ export default function useSSE(
attachmentHandler, attachmentHandler,
abortConversation, abortConversation,
} = useEventHandlers({ } = useEventHandlers({
genTitle,
setMessages, setMessages,
getMessages, getMessages,
setCompleted, setCompleted,

View file

@ -101,7 +101,8 @@ export const conversations = (params: q.ConversationListParams) => {
export const conversationById = (id: string) => `${conversationsRoot}/${id}`; 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`; export const updateConversation = () => `${conversationsRoot}/update`;

View file

@ -724,7 +724,7 @@ export function archiveConversation(
} }
export function genTitle(payload: m.TGenTitleRequest): Promise<m.TGenTitleResponse> { 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> => { export const listMessages = (params?: q.MessagesListParams): Promise<q.MessagesListResponse> => {