From c006ccd97aaa1b3f08525f5ebcc99a7fee4c3f97 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Mon, 6 Mar 2023 22:50:03 +0800
Subject: [PATCH] feat: adaptive memory

fixes 9, fixed 10
---
 src/hooks/useSubmit.ts    | 11 +++++++++--
 src/utils/messageUtils.ts | 19 +++++++++++++++++++
 tsconfig.json             |  3 ++-
 vite.config.ts            |  1 +
 4 files changed, 31 insertions(+), 3 deletions(-)
 create mode 100644 src/utils/messageUtils.ts

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,
     },
   },
 });