diff --git a/src/hooks/useSubmit.ts b/src/hooks/useSubmit.ts index 52e84d1..a013d0b 100644 --- a/src/hooks/useSubmit.ts +++ b/src/hooks/useSubmit.ts @@ -4,6 +4,7 @@ import { ChatInterface } from '@type/chat'; import { getChatCompletionStream as getChatCompletionStreamFree } from '@api/freeApi'; import { getChatCompletionStream as getChatCompletionStreamCustom } from '@api/customApi'; import { parseEventSource } from '@api/helper'; +import { limitMessageTokens } from '@utils/messageUtils'; const useSubmit = () => { const error = useStore((state) => state.error); @@ -31,15 +32,21 @@ const useSubmit = () => { let stream; try { + const messages = limitMessageTokens( + chats[currentChatIndex].messages, + 4000 + ); + if (messages.length === 0) throw new Error('Message exceed max token!'); + if (apiFree) { stream = await getChatCompletionStreamFree( - chats[currentChatIndex].messages, + messages, chats[currentChatIndex].config ); } else if (apiKey) { stream = await getChatCompletionStreamCustom( apiKey, - chats[currentChatIndex].messages, + messages, chats[currentChatIndex].config ); } diff --git a/src/utils/messageUtils.ts b/src/utils/messageUtils.ts new file mode 100644 index 0000000..cef477a --- /dev/null +++ b/src/utils/messageUtils.ts @@ -0,0 +1,19 @@ +import { MessageInterface } from '@type/chat'; +import countTokens from './countTokens'; + +export const limitMessageTokens = ( + messages: MessageInterface[], + limit: number = 4096 +): MessageInterface[] => { + const limitedMessages: MessageInterface[] = []; + let tokenCount = 0; + + for (let i = messages.length - 1; i >= 0; i--) { + const count = countTokens(messages[i].content); + if (count + tokenCount > limit) break; + tokenCount += count; + limitedMessages.unshift({ ...messages[i] }); + } + + return limitedMessages; +}; diff --git a/tsconfig.json b/tsconfig.json index 1932a6d..77430d0 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -23,7 +23,8 @@ "@hooks/*": ["./src/hooks/*"], "@constants/*": ["./src/constants/*"], "@api/*": ["./src/api/*"], - "@components/*": ["./src/components/*"] + "@components/*": ["./src/components/*"], + "@utils/*": ["./src/utils/*"] } }, "include": ["src"], diff --git a/vite.config.ts b/vite.config.ts index 02fb56e..fe439a6 100644 --- a/vite.config.ts +++ b/vite.config.ts @@ -13,6 +13,7 @@ export default defineConfig({ '@constants/': new URL('./src/constants/', import.meta.url).pathname, '@api/': new URL('./src/api/', import.meta.url).pathname, '@components/': new URL('./src/components/', import.meta.url).pathname, + '@utils/': new URL('./src/utils/', import.meta.url).pathname, }, }, });