diff --git a/client/src/components/Chat/Input/TokenUsageIndicator.tsx b/client/src/components/Chat/Input/TokenUsageIndicator.tsx index 1e8658bfdd..3008244d56 100644 --- a/client/src/components/Chat/Input/TokenUsageIndicator.tsx +++ b/client/src/components/Chat/Input/TokenUsageIndicator.tsx @@ -4,13 +4,10 @@ import { useLocalize, useTokenUsage } from '~/hooks'; import { cn } from '~/utils'; function formatTokens(n: number): string { - if (n >= 1000000) { - return `${(n / 1000000).toFixed(1).replace(/\.0$/, '')}M`; - } - if (n >= 1000) { - return `${(n / 1000).toFixed(1).replace(/\.0$/, '')}K`; - } - return n.toString(); + return new Intl.NumberFormat(undefined, { + notation: 'compact', + maximumFractionDigits: 1, + }).format(n); } interface ProgressBarProps { @@ -19,30 +16,48 @@ interface ProgressBarProps { colorClass: string; label: string; showPercentage?: boolean; + indeterminate?: boolean; } -function ProgressBar({ value, max, colorClass, label, showPercentage = false }: ProgressBarProps) { +function ProgressBar({ + value, + max, + colorClass, + label, + showPercentage = false, + indeterminate = false, +}: ProgressBarProps) { const percentage = max > 0 ? Math.min((value / max) * 100, 100) : 0; return (
-
+ {indeterminate ? (
-
-
+ ) : ( +
+
+
+
+ )}
- {showPercentage && ( + {showPercentage && !indeterminate && ( @@ -80,7 +95,7 @@ function TokenRow({ label, value, total, colorClass, ariaLabel }: TokenRowProps) function TokenUsageContent() { const localize = useLocalize(); - const { inputTokens, outputTokens, maxContext } = useTokenUsage(); + const { inputTokens = 0, outputTokens = 0, maxContext = null } = useTokenUsage() ?? {}; const totalUsed = inputTokens + outputTokens; const hasMaxContext = maxContext !== null && maxContext > 0; @@ -127,20 +142,23 @@ function TokenUsageContent() {
{/* Main Progress Bar */} - {hasMaxContext && ( -
- - +
+ + - )} +
{/* Divider */}
@@ -168,7 +186,7 @@ function TokenUsageContent() { const TokenUsageIndicator = memo(function TokenUsageIndicator() { const localize = useLocalize(); - const { inputTokens, outputTokens, maxContext } = useTokenUsage(); + const { inputTokens = 0, outputTokens = 0, maxContext = null } = useTokenUsage() ?? {}; const totalUsed = inputTokens + outputTokens; const hasMaxContext = maxContext !== null && maxContext > 0; @@ -249,13 +267,7 @@ const TokenUsageIndicator = memo(function TokenUsageIndicator() { - + diff --git a/packages/data-provider/specs/tokens.spec.ts b/packages/data-provider/specs/tokens.spec.ts new file mode 100644 index 0000000000..37eeecbea6 --- /dev/null +++ b/packages/data-provider/specs/tokens.spec.ts @@ -0,0 +1,152 @@ +import { + findMatchingPattern, + getModelMaxTokens, + getModelMaxOutputTokens, + matchModelName, + maxTokensMap, +} from '../src/tokens'; +import { EModelEndpoint } from '../src/schemas'; + +describe('Token Pattern Matching', () => { + describe('findMatchingPattern', () => { + const testMap: Record = { + 'claude-': 100000, + 'claude-3': 200000, + 'claude-3-opus': 200000, + 'gpt-4': 8000, + 'gpt-4-turbo': 128000, + }; + + it('should match exact model names', () => { + expect(findMatchingPattern('claude-3-opus', testMap)).toBe('claude-3-opus'); + expect(findMatchingPattern('gpt-4-turbo', testMap)).toBe('gpt-4-turbo'); + }); + + it('should match more specific patterns first (reverse order)', () => { + // claude-3-opus-20240229 should match 'claude-3-opus' not 'claude-3' or 'claude-' + expect(findMatchingPattern('claude-3-opus-20240229', testMap)).toBe('claude-3-opus'); + }); + + it('should fall back to broader patterns when no specific match', () => { + // claude-3-haiku should match 'claude-3' (not 'claude-3-opus') + expect(findMatchingPattern('claude-3-haiku', testMap)).toBe('claude-3'); + }); + + it('should be case-insensitive', () => { + expect(findMatchingPattern('Claude-3-Opus', testMap)).toBe('claude-3-opus'); + expect(findMatchingPattern('GPT-4-TURBO', testMap)).toBe('gpt-4-turbo'); + }); + + it('should return null for unmatched models', () => { + expect(findMatchingPattern('unknown-model', testMap)).toBeNull(); + expect(findMatchingPattern('llama-2', testMap)).toBeNull(); + }); + + it('should NOT match when pattern appears in middle of model name (startsWith behavior)', () => { + // This is the key fix: "my-claude-wrapper" should NOT match "claude-" + expect(findMatchingPattern('my-claude-wrapper', testMap)).toBeNull(); + expect(findMatchingPattern('openai-gpt-4-proxy', testMap)).toBeNull(); + expect(findMatchingPattern('custom-claude-3-service', testMap)).toBeNull(); + }); + + it('should handle empty string model name', () => { + expect(findMatchingPattern('', testMap)).toBeNull(); + }); + + it('should handle empty tokens map', () => { + expect(findMatchingPattern('claude-3', {})).toBeNull(); + }); + }); + + describe('getModelMaxTokens', () => { + it('should return exact match tokens', () => { + expect(getModelMaxTokens('gpt-4o', EModelEndpoint.openAI)).toBe(127500); + expect(getModelMaxTokens('claude-3-opus', EModelEndpoint.anthropic)).toBe(200000); + }); + + it('should return pattern-matched tokens', () => { + // claude-3-opus-20240229 should match claude-3-opus pattern + expect(getModelMaxTokens('claude-3-opus-20240229', EModelEndpoint.anthropic)).toBe(200000); + }); + + it('should return undefined for unknown models', () => { + expect(getModelMaxTokens('completely-unknown-model', EModelEndpoint.openAI)).toBeUndefined(); + }); + + it('should fall back to openAI for unknown endpoints', () => { + const result = getModelMaxTokens('gpt-4o', 'unknown-endpoint'); + expect(result).toBe(127500); + }); + + it('should handle non-string input gracefully', () => { + expect(getModelMaxTokens(null as unknown as string)).toBeUndefined(); + expect(getModelMaxTokens(undefined as unknown as string)).toBeUndefined(); + expect(getModelMaxTokens(123 as unknown as string)).toBeUndefined(); + }); + + it('should NOT match model names with pattern in middle', () => { + // A model like "my-gpt-4-wrapper" should not match "gpt-4" + expect(getModelMaxTokens('my-gpt-4-wrapper', EModelEndpoint.openAI)).toBeUndefined(); + }); + }); + + describe('getModelMaxOutputTokens', () => { + it('should return exact match output tokens', () => { + expect(getModelMaxOutputTokens('o1', EModelEndpoint.openAI)).toBe(32268); + expect(getModelMaxOutputTokens('claude-3-opus', EModelEndpoint.anthropic)).toBe(4096); + }); + + it('should return pattern-matched output tokens', () => { + expect(getModelMaxOutputTokens('claude-3-opus-20240229', EModelEndpoint.anthropic)).toBe( + 4096, + ); + }); + + it('should return system_default for unknown models (openAI endpoint)', () => { + expect(getModelMaxOutputTokens('unknown-model', EModelEndpoint.openAI)).toBe(32000); + }); + + it('should handle non-string input gracefully', () => { + expect(getModelMaxOutputTokens(null as unknown as string)).toBeUndefined(); + expect(getModelMaxOutputTokens(undefined as unknown as string)).toBeUndefined(); + }); + }); + + describe('matchModelName', () => { + it('should return exact match model name', () => { + expect(matchModelName('gpt-4o', EModelEndpoint.openAI)).toBe('gpt-4o'); + }); + + it('should return pattern key for pattern matches', () => { + expect(matchModelName('claude-3-opus-20240229', EModelEndpoint.anthropic)).toBe( + 'claude-3-opus', + ); + }); + + it('should return input for unknown models', () => { + expect(matchModelName('unknown-model', EModelEndpoint.openAI)).toBe('unknown-model'); + }); + + it('should handle non-string input gracefully', () => { + expect(matchModelName(null as unknown as string)).toBeUndefined(); + }); + }); + + describe('maxTokensMap structure', () => { + it('should have entries for all major endpoints', () => { + expect(maxTokensMap[EModelEndpoint.openAI]).toBeDefined(); + expect(maxTokensMap[EModelEndpoint.anthropic]).toBeDefined(); + expect(maxTokensMap[EModelEndpoint.google]).toBeDefined(); + expect(maxTokensMap[EModelEndpoint.azureOpenAI]).toBeDefined(); + expect(maxTokensMap[EModelEndpoint.bedrock]).toBeDefined(); + }); + + it('should have positive token values', () => { + Object.values(maxTokensMap).forEach((endpointMap) => { + Object.entries(endpointMap).forEach(([model, tokens]) => { + expect(tokens).toBeGreaterThan(0); + }); + }); + }); + }); +}); diff --git a/packages/data-provider/src/tokens.ts b/packages/data-provider/src/tokens.ts index c5bbbb233b..f5c6d6eedc 100644 --- a/packages/data-provider/src/tokens.ts +++ b/packages/data-provider/src/tokens.ts @@ -313,6 +313,10 @@ export const maxTokensMap: Record> = { /** * Finds the first matching pattern in the tokens map. * Searches in reverse order to match more specific patterns first. + * + * Note: This relies on the insertion order of keys in the tokensMap object. + * More specific patterns must be defined later in the object to be matched first. + * If the order of keys is changed, the matching behavior may be affected. */ export function findMatchingPattern( modelName: string, @@ -322,7 +326,7 @@ export function findMatchingPattern( const lowerModelName = modelName.toLowerCase(); for (let i = keys.length - 1; i >= 0; i--) { const modelKey = keys[i]; - if (lowerModelName.includes(modelKey)) { + if (lowerModelName.startsWith(modelKey)) { return modelKey; } } @@ -510,7 +514,6 @@ export function getModelMaxOutputTokens( /** * Centralized token-related default values. - * These replace hardcoded magic numbers throughout the codebase. */ export const TOKEN_DEFAULTS = { /** Fallback context window for agents when model lookup fails */