From 591f61d35459ec2817ff251aa2392b4f7f7cdeaf Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Mon, 7 Oct 2024 01:32:46 +0300
Subject: [PATCH] Don't split tool calls and tool results

---
 public/scripts/openai.js | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 3d8355216..90334f299 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -732,19 +732,15 @@ async function populateChatHistory(messages, prompts, chatCompletion, type = nul
             /** @type {import('./tool-calling.js').ToolInvocation[]} */
             const invocations = chatPrompt.invocations;
             const toolCallMessage = new Message(chatMessage.role, undefined, 'toolCall-' + chatMessage.identifier);
+            const toolResultMessages = invocations.slice().reverse().map((invocation) => new Message('tool', invocation.result || '[No content]', invocation.id));
             toolCallMessage.setToolCalls(invocations);
-            if (chatCompletion.canAfford(toolCallMessage)) {
-                chatCompletion.reserveBudget(toolCallMessage);
-                for (const invocation of invocations.slice().reverse()) {
-                    const toolResultMessage = new Message('tool', invocation.result || '[No content]', invocation.id);
-                    const canAfford = chatCompletion.canAfford(toolResultMessage);
-                    if (!canAfford) {
-                        break;
-                    }
-                    chatCompletion.insertAtStart(toolResultMessage, 'chatHistory');
+            if (chatCompletion.canAffordAll([toolCallMessage, ...toolResultMessages])) {
+                for (const resultMessage of toolResultMessages) {
+                    chatCompletion.insertAtStart(resultMessage, 'chatHistory');
                 }
-                chatCompletion.freeBudget(toolCallMessage);
                 chatCompletion.insertAtStart(toolCallMessage, 'chatHistory');
+            } else {
+                break;
             }
 
             continue;
@@ -2652,6 +2648,15 @@ export class ChatCompletion {
         return 0 <= this.tokenBudget - message.getTokens();
     }
 
+    /**
+     * Checks if the token budget can afford the tokens of all the specified messages.
+     * @param {Message[]} messages - The messages to check for affordability.
+     * @returns {boolean} True if the budget can afford all the messages, false otherwise.
+     */
+    canAffordAll(messages) {
+        return 0 <= this.tokenBudget - messages.reduce((total, message) => total + message.getTokens(), 0);
+    }
+
     /**
      * Checks if a message with the specified identifier exists in the collection.
      *