Merge pull request #966 from WolframResearch/bugfix/454827

Bugfix: Use default `EvaluateWithProgress` for normal chat notebook evaluations

Author: rhennigan

PacletInfo.wl

@@ -1,7 +1,7 @@
 PacletObject[ <|
     "Name"           -> "Wolfram/Chatbook",
     "PublisherID"    -> "Wolfram",
-    "Version"        -> "1.5.2.16",
+    "Version"        -> "1.5.2.17",
     "WolframVersion" -> "14.1+",
     "Description"    -> "Wolfram Notebooks + LLMs",
     "License"        -> "MIT",

Source/Chatbook/PromptGenerators/VectorDatabases.wl

@@ -128,21 +128,27 @@ downloadVectorDatabases[ ] :=
     downloadVectorDatabases[ $localVectorDBDirectory, $vectorDBDownloadURLs ];
 
 downloadVectorDatabases[ dir0_, urls_Association ] := Enclose[
-    Module[ { names, sizes, dir, tasks },
+    Module[ { dir, lock, names, sizes, tasks },
 
-        dir = ConfirmBy[ GeneralUtilities`EnsureDirectory @ dir0, DirectoryQ, "Directory" ];
+        dir   = ConfirmBy[ GeneralUtilities`EnsureDirectory @ dir0, DirectoryQ, "Directory" ];
+        lock  = FileNameJoin @ { dir, "download.lock" };
         names = ConfirmMatch[ Keys @ urls, { __String }, "Names" ];
         sizes = ConfirmMatch[ getDownloadSize /@ Values @ urls, { __? Positive }, "Sizes" ];
 
         $downloadProgress = AssociationMap[ 0 &, names ];
         $progressText = "Downloading semantic search indices\[Ellipsis]";
 
         evaluateWithProgress[
-
-            tasks = ConfirmMatch[ KeyValueMap[ downloadVectorDatabase @ dir, urls ], { __TaskObject }, "Download" ];
-            ConfirmMatch[ taskWait @ tasks, { __TaskObject }, "TaskWait" ];
-            $progressText = "Unpacking files\[Ellipsis]";
-            ConfirmBy[ unpackVectorDatabases @ dir, DirectoryQ, "Unpacked" ],
+            WithLock[
+                File @ lock
+                ,
+                tasks = ConfirmMatch[ KeyValueMap[ downloadVectorDatabase @ dir, urls ], { __TaskObject }, "Download" ];
+                ConfirmMatch[ taskWait @ tasks, { (_TaskObject|$Failed).. }, "TaskWait" ];
+                $progressText = "Unpacking files\[Ellipsis]";
+                ConfirmBy[ unpackVectorDatabases @ dir, DirectoryQ, "Unpacked" ]
+                ,
+                PersistenceTime -> 180
+            ],
 
             <|
                 "Text"             :> $progressText,
@@ -205,7 +211,7 @@ unpackVectorDatabase // endDefinition;
 (* ::Subsubsection::Closed:: *)
 (*taskWait*)
 taskWait // beginDefinition;
-taskWait[ tasks_List ] := taskWait /@ tasks;
+taskWait[ tasks_List ] := CheckAbort[ taskWait /@ tasks, Quiet[ TaskRemove /@ tasks ], PropagateAborts -> True ];
 taskWait[ task_TaskObject ] := taskWait[ task, task[ "TaskStatus" ] ];
 taskWait[ task_TaskObject, "Removed" ] := task;
 taskWait[ task_TaskObject, _ ] := TaskWait @ task;
@@ -220,17 +226,26 @@ downloadVectorDatabase[ dir_ ] :=
     downloadVectorDatabase[ dir, ## ] &;
 
 downloadVectorDatabase[ dir_, name_String, url_String ] := Enclose[
-    Module[ { file },
+    Module[ { file, tmp },
+
         file = ConfirmBy[ FileNameJoin @ { dir, name<>".zip" }, StringQ, "File" ];
-        ConfirmMatch[
-            URLDownloadSubmit[
-                url,
-                file,
-                HandlerFunctions     -> <| "TaskProgress" -> setDownloadProgress @ name |>,
-                HandlerFunctionsKeys -> { "ByteCountDownloaded" }
-            ],
-            _TaskObject,
-            "Task"
+        tmp = file <> ".tmp";
+        Quiet[ DeleteFile /@ { file, tmp } ];
+
+        With[ { tmp = tmp, file = file },
+            ConfirmMatch[
+                URLDownloadSubmit[
+                    url,
+                    tmp,
+                    HandlerFunctions -> <|
+                        "TaskProgress" -> setDownloadProgress @ name,
+                        "TaskFinished" -> (RenameFile[ tmp, file ] &)
+                    |>,
+                    HandlerFunctionsKeys -> { "ByteCountDownloaded" }
+                ],
+                _TaskObject,
+                "Task"
+            ]
         ]
     ] // LogChatTiming[ { "DownloadVectorDatabase", name } ],
     throwInternalFailure

Source/Chatbook/Utils.wl

@@ -822,7 +822,8 @@ evaluateWithProgress // beginDefinition;
 evaluateWithProgress // Attributes = { HoldFirst };
 evaluateWithProgress[ args___ ] /; $WorkspaceChat := evaluateWithWorkspaceProgress @ args;
 evaluateWithProgress[ args___ ] /; $InlineChat := evaluateWithInlineProgress @ args;
-evaluateWithProgress[ args___ ] := evaluateWithProgressContainer[ $progressContainer, args ];
+evaluateWithProgress[ args___ ] /; $ContentSuggestions := evaluateWithProgressContainer[ $progressContainer, args ];
+evaluateWithProgress[ args___ ] := Progress`EvaluateWithProgress[ args, "Delay" -> 0 ];
 evaluateWithProgress // endDefinition;
 
 (* ::**************************************************************************************************************:: *)