From 68325c1d9167af0a654d1ca40b4c423da1e06dbc Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Tue, 8 Jul 2025 09:58:30 -0500 Subject: [PATCH 1/2] feat: implement accurate token-based counting with progressive warnings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace rough token estimation (text.length/4) with accurate gpt-tokenizer library - Add model-specific token limits (Llama/Gemma: 70k, DeepSeek: 64k, default: 64k) - Implement progressive warning thresholds: - 50%: "Tip" to compress chat - 95%: "Warning" with compress option available - 99%: "Error" with disabled submit button (textarea remains editable) - Add debouncing (300ms) for token calculations to prevent typing lag - Optimize token calculations with useMemo to avoid unnecessary recalculations - Extract shared calculateTotalTokens utility function - Update UI styling for different severity levels - Reorder TokenWarning to appear directly above chatbox This is a reimplementation of PR #135 to resolve merge conflicts on the current master branch. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- frontend/bun.lock | 3 + frontend/package.json | 1 + frontend/src/components/ChatBox.tsx | 161 +++++++++++++++------- frontend/src/components/ModelSelector.tsx | 23 +++- 4 files changed, 130 insertions(+), 58 deletions(-) diff --git a/frontend/bun.lock b/frontend/bun.lock index b864a2535..4a8acd536 100644 --- a/frontend/bun.lock +++ b/frontend/bun.lock @@ -23,6 +23,7 @@ "@tauri-apps/plugin-os": "^2.2.1", "class-variance-authority": "^0.7.0", "clsx": "^2.1.1", + "gpt-tokenizer": "^3.0.1", "lucide-react": "^0.436.0", "openai": "^4.56.1", "react": "^18.3.1", @@ -709,6 +710,8 @@ "gopd": ["gopd@1.2.0", "", {}, "sha512-ZUKRh6/kUFoAiTAtTYPZJ3hw9wNxx+BIBOijnlG9PnrJsCcSjs1wyyD6vJpaYtgnzDrKYRSqf3OO6Rfa93xsRg=="], + "gpt-tokenizer": ["gpt-tokenizer@3.0.1", "", {}, "sha512-5jdaspBq/w4sWw322SvQj1Fku+CN4OAfYZeeEg8U7CWtxBz+zkxZ3h0YOHD43ee+nZYZ5Ud70HRN0ANcdIj4qg=="], + "graphemer": ["graphemer@1.4.0", "", {}, "sha512-EtKwoO6kxCL9WO5xipiHTZlSzBm7WLT627TqC/uVRd0HKmq8NXyebnNYxDoBi7wt8eTWrUrKXCOVaFq9x1kgag=="], "has-flag": ["has-flag@4.0.0", "", {}, "sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ=="], diff --git a/frontend/package.json b/frontend/package.json index 2d8f2436f..3a81dc345 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -35,6 +35,7 @@ "@tauri-apps/plugin-os": "^2.2.1", "class-variance-authority": "^0.7.0", "clsx": "^2.1.1", + "gpt-tokenizer": "^3.0.1", "lucide-react": "^0.436.0", "openai": "^4.56.1", "react": "^18.3.1", diff --git a/frontend/src/components/ChatBox.tsx b/frontend/src/components/ChatBox.tsx index 88cba3044..0c25d64e5 100644 --- a/frontend/src/components/ChatBox.tsx +++ b/frontend/src/components/ChatBox.tsx @@ -7,18 +7,18 @@ import { DropdownMenuItem, DropdownMenuTrigger } from "@/components/ui/dropdown-menu"; -import { useEffect, useRef, useState } from "react"; +import { useEffect, useRef, useState, useMemo } from "react"; import { useLocalState } from "@/state/useLocalState"; import { cn, useIsMobile } from "@/utils/utils"; import { useQuery } from "@tanstack/react-query"; import { getBillingService } from "@/billing/billingService"; -import { BillingStatus } from "@/billing/billingApi"; import { Route as ChatRoute } from "@/routes/_auth.chat.$chatId"; import { ChatMessage } from "@/state/LocalStateContext"; import { useNavigate, useRouter } from "@tanstack/react-router"; -import { ModelSelector, MODEL_CONFIG } from "@/components/ModelSelector"; +import { ModelSelector, MODEL_CONFIG, getModelTokenLimit } from "@/components/ModelSelector"; import { useOpenSecret } from "@opensecret/react"; import type { DocumentResponse } from "@opensecret/react"; +import { encode } from "gpt-tokenizer"; interface ParsedDocument { document: { @@ -35,30 +35,15 @@ interface ParsedDocument { timings: Record; } -// Rough token estimation function +// Accurate token counting using gpt-tokenizer function estimateTokenCount(text: string): number { - // A very rough estimation: ~4 characters per token on average - return Math.ceil(text.length / 4); + // Use gpt-tokenizer for accurate token counting + return encode(text).length; } -function TokenWarning({ - messages, - currentInput, - chatId, - className, - billingStatus, - onCompress, - isCompressing = false -}: { - messages: ChatMessage[]; - currentInput: string; - chatId?: string; - className?: string; - billingStatus?: BillingStatus; - onCompress?: () => void; - isCompressing?: boolean; -}) { - const totalTokens = +// Calculate total tokens for messages and current input +function calculateTotalTokens(messages: ChatMessage[], currentInput: string): number { + return ( messages.reduce((acc, msg) => { if (typeof msg.content === "string") { return acc + estimateTokenCount(msg.content); @@ -75,22 +60,48 @@ function TokenWarning({ }, 0) ); } - }, 0) + (currentInput ? estimateTokenCount(currentInput) : 0); + }, 0) + (currentInput ? estimateTokenCount(currentInput) : 0) + ); +} - const navigate = useNavigate(); +// Custom hook for debouncing values +function useDebounce(value: T, delay: number): T { + const [debouncedValue, setDebouncedValue] = useState(value); - // Check if user is on starter plan - const isStarter = billingStatus?.product_name?.toLowerCase().includes("starter") || false; + useEffect(() => { + const handler = setTimeout(() => { + setDebouncedValue(value); + }, delay); - // Token thresholds for different plan types - const STARTER_WARNING_THRESHOLD = 4000; - const PRO_WARNING_THRESHOLD = 10000; + return () => { + clearTimeout(handler); + }; + }, [value, delay]); - // Different thresholds for starter vs pro users - const warningThreshold = isStarter ? STARTER_WARNING_THRESHOLD : PRO_WARNING_THRESHOLD; + return debouncedValue; +} - // Only show warning if above the threshold - if (totalTokens < warningThreshold) return null; +function TokenWarning({ + chatId, + className, + onCompress, + isCompressing = false, + tokenPercentage +}: { + chatId?: string; + className?: string; + onCompress?: () => void; + isCompressing?: boolean; + tokenPercentage: number; +}) { + const navigate = useNavigate(); + + // Only show warning if above 50% + if (tokenPercentage < 50) return null; + + // Determine the severity and behavior based on percentage + const isAt95Percent = tokenPercentage >= 95; + const isAt99Percent = tokenPercentage >= 99; const handleNewChat = async (e: React.MouseEvent) => { e.preventDefault(); @@ -103,7 +114,17 @@ function TokenWarning({ } }; - // Determine button text based on compression state + // Get appropriate message and styling based on threshold + const getMessage = () => { + if (isAt99Percent) { + return "This chat is too long to continue."; + } else if (isAt95Percent) { + return "Chat is at capacity. Compress to continue."; + } else { + return "This chat is getting long. Compress it to save tokens."; + } + }; + const getButtonText = () => { if (isCompressing) { return { desktop: "Compressing...", mobile: "Compressing..." }; @@ -116,26 +137,46 @@ function TokenWarning({ const buttonText = getButtonText(); + // Determine background color based on severity + const bgClass = isAt99Percent + ? "bg-destructive/20 border border-destructive/30" + : isAt95Percent + ? "bg-warning/20 border border-warning/30" + : "bg-muted/50"; + return (
- Tip: - This chat is getting long. Compress it to save tokens. + + {isAt99Percent ? "Error:" : isAt95Percent ? "Warning:" : "Tip:"} + + {getMessage()}
- {chatId && ( + {chatId && !isAt99Percent && (