From 249795c70b4a35e7e570621154149c5fe9cd8118 Mon Sep 17 00:00:00 2001
From: sxjeru <sxjeru@gmail.com>
Date: Wed, 11 Sep 2024 21:53:41 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=92=84=20style:=20Optimize=20model=20toke?=
 =?UTF-8?q?n=20display=20method=20(#3697)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 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
---
 src/components/ModelSelect/index.tsx | 10 +------
 src/utils/format.test.ts             | 43 +++++++++++++++++++++++++++-
 src/utils/format.ts                  | 11 +++++++
 3 files changed, 54 insertions(+), 10 deletions(-)

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`;
+};