diff --git a/client/src/components/Chat/Input/ChatForm.tsx b/client/src/components/Chat/Input/ChatForm.tsx index 9e0ad7f382..7030d8fe35 100644 --- a/client/src/components/Chat/Input/ChatForm.tsx +++ b/client/src/components/Chat/Input/ChatForm.tsx @@ -3,7 +3,7 @@ import { useWatch } from 'react-hook-form'; import { TextareaAutosize } from '@librechat/client'; import { useRecoilState, useRecoilValue } from 'recoil'; import { Constants, isAssistantsEndpoint, isAgentsEndpoint } from 'librechat-data-provider'; -import type { TConversation } from 'librechat-data-provider'; +import type { TConversation, TMessage } from 'librechat-data-provider'; import type { ExtendedFile, FileSetter, ConvoGenerator } from '~/common'; import { useChatContext, @@ -32,6 +32,7 @@ import CollapseChat from './CollapseChat'; import StreamAudio from './StreamAudio'; import StopButton from './StopButton'; import SendButton from './SendButton'; +import ContextTracker from './ContextTracker'; import EditBadges from './EditBadges'; import BadgeRow from './BadgeRow'; import Mention from './Mention'; @@ -48,6 +49,7 @@ interface ChatFormProps { setFilesLoading: React.Dispatch>; newConversation: ConvoGenerator; handleStopGenerating: (e: React.MouseEvent) => void; + getMessages: () => TMessage[] | undefined; } const ChatForm = memo(function ChatForm({ @@ -60,6 +62,7 @@ const ChatForm = memo(function ChatForm({ setFilesLoading, newConversation, handleStopGenerating, + getMessages, }: ChatFormProps) { const submitButtonRef = useRef(null); const textAreaRef = useRef(null); @@ -372,7 +375,10 @@ const ChatForm = memo(function ChatForm({ isSubmitting={isSubmitting} /> )} -
+
+ {endpoint && ( + + )} {isSubmitting && showStopButton ? ( ) : ( @@ -410,6 +416,7 @@ function ChatFormWrapper({ index = 0 }: { index?: number }) { setFilesLoading, newConversation, handleStopGenerating, + getMessages, } = useChatContext(); /** @@ -449,6 +456,10 @@ function ChatFormWrapper({ index = 0 }: { index?: number }) { [], ); + const getMessagesRef = useRef(getMessages); + getMessagesRef.current = getMessages; + const stableGetMessages = useCallback(() => getMessagesRef.current(), []); + return ( ); } diff --git a/client/src/components/Chat/Input/ContextTracker.tsx b/client/src/components/Chat/Input/ContextTracker.tsx new file mode 100644 index 0000000000..7db5767c28 --- /dev/null +++ b/client/src/components/Chat/Input/ContextTracker.tsx @@ -0,0 +1,198 @@ +import { useEffect, useMemo, useRef, useState } from 'react'; +import { useQueryClient } from '@tanstack/react-query'; +import { QueryKeys, Constants } from 'librechat-data-provider'; +import type { TConversation, TMessage, TModelSpec, TStartupConfig } from 'librechat-data-provider'; +import { useRecoilValue } from 'recoil'; +import { TooltipAnchor } from '@librechat/client'; +import { useGetStartupConfig } from '~/data-provider'; +import { useLocalize } from '~/hooks'; +import store from '~/store'; +import { cn } from '~/utils'; + +type ContextTrackerProps = { + conversation: TConversation | null; + getMessages: () => TMessage[] | undefined; + isSubmitting: boolean; +}; + +type MessageWithTokenCount = TMessage & { tokenCount?: number }; + +const TRACKER_SIZE = 24; +const TRACKER_STROKE = 2.5; + +const formatTokenCount = (count: number): string => { + if (count >= 1_000_000) { + return `${(count / 1_000_000).toFixed(1)}M`; + } + if (count >= 1_000) { + return `${(count / 1_000).toFixed(1)}K`; + } + return count.toString(); +}; + +const getUsedTokens = (messages: TMessage[] | undefined): number => { + if (!messages?.length) { + return 0; + } + + return messages.reduce((totalTokens, message) => { + const tokenCount = (message as MessageWithTokenCount).tokenCount; + if (typeof tokenCount !== 'number' || !Number.isFinite(tokenCount) || tokenCount <= 0) { + return totalTokens; + } + + return totalTokens + tokenCount; + }, 0); +}; + +const getSpecMaxContextTokens = ( + startupConfig: TStartupConfig | undefined, + specName: string | null | undefined, +): number | null => { + if (!specName) { + return null; + } + + const modelSpec = startupConfig?.modelSpecs?.list?.find( + (spec: TModelSpec) => spec.name === specName, + ); + const maxContextTokens = modelSpec?.preset?.maxContextTokens; + if ( + typeof maxContextTokens !== 'number' || + !Number.isFinite(maxContextTokens) || + maxContextTokens <= 0 + ) { + return null; + } + + return maxContextTokens; +}; + +export default function ContextTracker({ + conversation, + getMessages, + isSubmitting, +}: ContextTrackerProps) { + const localize = useLocalize(); + const queryClient = useQueryClient(); + const { data: startupConfig } = useGetStartupConfig(); + const showContextTracker = useRecoilValue(store.showContextTracker); + const conversationId = conversation?.conversationId ?? Constants.NEW_CONVO; + const [usedTokens, setUsedTokens] = useState(() => getUsedTokens(getMessages())); + + useEffect(() => { + setUsedTokens(getUsedTokens(getMessages())); + }, [conversationId, getMessages]); + + useEffect(() => { + const unsubscribe = queryClient.getQueryCache().subscribe((event) => { + const queryKey = event?.query?.queryKey; + if (!Array.isArray(queryKey) || queryKey[0] !== QueryKeys.messages) { + return; + } + + setUsedTokens(getUsedTokens(getMessages())); + }); + + return unsubscribe; + }, [getMessages, queryClient]); + + const prevIsSubmitting = useRef(isSubmitting); + useEffect(() => { + if (prevIsSubmitting.current && !isSubmitting) { + setUsedTokens(getUsedTokens(getMessages())); + } + prevIsSubmitting.current = isSubmitting; + }, [isSubmitting, getMessages]); + + const maxContextTokens = + typeof conversation?.maxContextTokens === 'number' && + Number.isFinite(conversation.maxContextTokens) && + conversation.maxContextTokens > 0 + ? conversation.maxContextTokens + : getSpecMaxContextTokens(startupConfig, conversation?.spec); + + const usageRatio = useMemo(() => { + if (maxContextTokens == null || maxContextTokens <= 0) { + return 0; + } + + return Math.min(usedTokens / maxContextTokens, 1); + }, [maxContextTokens, usedTokens]); + + const trackerRadius = useMemo(() => (TRACKER_SIZE - TRACKER_STROKE) / 2, []); + const circumference = useMemo(() => 2 * Math.PI * trackerRadius, [trackerRadius]); + const dashOffset = useMemo(() => circumference * (1 - usageRatio), [circumference, usageRatio]); + + let ringColorClass = 'text-text-primary'; + if (maxContextTokens == null) { + ringColorClass = 'text-text-secondary'; + } else if (usageRatio > 0.9) { + ringColorClass = 'text-red-500'; + } else if (usageRatio > 0.75) { + ringColorClass = 'text-yellow-500'; + } + + const tooltipDescription = useMemo(() => { + if (maxContextTokens == null) { + return localize('com_ui_context_usage_unknown_max', { + 0: formatTokenCount(usedTokens), + }); + } + + const percentage = (usageRatio * 100).toFixed(1) + '%'; + return localize('com_ui_context_usage_with_max', { + 0: percentage, + 1: formatTokenCount(usedTokens), + 2: formatTokenCount(maxContextTokens), + }); + }, [localize, maxContextTokens, usedTokens, usageRatio]); + + if (!showContextTracker) { + return null; + } + + return ( + + + + + + + } + /> + ); +} diff --git a/client/src/components/Nav/SettingsTabs/Chat/Chat.tsx b/client/src/components/Nav/SettingsTabs/Chat/Chat.tsx index 8d4f1817f3..00769c8ff0 100644 --- a/client/src/components/Nav/SettingsTabs/Chat/Chat.tsx +++ b/client/src/components/Nav/SettingsTabs/Chat/Chat.tsx @@ -84,6 +84,13 @@ const toggleSwitchConfigs = [ hoverCardText: 'com_nav_info_save_badges_state' as const, key: 'showBadges', }, + { + stateAtom: store.showContextTracker, + localizationKey: 'com_nav_show_context_tracker' as const, + switchId: 'showContextTracker', + hoverCardText: undefined, + key: 'showContextTracker', + }, { stateAtom: store.modularChat, localizationKey: 'com_nav_modular_chat' as const, diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 3d19f65ad6..55a6b8fefe 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -570,6 +570,7 @@ "com_nav_plus_command_description": "Toggle command \"+\" for adding a multi-response setting", "com_nav_profile_picture": "Profile Picture", "com_nav_save_badges_state": "Save badges state", + "com_nav_show_context_tracker": "Show context tracker", "com_nav_save_drafts": "Save drafts locally", "com_nav_scroll_button": "Scroll to the end button", "com_nav_search_placeholder": "Search messages", @@ -1522,6 +1523,9 @@ "com_ui_upload_provider": "Upload to Provider", "com_ui_upload_success": "Successfully uploaded file", "com_ui_upload_type": "Select Upload Type", + "com_ui_context_usage": "Context usage", + "com_ui_context_usage_unknown_max": "{{0}} tokens used (max unavailable)", + "com_ui_context_usage_with_max": "{{0}} ยท {{1}} / {{2}} context used", "com_ui_usage": "Usage", "com_ui_use_2fa_code": "Use 2FA Code Instead", "com_ui_use_backup_code": "Use Backup Code Instead", diff --git a/client/src/store/settings.ts b/client/src/store/settings.ts index 2a2796ad59..4d163a9621 100644 --- a/client/src/store/settings.ts +++ b/client/src/store/settings.ts @@ -39,6 +39,7 @@ const localStorageAtoms = { rememberDefaultFork: atomWithLocalStorage(LocalStorageKeys.REMEMBER_FORK_OPTION, false), showThinking: atomWithLocalStorage('showThinking', false), saveBadgesState: atomWithLocalStorage('saveBadgesState', false), + showContextTracker: atomWithLocalStorage('showContextTracker', true), // Beta features settings modularChat: atomWithLocalStorage('modularChat', true),