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 24263c9874
commit c7bc5548bc
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) => {
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;

View file

@ -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)');

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 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,
});
}

View file

@ -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<

View file

@ -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,

View file

@ -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');

View file

@ -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,

View file

@ -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`;

View file

@ -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> => {