Context tracker init

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<React.SetStateAction<boolean>>;
   newConversation: ConvoGenerator;
   handleStopGenerating: (e: React.MouseEvent<HTMLButtonElement>) => 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<HTMLButtonElement>(null);
   const textAreaRef = useRef<HTMLTextAreaElement>(null);
@@ -372,7 +375,10 @@ const ChatForm = memo(function ChatForm({
                   isSubmitting={isSubmitting}
                 />
               )}
-              <div className={`${isRTL ? 'ml-2' : 'mr-2'}`}>
+              <div className={cn('flex items-center gap-2', isRTL ? 'ml-2' : 'mr-2')}>
+                {endpoint && (
+                  <ContextTracker conversation={conversation} getMessages={getMessages} isSubmitting={isSubmitting} />
+                )}
                 {isSubmitting && showStopButton ? (
                   <StopButton stop={handleStopGenerating} setShowStopButton={setShowStopButton} />
                 ) : (
@@ -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 (
     <ChatForm
       index={index}
@@ -460,6 +471,7 @@ function ChatFormWrapper({ index = 0 }: { index?: number }) {
       setFilesLoading={setFilesLoading}
       newConversation={stableNewConversation}
       handleStopGenerating={stableHandleStop}
+      getMessages={stableGetMessages}
     />
   );
 }

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<number>(() => 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 (
+    <TooltipAnchor
+      description={tooltipDescription}
+      side="top"
+      render={
+        <button
+          type="button"
+          aria-label={localize('com_ui_context_usage')}
+          className={cn('rounded-full p-1.5', ringColorClass)}
+          data-testid="context-tracker"
+        >
+          <svg
+            width={TRACKER_SIZE}
+            height={TRACKER_SIZE}
+            viewBox={`0 0 ${TRACKER_SIZE} ${TRACKER_SIZE}`}
+          >
+            <circle
+              cx={TRACKER_SIZE / 2}
+              cy={TRACKER_SIZE / 2}
+              r={trackerRadius}
+              stroke="currentColor"
+              strokeWidth={TRACKER_STROKE}
+              fill="transparent"
+              className="opacity-20"
+            />
+            <circle
+              cx={TRACKER_SIZE / 2}
+              cy={TRACKER_SIZE / 2}
+              r={trackerRadius}
+              stroke="currentColor"
+              strokeWidth={TRACKER_STROKE}
+              strokeLinecap="round"
+              strokeDasharray={`${circumference} ${circumference}`}
+              strokeDashoffset={dashOffset}
+              fill="transparent"
+              transform={`rotate(-90 ${TRACKER_SIZE / 2} ${TRACKER_SIZE / 2})`}
+              className="transition-all duration-200"
+            />
+          </svg>
+        </button>
+      }
+    />
+  );
+}

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,

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

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