Merge pull request #939 from WolframResearch/bugfix/more-token-optimi…

…zations More token usage optimizations
WolframResearch · Nov 20, 2024 · aff6ee2 · aff6ee2
2 parents 8ae59f2 + e89b539
commit aff6ee2
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 40 deletions.
diff --git a/Source/Chatbook/ChatModes/ContentSuggestions.wl b/Source/Chatbook/ChatModes/ContentSuggestions.wl
@@ -35,6 +35,11 @@ $$generatedStyle = "Echo"|"EchoAfter"|"EchoBefore"|"EchoTiming"|"Message"|"Outpu
 $$whitespaceString = _String? (StringMatchQ[ WhitespaceCharacter... ])
 $$whitespace       = $$whitespaceString | TextData @ $$whitespaceString | TextData @ { $$whitespaceString... };
 
+
+(* TODO: stop tokens for resource notebooks:
+    CurrentValue[ EvaluationNotebook[ ], { TaggingRules, "ResourceCreateNotebook" } ]
+*)
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsection::Closed:: *)
 (*Wolfram Language Suggestions*)

diff --git a/Source/Chatbook/ChatModes/Context.wl b/Source/Chatbook/ChatModes/Context.wl
@@ -11,9 +11,11 @@ Needs[ "Wolfram`Chatbook`ChatModes`Common`" ];
 (* ::Section::Closed:: *)
 (*Configuration*)
 $maxCellsBeforeSelection    = 100;
-$maxCellsAfterSelection     = 20;
+$maxCellsAfterSelection     = 50;
 $notebookInstructionsPrompt = True;
 $currentSelectionIndicator  = { $leftSelectionIndicator, $rightSelectionIndicator };
+$notebookContextLimitScale  = 0.25;
+$downScaledSettings         = { "MaxCellStringLength", "MaxContextTokens", "MaxOutputCellStringLength" };
 
 $notebookContextTemplate = StringTemplate[ "\
 IMPORTANT: Below is some context from the user's currently selected notebook. \
@@ -221,14 +223,15 @@ getContextFromSelection[ chatNB_NotebookObject, settings_Association, opts: Opti
 getContextFromSelection[ chatNB_NotebookObject, None, settings_Association, opts: OptionsPattern[ ] ] :=
     None;
 
-getContextFromSelection[ chatNB_, nbo_NotebookObject, settings_Association, opts: OptionsPattern[ ] ] := Enclose[
+getContextFromSelection[ chatNB_, nbo_NotebookObject, settings0_Association, opts: OptionsPattern[ ] ] := Enclose[
     Catch @ Block[
         {
             $notebookInstructionsPrompt = OptionValue[ "NotebookInstructionsPrompt" ],
             $maxCellsBeforeSelection    = OptionValue[ "MaxCellsBeforeSelection"    ],
             $maxCellsAfterSelection     = OptionValue[ "MaxCellsAfterSelection"     ]
         },
-        Module[ { selectionData, string },
+        Module[ { settings, selectionData, string },
+            settings = ConfirmBy[ downScaledSettings @ settings0, AssociationQ, "Settings" ];
 
             selectionData = ConfirmMatch[
                 LogChatTiming @ selectContextCells @ nbo,
@@ -291,6 +294,18 @@ getContextFromSelection0[ selectionData_Association, settings_ ] := Enclose[
 
 getContextFromSelection0 // endDefinition;
 
+(* ::**************************************************************************************************************:: *)
+(* ::Subsubsection::Closed:: *)
+(*downScaledSettings*)
+downScaledSettings // beginDefinition;
+
+downScaledSettings[ settings_Association ] := <|
+    settings,
+    Floor[ $notebookContextLimitScale * KeyTake[ settings, $downScaledSettings ] ]
+|>;
+
+downScaledSettings // endDefinition;
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsubsection::Closed:: *)
 (*extractSelectionPrompt*)
@@ -401,7 +416,16 @@ insertSelectionIndicator[ cells_ ] :=
     insertSelectionIndicator[ cells, $currentSelectionIndicator ];
 
 insertSelectionIndicator[
-    { { beforeCells___Cell }, { selectedCells___Cell }, { afterCells___Cell } },
+    { { beforeCells___Cell }, { }, { afterCells___Cell } },
+    { before_String, after_String }
+] := {
+    beforeCells,
+    Cell @ Verbatim[ before<>after ],
+    afterCells
+};
+
+insertSelectionIndicator[
+    { { beforeCells___Cell }, { selectedCells__Cell }, { afterCells___Cell } },
     { before_String, after_String }
 ] := {
     beforeCells,

diff --git a/Source/Chatbook/ChatModes/ShowNotebookAssistance.wl b/Source/Chatbook/ChatModes/ShowNotebookAssistance.wl
@@ -13,14 +13,18 @@ Needs[ "Wolfram`Chatbook`ChatModes`Common`" ];
 $workspaceChatWidth := $workspaceChatWidth = Switch[ $OperatingSystem, "MacOSX", 450, _, 360 ];
 
 $notebookAssistanceBaseSettings = <|
-    "AppName"           -> "NotebookAssistance",
-    "LLMEvaluator"      -> "NotebookAssistant",
-    "Model"             -> <| "Service" -> "LLMKit", "Name" -> Automatic |>,
-    "PromptGenerators"  -> { "RelatedDocumentation" },
-    "ServiceCaller"     -> "NotebookAssistance",
-    "ToolOptions"       -> <| "WolframLanguageEvaluator" -> <| "AppendURIPrompt" -> True, "Method" -> "Session" |> |>,
-    "Tools"             -> { "NotebookEditor" },
-    "ToolSelectionType" -> <| "DocumentationLookup" -> None, "DocumentationSearcher" -> None |>
+    "AppName"                   -> "NotebookAssistance",
+    "LLMEvaluator"              -> "NotebookAssistant",
+    "MaxContextTokens"          -> 2^15,
+    "MaxToolResponses"          -> 3,
+    "Model"                     -> <| "Service" -> "LLMKit", "Name" -> Automatic |>,
+    "PromptGenerators"          -> { "RelatedDocumentation" },
+    "ServiceCaller"             -> "NotebookAssistance",
+    "Tools"                     -> { "NotebookEditor" },
+    "ToolSelectionType"         -> <| "DocumentationLookup" -> None, "DocumentationSearcher" -> None |>,
+    "ToolOptions"               -> <|
+        "WolframLanguageEvaluator" -> <| "AppendURIPrompt" -> True, "Method" -> "Session" |>
+    |>
 |>;
 
 $notebookAssistanceWorkspaceSettings := <|

diff --git a/Source/Chatbook/ChatTitle.wl b/Source/Chatbook/ChatTitle.wl
@@ -26,6 +26,11 @@ Here is the chat transcript:
 
 Remember, respond only with the title and nothing else.";
 
+(* TODO:
+    Detect when topic has changed and display a message:
+        Is this a new topic? You'll get more accurate results with a new chat when you change topics. <Start New Chat>
+*)
+
 (* ::**************************************************************************************************************:: *)
 (* ::Section::Closed:: *)
 (*GenerateChatTitle*)

diff --git a/Source/Chatbook/Prompting.wl b/Source/Chatbook/Prompting.wl
@@ -114,20 +114,21 @@ $basePromptComponents[ "GeneralInstructionsHeader" ] = "\
 # General Instructions
 ";
 
-$basePromptComponents[ "NotebooksPreamble" ] := If[ TrueQ @ $WorkspaceChat,
-    "You are interacting with a user through a special Wolfram Chat interface alongside normal notebooks. \
-    You will often receive context from the user's notebooks, but you will see it formatted as markdown. \
-    Similarly, your responses are automatically converted from plain text before being displayed to the user. \
-    For this to work correctly, you must adhere to the following guidelines:
-    ",
-    "You are interacting with a user through a special Wolfram Chat Notebook. \
-    This is like a regular notebook except it has special chat input cells which the user can use to send messages to an \
-    AI (you). \
-    The messages you receive from the user have been converted to plain text from notebook content. \
-    Similarly, your messages are automatically converted from plain text before being displayed to the user. \
-    For this to work correctly, you must adhere to the following guidelines:
-    "
-];
+$basePromptComponents[ "NotebooksPreamble" ] :=
+If[ TrueQ @ $WorkspaceChat,
+"You are interacting with a user through a special Wolfram Chat interface alongside normal notebooks. \
+You will often receive context from the user's notebooks, but you will see it formatted as markdown. \
+Similarly, your responses are automatically converted from plain text before being displayed to the user. \
+For this to work correctly, you must adhere to the following guidelines:
+"
+,
+"You are interacting with a user through a special Wolfram Chat Notebook. \
+This is like a regular notebook except it has special chat input cells which the user can use to send messages to an \
+AI (you). \
+The messages you receive from the user have been converted to plain text from notebook content. \
+Similarly, your messages are automatically converted from plain text before being displayed to the user. \
+For this to work correctly, you must adhere to the following guidelines:
+" ];
 
 $basePromptComponents[ "AutoAssistant" ] = "\
 * ALWAYS begin your response with one of the following tags to indicate the type of response: [INFO], [WARNING], or [ERROR]

diff --git a/Source/Chatbook/SendChat.wl b/Source/Chatbook/SendChat.wl
@@ -895,6 +895,7 @@ autoCorrect[ string_String ] := StringReplace[ string, $llmAutoCorrectRules ];
 autoCorrect // endDefinition;
 
 $llmAutoCorrectRules := $llmAutoCorrectRules = Flatten @ {
+    StartOfLine ~~ WhitespaceCharacter... ~~ "/command\ncode:" :> "/wl\ncode:",
     "```" ~~ code: Except[ "\n" ].. ~~ "```" :> "``"<>code<>"``",
     "wolfram_language_evaliator" -> "wolfram_language_evaluator",
     "\\!\\(\\*MarkdownImageBox[\"" ~~ Shortest[ uri__ ] ~~ "\"]\\)" :> uri,
@@ -909,6 +910,35 @@ $llmAutoCorrectRules := $llmAutoCorrectRules = Flatten @ {
     $longNameCharacters
 };
 
+(* TODO:
+Automatically rewrite these as WL tool calls:
+
+
+    Sure! I will use `EmbeddedService` to show a map of Tokyo by utilizing the OpenStreetMap service.
+
+    ```wl
+    EmbeddedService[{\"OpenStreetMap\", GeoPosition[\[FreeformPrompt][\"Tokyo\"]]}]
+    ``` /exec
+
+=====
+
+    To show a map of the United States with all its state capitals, we can use the [GeoGraphics](paclet:ref/GeoGraphics) function along with [Entity](paclet:ref/Entity) to get the positions of the capitals. Here's how you can do it:
+
+    ```wl
+    GeoGraphics[
+        {Red, PointSize[Large],
+        Point[GeoPosition /@ EntityValue[
+            EntityClass["AdministrativeDivision", "USStates"],
+            "CapitalLocation"
+        ]]
+        },
+        GeoRange -> Entity["Country", "UnitedStates"]
+    ]
+    ```
+
+    Let me create this map for you. /wl
+*)
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsubsection::Closed:: *)
 (*splitDynamicContent*)
@@ -1240,17 +1270,15 @@ toolEvaluation[ settings_, container_Symbol, cell_, as_Association ] := Enclose[
                 ToString @ output
             ];
 
-        newMessages = Join[
+        newMessages = Flatten @ {
             messages,
-            {
-                <|
-                    "Role"        -> "Assistant",
-                    "Content"     -> appendToolCallEndToken[ settings, StringTrim @ string ],
-                    "ToolRequest" -> True
-                |>,
-                makeToolResponseMessage[ settings, response ]
-            }
-        ];
+            <|
+                "Role"        -> "Assistant",
+                "Content"     -> appendToolCallEndToken[ settings, StringTrim @ string ],
+                "ToolRequest" -> True
+            |>,
+            makeToolResponseMessage[ settings, response ]
+        };
 
         $finishReason = None;
 
@@ -1338,8 +1366,10 @@ makeToolResponseMessage // beginDefinition;
 makeToolResponseMessage[ settings_, response_ ] :=
     makeToolResponseMessage[ settings, settings[ "Model" ], response ];
 
-makeToolResponseMessage[ settings_, model_Association, response_ ] :=
-    makeToolResponseMessage0[ model[ "Service" ], model[ "Family" ], response ];
+makeToolResponseMessage[ settings_, model_Association, response_ ] := {
+    If[ TrueQ @ settings[ "ToolCallRetryMessage" ], $toolCallRetryMessage, Nothing ],
+    makeToolResponseMessage0[ model[ "Service" ], model[ "Family" ], response ]
+};
 
 makeToolResponseMessage // endDefinition;
 
@@ -1375,6 +1405,17 @@ makeToolResponseMessage0[ service_String, family_, response_ ] :=
 
 makeToolResponseMessage0 // endDefinition;
 
+
+$toolCallRetryPrompt = "\
+IMPORTANT: If a tool call does not give the expected output, \
+ask the user before retrying unless you are ABSOLUTELY SURE you know how to fix the issue.";
+
+$toolCallRetryMessage = <|
+    "Role"      -> "System",
+    "Content"   -> $toolCallRetryPrompt,
+    "Temporary" -> True
+|>;
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsubsection::Closed:: *)
 (*wrapResponse*)

diff --git a/Source/Chatbook/Serialization.wl b/Source/Chatbook/Serialization.wl
@@ -1079,8 +1079,7 @@ toMarkdownImageBox // beginDefinition;
 
 toMarkdownImageBox[ graphics_ ] := Enclose[
     Catch @ Module[ { img, uri },
-        img = ConfirmMatch[ rasterizeGraphics @ graphics, _Image | Missing[ "OutOfMemory" ], "RasterizeGraphics" ];
-        If[ MissingQ @ img, Block[ { $multiModalImages = False }, Throw @ fasterCellToString0 @ graphics ] ];
+        img = Flatten @ RawBoxes @ graphics;
         uri = ConfirmBy[ MakeExpressionURI[ "image", img ], StringQ, "RasterID" ];
         needsBasePrompt[ "MarkdownImageBox" ];
         If[ toolSelectedQ[ "WolframLanguageEvaluator" ], needsBasePrompt[ "MarkdownImageBoxImporting" ] ];

diff --git a/Source/Chatbook/Settings.wl b/Source/Chatbook/Settings.wl
@@ -73,6 +73,7 @@ $defaultChatSettings = <|
     "Tokenizer"                      -> Automatic,
     "ToolCallExamplePromptStyle"     -> Automatic,
     "ToolCallFrequency"              -> Automatic,
+    "ToolCallRetryMessage"           -> Automatic,
     "ToolExamplePrompt"              -> Automatic,
     "ToolMethod"                     -> Automatic,
     "ToolOptions"                    :> $DefaultToolOptions,
@@ -331,6 +332,7 @@ $gpt4oTextToolOverrides = <| "Model" -> <| "Service" -> "OpenAI", "Name" -> "gpt
 llmKitQ // beginDefinition;
 
 llmKitQ[ as_Association ] := TrueQ @ Or[
+    as[ "Authentication"          ] === "LLMKit",
     as[ "Model", "Service"        ] === "LLMKit",
     as[ "Model", "Authentication" ] === "LLMKit"
 ];
@@ -394,6 +396,7 @@ resolveAutoSetting0[ as_, "Tokenizer"                      ] := getTokenizer @ a
 resolveAutoSetting0[ as_, "TokenizerName"                  ] := getTokenizerName @ as;
 resolveAutoSetting0[ as_, "ToolCallExamplePromptStyle"     ] := chooseToolExamplePromptStyle @ as;
 resolveAutoSetting0[ as_, "ToolCallFrequency"              ] := Automatic;
+resolveAutoSetting0[ as_, "ToolCallRetryMessage"           ] := toolCallRetryMessageQ @ as;
 resolveAutoSetting0[ as_, "ToolExamplePrompt"              ] := chooseToolExamplePromptSpec @ as;
 resolveAutoSetting0[ as_, "ToolsEnabled"                   ] := toolsEnabledQ @ as;
 resolveAutoSetting0[ as_, "TrackScrollingWhenPlaced"       ] := scrollOutputQ @ as;
@@ -416,6 +419,7 @@ $autoSettingKeyDependencies = <|
     "Tokenizer"                  -> "TokenizerName",
     "TokenizerName"              -> "Model",
     "ToolCallExamplePromptStyle" -> { "Model", "ToolsEnabled" },
+    "ToolCallRetryMessage"       -> { "Authentication", "Model" },
     "ToolExamplePrompt"          -> "Model",
     "Tools"                      -> { "LLMEvaluator", "ToolsEnabled" },
     "ToolsEnabled"               -> { "Model", "ToolCallFrequency" }
@@ -438,6 +442,13 @@ $autoSettingKeyPriority := Enclose[
     * ChatContextPreprompt
 *)
 
+(* ::**************************************************************************************************************:: *)
+(* ::Subsubsection::Closed:: *)
+(*toolCallRetryMessageQ*)
+toolCallRetryMessageQ // beginDefinition;
+toolCallRetryMessageQ[ as_Association ] := llmKitQ @ as;
+toolCallRetryMessageQ // endDefinition;
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsubsection::Closed:: *)
 (*autoAuthentication*)