From 614ec4b704517784acca54cc11d7f78f03968f7c Mon Sep 17 00:00:00 2001 From: arvinxx Date: Fri, 8 Sep 2023 21:44:17 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=A8=20chore:=20refactor=20token=20coun?= =?UTF-8?q?t=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../session/slices/chat/selectors/token.ts | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/store/session/slices/chat/selectors/token.ts b/src/store/session/slices/chat/selectors/token.ts index b379aeefcf8a3..b6ef52c83b507 100644 --- a/src/store/session/slices/chat/selectors/token.ts +++ b/src/store/session/slices/chat/selectors/token.ts @@ -5,25 +5,24 @@ import { agentSelectors } from '@/store/session/slices/agentConfig'; import type { SessionStore } from '../../../store'; import { currentChatsWithHistoryConfig } from './chat'; +// TODO: We should find a smaller package to do this. gpt-tokenizer has 2M Size +const count = (text: string) => encode(text).length; + export const systemRoleSel = (s: SessionStore): string => { const config = agentSelectors.currentAgentConfig(s); return config.systemRole; }; -const systemRoleTokens = (s: SessionStore): number[] => { +export const systemRoleTokenCount = (s: SessionStore): number => { const systemRole = systemRoleSel(s); - return encode(systemRole || ''); + return count(systemRole || ''); }; -const chatsTokens = (s: SessionStore): number[] => { +export const chatsTokenCount = (s: SessionStore): number => { const chats = currentChatsWithHistoryConfig(s); - return encode(chats.map((m) => m.content).join('')); + return count(chats.map((m) => m.content).join('')); }; -export const chatsTokenCount = (s: SessionStore) => chatsTokens(s).length; - -export const systemRoleTokenCount = (s: SessionStore) => systemRoleTokens(s).length; - -export const totalTokenCount = (s: SessionStore) => chatsTokens(s).length + systemRoleTokenCount(s); +export const totalTokenCount = (s: SessionStore) => chatsTokenCount(s) + systemRoleTokenCount(s);