From 9c61d73076052f7eeee16652d25736eff760fe33 Mon Sep 17 00:00:00 2001
From: Marco Beretta <81851188+berry-13@users.noreply.github.com>
Date: Tue, 16 Dec 2025 00:20:57 +0100
Subject: [PATCH] feat: enhance token usage visualization and matching logic in
TokenUsageIndicator and tokens module
---
.../Chat/Input/TokenUsageIndicator.tsx | 86 +++++-----
packages/data-provider/specs/tokens.spec.ts | 152 ++++++++++++++++++
packages/data-provider/src/tokens.ts | 7 +-
3 files changed, 206 insertions(+), 39 deletions(-)
create mode 100644 packages/data-provider/specs/tokens.spec.ts
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 (
- {showPercentage && (
+ {showPercentage && !indeterminate && (
{Math.round(percentage)}%
@@ -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 && (
-
-
-
- {formatTokens(totalUsed)}
- {formatTokens(maxContext)}
-
+
+
+
+ {formatTokens(totalUsed)}
+ {hasMaxContext ? formatTokens(maxContext) : 'N/A'}
- )}
+
{/* 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 */