From 6c8bd06308c69d51e2eb174541792a870a83d2d6 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Mon, 11 Sep 2023 17:22:31 +0300
Subject: [PATCH] Reserve 3 extra tokens for each chat completion

---
 public/scripts/openai.js | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 3ac66f0df..d4764f02b 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -576,6 +576,7 @@ function populateChatCompletion(prompts, chatCompletion, { bias, quietPrompt, ty
         chatCompletion.add(collection, index);
     };
 
+    chatCompletion.reserveBudget(3); // every reply is primed with <|start|>assistant<|message|>
     // Character and world information
     addToChatCompletion('worldInfoBefore');
     addToChatCompletion('main');
@@ -1761,9 +1762,12 @@ class ChatCompletion {
     /**
      * Reserves the tokens required by the given message from the token budget.
      *
-     * @param {Message|MessageCollection} message - The message whose tokens to reserve.
+     * @param {Message|MessageCollection|number} message - The message whose tokens to reserve.
      */
-    reserveBudget(message) { this.decreaseTokenBudgetBy(message.getTokens()) };
+    reserveBudget(message) {
+        const tokens = typeof message === 'number' ? message : message.getTokens();
+        this.decreaseTokenBudgetBy(tokens);
+    };
 
     /**
      * Frees up the tokens used by the given message from the token budget.