Skip to content

Commit 7b9f490

Browse files
authored
Update TokenTag.tsx
1 parent de8df94 commit 7b9f490

File tree

1 file changed

+13
-3
lines changed

1 file changed

+13
-3
lines changed

src/features/ChatInput/ActionBar/Token/TokenTag.tsx

+13-3
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { TokenTag, Tooltip } from '@lobehub/ui';
22
import { Popover } from 'antd';
33
import { useTheme } from 'antd-style';
44
import numeral from 'numeral';
5-
import { memo } from 'react';
5+
import { memo, useMemo } from 'react';
66
import { useTranslation } from 'react-i18next';
77
import { Center, Flexbox } from 'react-layout-kit';
88

@@ -12,7 +12,7 @@ import { useTokenCount } from '@/hooks/useTokenCount';
1212
import { useAgentStore } from '@/store/agent';
1313
import { agentChatConfigSelectors, agentSelectors } from '@/store/agent/selectors';
1414
import { useChatStore } from '@/store/chat';
15-
import { topicSelectors } from '@/store/chat/selectors';
15+
import { chatSelectors, topicSelectors } from '@/store/chat/selectors';
1616
import { useToolStore } from '@/store/tool';
1717
import { toolSelectors } from '@/store/tool/selectors';
1818

@@ -41,6 +41,11 @@ const Token = memo<TokenTagProps>(({ total: messageString }) => {
4141
];
4242
});
4343

44+
const [historyCount, enableHistoryCount] = useAgentStore((s) => [
45+
agentChatConfigSelectors.historyCount(s),
46+
agentChatConfigSelectors.enableHistoryCount(s),
47+
]);
48+
4449
const maxTokens = useModelContextWindowTokens(model, provider);
4550

4651
// Tool usage token
@@ -60,7 +65,12 @@ const Token = memo<TokenTagProps>(({ total: messageString }) => {
6065
// Chat usage token
6166
const inputTokenCount = useTokenCount(input);
6267

63-
const chatsToken = useTokenCount(messageString) + inputTokenCount;
68+
const chatsString = useMemo(() => {
69+
const chats = chatSelectors.mainAIChatsWithHistoryConfig(useChatStore.getState());
70+
return chats.map(chat => chat.content).join('');
71+
}, [messageString, historyCount, enableHistoryCount]);
72+
73+
const chatsToken = useTokenCount(chatsString) + inputTokenCount;
6474

6575
// SystemRole token
6676
const systemRoleToken = useTokenCount(systemRole);

0 commit comments

Comments
 (0)