💄 style: Optimize model token display method (lobehub#3697)

* Update index.tsx * Update index.tsx * Update index.tsx * Update index.tsx * Update index.tsx * Update format.ts * Update format.test.ts * Update format.test.ts * Update format.ts * Update format.test.ts * Update format.test.ts
LovelyGuYiMeng · Sep 11, 2024 · 249795c · 249795c
1 parent 0610253
commit 249795c
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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`;
+};