✨ feat: 优化使用 Token 数量计算与显示

lobehub · May 2, 2024 · b549b4d · b549b4d
1 parent fad7269
commit b549b4d
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/src/features/Actions/TokenMini.tsx b/src/features/Actions/TokenMini.tsx
@@ -0,0 +1,22 @@
+import { Tooltip, Typography } from 'antd';
+import { isEqual } from 'lodash-es';
+
+import { OPENAI_MODEL_LIST } from '@/constants/openai';
+import { useCalculateToken } from '@/hooks/useCalculateToken';
+import { configSelectors, useConfigStore } from '@/store/config';
+
+const TokenMini = () => {
+  const config = useConfigStore((s) => configSelectors.currentOpenAIConfig(s), isEqual);
+  const usedTokens = useCalculateToken();
+  const maxValue = OPENAI_MODEL_LIST.find((item) => item.name === config?.model)?.maxToken || 4096;
+
+  return (
+    <Tooltip title="上下文消耗 Token 数量计算，数字代表 {总计使用} / {总计可用}]">
+      <Typography.Text type={'secondary'}>
+        {usedTokens} / {maxValue}
+      </Typography.Text>
+    </Tooltip>
+  );
+};
+
+export default TokenMini;
diff --git a/src/features/ChatInfo/index.tsx b/src/features/ChatInfo/index.tsx
@@ -7,7 +7,6 @@ import { Flexbox } from 'react-layout-kit';
 
 import AgentCard from '@/components/agent/AgentCard';
 import { SIDEBAR_WIDTH } from '@/constants/common';
-import Token from '@/features/Actions/Token';
 import { sessionSelectors, useSessionStore } from '@/store/session';
 
 import Operations from './Operations';
@@ -36,7 +35,7 @@ const Header = () => {
       mode={'fixed'}
       placement={'right'}
     >
-      <AgentCard agent={currentAgent} actions={[<Token key="token" />]} />
+      <AgentCard agent={currentAgent} />
       <Flexbox gap={8} style={{ padding: 8 }}>
         <Operations />
       </Flexbox>

diff --git a/src/features/ChatInput/MessageInput/index.tsx b/src/features/ChatInput/MessageInput/index.tsx
@@ -6,6 +6,7 @@ import React, { memo, useRef } from 'react';
 import { Flexbox } from 'react-layout-kit';
 
 import Record from '@/features/Actions/Record';
+import TokenMini from '@/features/Actions/TokenMini';
 import useChatInput from '@/hooks/useSendMessage';
 import { useSessionStore } from '@/store/session';
 import { isCommandPressed } from '@/utils/keyboard';
@@ -58,6 +59,7 @@ const InputArea = memo((props: InputAreaProps) => {
         }}
         placeholder="请输入内容开始聊天"
         ref={ref}
+        suffix={<TokenMini />}
         type={'block'}
         value={messageInput}
       />