fix todo: estimate user's token

Author: heimoshuiyu

src/app.tsx

@@ -108,13 +108,31 @@ export function App() {
   const [chatStore, _setChatStore] = useState(
     getChatStoreByIndex(selectedChatIndex)
   );
-  const setChatStore = (cs: ChatStore) => {
+  const setChatStore = (chatStore: ChatStore) => {
     console.log("saved chat", selectedChatIndex, chatStore);
     localStorage.setItem(
       `${STORAGE_NAME}-${selectedChatIndex}`,
-      JSON.stringify(cs)
+      JSON.stringify(chatStore)
     );
-    _setChatStore(cs);
+
+    console.log("recalculate postBeginIndex");
+    const max = chatStore.maxTokens - chatStore.tokenMargin;
+    let sum = 0;
+    chatStore.postBeginIndex = chatStore.history.filter(
+      ({ hide }) => !hide
+    ).length;
+    for (const msg of chatStore.history
+      .filter(({ hide }) => !hide)
+      .slice()
+      .reverse()) {
+      if (sum + msg.token > max) break;
+      sum += msg.token;
+      chatStore.postBeginIndex -= 1;
+    }
+    chatStore.postBeginIndex =
+      chatStore.postBeginIndex < 0 ? 0 : chatStore.postBeginIndex;
+
+    _setChatStore(chatStore);
   };
   useEffect(() => {
     _setChatStore(getChatStoreByIndex(selectedChatIndex));

src/chatbox.tsx

@@ -127,6 +127,20 @@ export default function ChatBOX(props: {
         (models[data.model]?.price?.completion ?? 0);
     }
     const content = client.processFetchResponse(data);
+
+    // estimate user's input message token
+    let aboveToken = 0;
+    for (const msg of chatStore.history
+      .filter(({ hide }) => !hide)
+      .slice(chatStore.postBeginIndex, -1)) {
+      aboveToken += msg.token;
+    }
+    if (data.usage.prompt_tokens) {
+      const userMessageToken = data.usage.prompt_tokens - aboveToken;
+      console.log("set user message token");
+      chatStore.history.slice(-1)[0].token = userMessageToken;
+    }
+
     chatStore.history.push({
       role: "assistant",
       content,
@@ -143,9 +157,9 @@ export default function ChatBOX(props: {
     client.sysMessageContent = chatStore.systemMessageContent;
     client.tokens_margin = chatStore.tokenMargin;
     client.messages = chatStore.history
-      .slice(chatStore.postBeginIndex)
       // only copy non hidden message
       .filter(({ hide }) => !hide)
+      .slice(chatStore.postBeginIndex)
       // only copy content and role attribute to client for posting
       .map(({ content, role }) => {
         return {
@@ -156,20 +170,6 @@ export default function ChatBOX(props: {
     client.model = chatStore.model;
     client.max_tokens = chatStore.maxTokens;
 
-    // todo move code
-    const max = chatStore.maxTokens - chatStore.tokenMargin;
-    let sum = 0;
-    chatStore.postBeginIndex = chatStore.history.filter(
-      ({ hide }) => !hide
-    ).length;
-    for (const msg of chatStore.history.slice().reverse()) {
-      sum += msg.token;
-      if (sum > max) break;
-      chatStore.postBeginIndex -= 1;
-    }
-    chatStore.postBeginIndex =
-      chatStore.postBeginIndex < 0 ? 0 : chatStore.postBeginIndex;
-
     try {
       setShowGenerating(true);
       const response = await client._fetch(chatStore.streamMode);
@@ -186,20 +186,6 @@ export default function ChatBOX(props: {
       chatStore.tokenMargin = client.tokens_margin;
       chatStore.totalTokens = client.total_tokens;
 
-      // todo move code
-      const max = chatStore.maxTokens - chatStore.tokenMargin;
-      let sum = 0;
-      chatStore.postBeginIndex = chatStore.history.filter(
-        ({ hide }) => !hide
-      ).length;
-      for (const msg of chatStore.history.slice().reverse()) {
-        sum += msg.token;
-        if (sum > max) break;
-        chatStore.postBeginIndex -= 1;
-      }
-      chatStore.postBeginIndex =
-        chatStore.postBeginIndex < 0 ? 0 : chatStore.postBeginIndex;
-
       console.log("postBeginIndex", chatStore.postBeginIndex);
       setChatStore({ ...chatStore });
     } catch (error) {

src/message.tsx

@@ -18,20 +18,6 @@ export default function Message(props: Props) {
         chatStore.history[messageIndex].hide =
           !chatStore.history[messageIndex].hide;
 
-        // todo move code
-        const max = chatStore.maxTokens - chatStore.tokenMargin;
-        let sum = 0;
-        chatStore.postBeginIndex = chatStore.history.filter(
-          ({ hide }) => !hide
-        ).length;
-        for (const msg of chatStore.history.slice().reverse()) {
-          sum += msg.token;
-          if (sum > max) break;
-          chatStore.postBeginIndex -= 1;
-        }
-        chatStore.postBeginIndex =
-          chatStore.postBeginIndex < 0 ? 0 : chatStore.postBeginIndex;
-
         //chatStore.totalTokens =
         chatStore.totalTokens = 0;
         for (const i of chatStore.history