diff --git a/src/components/ModelSelect/index.tsx b/src/components/ModelSelect/index.tsx index 12b61d5ee1a8..c963e3df5660 100644 --- a/src/components/ModelSelect/index.tsx +++ b/src/components/ModelSelect/index.tsx @@ -10,6 +10,7 @@ import { useTranslation } from 'react-i18next'; import { Center, Flexbox } from 'react-layout-kit'; import { ChatModelCard } from '@/types/llm'; +import { formatTokenNumber } from '@/utils/format'; const useStyles = createStyles(({ css, token }) => ({ custom: css` @@ -55,15 +56,6 @@ const useStyles = createStyles(({ css, token }) => ({ border-radius: 4px; `, })); -const formatTokenNumber = (num: number): string => { - if (num > 0 && num < 1024) return '1K'; - - let kiloToken = Math.floor(num / 1024); - if (num >= 128_000 && num < 1_024_000) { - kiloToken = Math.floor(num / 1000); - } - return kiloToken < 1000 ? `${kiloToken}K` : `${Math.floor(kiloToken / 1000)}M`; -}; interface ModelInfoTagsProps extends ChatModelCard { directionReverse?: boolean; diff --git a/src/utils/format.test.ts b/src/utils/format.test.ts index 58c1daa9c808..d70a9a296c61 100644 --- a/src/utils/format.test.ts +++ b/src/utils/format.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest'; -import { formatSize, formatSpeed, formatTime } from './format'; +import { formatSize, formatSpeed, formatTime, formatTokenNumber } from './format'; describe('formatSize', () => { it('should format bytes to KB correctly', () => { @@ -73,3 +73,44 @@ describe('formatTime', () => { expect(formatTime(3599.99)).toBe('60.0 min'); }); }); + +describe('formatTokenNumber', () => { + it('should return "1K" for numbers between 1 and 1023', () => { + expect(formatTokenNumber(500)).toBe('1K'); + expect(formatTokenNumber(1000)).toBe('1K'); + }); + + it('should format numbers between 1024 and 41,983 correctly', () => { + expect(formatTokenNumber(1024)).toBe('1K'); + expect(formatTokenNumber(2000)).toBe('2K'); + expect(formatTokenNumber(2048)).toBe('2K'); + expect(formatTokenNumber(4000)).toBe('4K'); + expect(formatTokenNumber(4096)).toBe('4K'); + expect(formatTokenNumber(32000)).toBe('32K'); + expect(formatTokenNumber(65536)).toBe('64K'); + }); + + it('should format numbers between 41,984 and 127,999 correctly', () => { + expect(formatTokenNumber(41984)).toBe('41K'); + expect(formatTokenNumber(100000)).toBe('97K'); + expect(formatTokenNumber(127999)).toBe('124K'); + }); + + it('should return "128K" for 131,072', () => { + expect(formatTokenNumber(131072)).toBe('128K'); // Qwen + }); + + it('should format numbers between 128,000 and 999,999 correctly', () => { + expect(formatTokenNumber(128000)).toBe('128K'); + expect(formatTokenNumber(200000)).toBe('200K'); // Claude + expect(formatTokenNumber(999999)).toBe('999K'); + }); + + it('should format numbers 1,000,000 and above correctly', () => { + expect(formatTokenNumber(1000000)).toBe('1M'); + expect(formatTokenNumber(1024000)).toBe('1M'); + expect(formatTokenNumber(1048576)).toBe('1M'); // Gemini Flash + expect(formatTokenNumber(2000000)).toBe('2M'); + expect(formatTokenNumber(2097152)).toBe('2M'); // Gemini Pro + }); +}); diff --git a/src/utils/format.ts b/src/utils/format.ts index 050919f25cc8..b3bd551dc0a9 100644 --- a/src/utils/format.ts +++ b/src/utils/format.ts @@ -54,3 +54,14 @@ export const formatTime = (timeInSeconds: number): string => { export const formatNumber = (num: any) => { return new Intl.NumberFormat('en-US').format(num); }; + +export const formatTokenNumber = (num: number): string => { + if (num > 0 && num < 1024) return '1K'; + + let kiloToken = Math.floor(num / 1024); + if (num >= 1024 && num < 1024 * 41 || num >= 128_000) { + kiloToken = Math.floor(num / 1000); + } + if (num === 131_072) return '128K'; + return kiloToken < 1000 ? `${kiloToken}K` : `${Math.floor(kiloToken / 1000)}M`; +};