From 6b9aa17623d18300e783401508da712a740ddad5 Mon Sep 17 00:00:00 2001
From: Rick Hennigan <richardh@wolfram.com>
Date: Tue, 2 Jan 2024 13:49:35 -0500
Subject: [PATCH] Initialize budget values earlier

---
 Source/Chatbook/ChatMessages.wl | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/Source/Chatbook/ChatMessages.wl b/Source/Chatbook/ChatMessages.wl
index 718ab3ad..8d0ef48a 100644
--- a/Source/Chatbook/ChatMessages.wl
+++ b/Source/Chatbook/ChatMessages.wl
@@ -217,11 +217,13 @@ makeChatMessages[ settings_, cells_ ] :=
             $multimodalMessages = TrueQ @ settings[ "Multimodal" ],
             $tokenBudget        = settings[ "MaxContextTokens" ],
             $tokenPressure      = 0.0,
-            $cellStringBudget   = Replace[
+            $initialCellStringBudget = Replace[
                 settings[ "MaxCellStringLength" ],
                 Except[ $$size ] -> $defaultMaxCellStringLength
-            ]
+            ],
+            $cellStringBudget
         },
+        $cellStringBudget = $initialCellStringBudget;
         If[ settings[ "BasePrompt" ] =!= None, tokenCheckedMessage[ settings, $fullBasePrompt ] ];
         (* FIXME: need to account for persona/tool prompting as well *)
         makeChatMessages0[ settings, cells ]
@@ -245,10 +247,6 @@ makeChatMessages0[ settings0_, cells_List ] := Enclose[
         toMessage0 = Confirm[ getCellMessageFunction @ settings, "CellMessageFunction" ];
 
         $tokenBudgetLog = Internal`Bag[ ];
-        $initialCellStringBudget = Replace[
-            settings[ "MaxCellStringLength" ],
-            Except[ $$size ] -> $defaultMaxCellStringLength
-        ];
 
         toMessage = Function @ With[
             { msg = toMessage0[ #1, <| #2, "TokenBudget" -> $tokenBudget, "TokenPressure" -> $tokenPressure |> ] },