From b0d0f2111b69a3d9561278bbcd3cbc6fdefad414 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sat, 12 Oct 2024 00:22:37 +0300
Subject: [PATCH] Tool Calling: Reserve prompt budget for tool definitions.

---
 public/scripts/openai.js | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 7fb23e711..b3938dc6d 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -998,6 +998,15 @@ async function populateChatCompletion(prompts, chatCompletion, { bias, quietProm
         chatCompletion.insert(Message.fromPrompt(prompt), 'main', prompt.position);
     }
 
+    // Pre-allocation of tokens for tool data
+    if (ToolManager.canPerformToolCalls(type)) {
+        const toolData = {};
+        await ToolManager.registerFunctionToolsOpenAI(toolData);
+        const toolMessage = [{ role: 'user', content: JSON.stringify(toolData) }];
+        const toolTokens = tokenHandler.count(toolMessage);
+        chatCompletion.reserveBudget(toolTokens);
+    }
+
     // Add in-chat injections
     messages = populationInjectionPrompts(absolutePrompts, messages);