Merge pull request #799 from WolframResearch/feature/search

Search Functionality
WolframResearch · Oct 2, 2024 · 6e93fca · 6e93fca
2 parents 160622d + df6c7aa
commit 6e93fca
Show file tree

Hide file tree

Showing 9 changed files with 608 additions and 84 deletions.
diff --git a/Source/Chatbook/ChatModes/ShowCodeAssistance.wl b/Source/Chatbook/ChatModes/ShowCodeAssistance.wl
@@ -14,6 +14,8 @@ $workspaceChatWidth = 325;
 
 $codeAssistanceBaseSettings = <|
     "AppName"           -> "CodeAssistance",
+    "Authentication"    -> Automatic, (* TODO *)
+    "Model"             -> <| "Service" -> "OpenAI", "Name" -> "gpt-4o" |>,
     "PromptGenerators"  -> { "RelatedDocumentation" },
     "ServiceCaller"     -> "CodeAssistance",
     "ToolOptions"       -> <| "WolframLanguageEvaluator" -> <| "AppendURIPrompt" -> True, "Method" -> "Session" |> |>,

diff --git a/Source/Chatbook/ChatModes/UI.wl b/Source/Chatbook/ChatModes/UI.wl
@@ -30,7 +30,7 @@ $inputFieldFrameOptions = Sequence[
     FrameStyle   -> Directive[ AbsoluteThickness[ 2 ], RGBColor[ "#a3c9f2" ] ]
 ];
 
-$userImageParams = <| "size" -> 50, "default" -> "identicon", "rating" -> "G" |>;
+$userImageParams = <| "size" -> 40, "default" -> "404", "rating" -> "G" |>;
 
 $defaultUserImage = Graphics[
     {
@@ -41,7 +41,7 @@ $defaultUserImage = Graphics[
         Disk[ { 0, 1 }, 1 ],
         Disk[ { 0, -1.8 }, { 1.65, 2 } ]
     },
-    ImageSize -> 25,
+    ImageSize -> 20,
     PlotRange -> { { -2.4, 2.4 }, { -2.0, 2.8 } }
 ];
 
@@ -51,7 +51,7 @@ $inputFieldBox = None;
 $inlineChatScrollPosition = 0.0;
 $lastScrollPosition       = 0.0;
 
-$maxHistoryItems = 50;
+$maxHistoryItems = 25;
 
 (* ::**************************************************************************************************************:: *)
 (* ::Section::Closed:: *)
@@ -717,12 +717,12 @@ userImage[ ] := userImage[ $CloudUserID ];
 
 userImage[ user_String ] := Enclose[
     Module[ { hash, url, image },
-        hash = Hash[ ToLowerCase @ StringTrim @ user, "MD5", "HexString" ];
-        url = ConfirmBy[ URLBuild[ { "https://www.gravatar.com/avatar/", hash }, $userImageParams ], StringQ, "URL" ];
-        image = ConfirmBy[ Import @ url, ImageQ, "Image" ];
-        userImage[ user ] = Show[ image, ImageSize -> 25 ]
+        hash  = Hash[ ToLowerCase @ StringTrim @ user, "MD5", "HexString" ];
+        url   = ConfirmBy[ URLBuild[ { "https://www.gravatar.com/avatar/", hash }, $userImageParams ], StringQ, "URL" ];
+        image = ConfirmBy[ Quiet @ Import @ url, ImageQ, "Image" ];
+        userImage[ user ] = Show[ image, ImageSize -> 20 ]
     ],
-    $defaultUserImage &
+    (userImage[ user ] = $defaultUserImage) &
 ];
 
 userImage[ other_ ] :=
@@ -858,7 +858,11 @@ createHistoryMenu[ nbo_NotebookObject ] := Enclose[
         appName = ConfirmBy[ CurrentChatSettings[ nbo, "AppName" ], StringQ, "AppName" ];
         chats = ConfirmMatch[ ListSavedChats @ appName, { ___Association }, "Chats" ];
         If[ chats === { }, Throw @ ActionMenu[ "History", { "Nothing here yet" :> Null } ] ];
-        ActionMenu[ "History", makeHistoryMenuItem[ nbo ] /@ Take[ chats, UpTo @ $maxHistoryItems ] ]
+        ActionMenu[
+            "History",
+            makeHistoryMenuItem[ nbo ] /@ Take[ chats, UpTo @ $maxHistoryItems ],
+            Method -> "Queued"
+        ]
     ],
     throwInternalFailure
 ];

diff --git a/Source/Chatbook/Common.wl b/Source/Chatbook/Common.wl
@@ -93,6 +93,7 @@ Needs[ "Wolfram`Chatbook`" ];
 $cloudNotebooks := TrueQ @ CloudSystem`$CloudNotebooks;
 
 $chatIndicatorSymbol  = "\|01f4ac";
+$defaultAppName       = "Default";
 
 $chatDelimiterStyles  = { "ChatBlockDivider", "ChatDelimiter", "ExcludedChatDelimiter" };
 $chatIgnoredStyles    = { "ChatExcluded" };
@@ -127,7 +128,9 @@ $resourceVersions = <|
     "ClickToCopy"             -> "1.0.0",
     "GPTTokenizer"            -> "1.1.0",
     "MessageFailure"          -> "1.0.0",
-    "ReplaceContext"          -> "1.0.0"
+    "RelativeTimeString"      -> "1.0.0",
+    "ReplaceContext"          -> "1.0.0",
+    "SelectByCurrentValue"    -> "1.0.1"
 |>;
 
 (* ::**************************************************************************************************************:: *)
@@ -541,7 +544,7 @@ importResourceFunction::failure = "[ERROR] Failed to import resource function `1
 importResourceFunction // Attributes = { HoldFirst };
 
 importResourceFunction[ symbol_Symbol, name_String ] :=
-    importResourceFunction[ symbol, name, Lookup[ $resourceVersions, name, "Latest" ] ];
+    importResourceFunction[ symbol, name, Lookup[ $resourceVersions, name ] ];
 
 importResourceFunction[ symbol_Symbol, name_String, version_String ] /; $mxFlag := Enclose[
     Block[ { PrintTemporary },

diff --git a/Source/Chatbook/CommonSymbols.wl b/Source/Chatbook/CommonSymbols.wl
@@ -31,6 +31,7 @@ BeginPackage[ "Wolfram`Chatbook`Common`" ];
 `$currentChatSettings;
 `$currentSettingsCache;
 `$customToolFormatter;
+`$defaultAppName;
 `$defaultChatSettings;
 `$defaultChatTools;
 `$defaultMaxCellStringLength;
@@ -151,7 +152,11 @@ BeginPackage[ "Wolfram`Chatbook`Common`" ];
 `functionTemplateBoxes;
 `getAvailableServiceNames;
 `getBoxObjectFromBoxID;
+`getChatConversationData;
 `getChatGroupSettings;
+`getChatMetadata;
+`getEmbedding;
+`getEmbeddings;
 `getHandlerFunctions;
 `getInlineChatPrompt;
 `getModelList;

diff --git a/Source/Chatbook/Main.wl b/Source/Chatbook/Main.wl
@@ -32,6 +32,7 @@ BeginPackage[ "Wolfram`Chatbook`" ];
 `$ToolFunctions;
 `$WorkspaceChat;
 `AbsoluteCurrentChatSettings;
+`AddChatToSearchIndex;
 `AppendURIInstructions;
 `BasePrompt;
 `CachedBoxes;
@@ -65,10 +66,12 @@ BeginPackage[ "Wolfram`Chatbook`" ];
 `LoadChat;
 `LogChatTiming;
 `MakeExpressionURI;
+`RebuildChatSearchIndex;
 `RelatedDocumentation;
 `RelatedWolframAlphaQueries;
 `SandboxLinguisticAssistantData;
 `SaveChat;
+`SearchChats;
 `SetModel;
 `SetToolOptions;
 `ShowCodeAssistance;
@@ -143,6 +146,7 @@ $ChatbookContexts = {
     "Wolfram`Chatbook`Prompting`",
     "Wolfram`Chatbook`ResourceInstaller`",
     "Wolfram`Chatbook`Sandbox`",
+    "Wolfram`Chatbook`Search`",
     "Wolfram`Chatbook`SendChat`",
     "Wolfram`Chatbook`Serialization`",
     "Wolfram`Chatbook`Services`",
@@ -184,6 +188,7 @@ $ChatbookProtectedNames = "Wolfram`Chatbook`" <> # & /@ {
     "$ToolFunctions",
     "$WorkspaceChat",
     "AbsoluteCurrentChatSettings",
+    "AddChatToSearchIndex",
     "AppendURIInstructions",
     "BasePrompt",
     "CachedBoxes",
@@ -216,10 +221,12 @@ $ChatbookProtectedNames = "Wolfram`Chatbook`" <> # & /@ {
     "LoadChat",
     "LogChatTiming",
     "MakeExpressionURI",
+    "RebuildChatSearchIndex",
     "RelatedDocumentation",
     "RelatedWolframAlphaQueries",
     "SandboxLinguisticAssistantData",
     "SaveChat",
+    "SearchChats",
     "SetModel",
     "SetToolOptions",
     "ShowCodeAssistance",

diff --git a/Source/Chatbook/PromptGenerators/VectorDatabases.wl b/Source/Chatbook/PromptGenerators/VectorDatabases.wl
@@ -15,9 +15,10 @@ HoldComplete[
 (* ::**************************************************************************************************************:: *)
 (* ::Section::Closed:: *)
 (*Configuration*)
-$vectorDBNames = { "DocumentationURIs", "WolframAlphaQueries" };
-$dbVersion     = "1.1.0";
-$allowDownload = True;
+$vectorDBNames   = { "DocumentationURIs", "WolframAlphaQueries" };
+$dbVersion       = "1.1.0";
+$allowDownload   = True;
+$cacheEmbeddings = True;
 
 $embeddingDimension      = 384;
 $maxNeighbors            = 50;
@@ -693,14 +694,15 @@ cacheVectorDBResult // endDefinition;
 (* ::Subsection::Closed:: *)
 (*getEmbedding*)
 getEmbedding // beginDefinition;
+getEmbedding // Options = { "CacheEmbeddings" -> $cacheEmbeddings };
 
-getEmbedding[ string_String ] :=
+getEmbedding[ string_String, opts: OptionsPattern[ ] ] :=
     With[ { embedding = $embeddingCache[ string ] },
         embedding /; NumericArrayQ @ embedding
     ];
 
-getEmbedding[ string_String ] := Enclose[
-    First @ ConfirmMatch[ getEmbeddings @ { string }, { _NumericArray }, "Embedding" ],
+getEmbedding[ string_String, opts: OptionsPattern[ ] ] := Enclose[
+    First @ ConfirmMatch[ getEmbeddings[ { string }, opts ], { _NumericArray }, "Embedding" ],
     throwInternalFailure
 ];
 
@@ -710,19 +712,31 @@ getEmbedding // endDefinition;
 (* ::Subsection::Closed:: *)
 (*getEmbeddings*)
 getEmbeddings // beginDefinition;
+getEmbeddings // Options = { "CacheEmbeddings" -> $cacheEmbeddings };
 
-getEmbeddings[ { } ] := { };
+getEmbeddings[ { }, opts: OptionsPattern[ ] ] := { };
 
-getEmbeddings[ strings: { __String } ] := Enclose[
+getEmbeddings[ strings: { __String }, opts: OptionsPattern[ ] ] :=
+    If[ TrueQ @ OptionValue[ "CacheEmbeddings" ],
+        getEmbeddings0 @ strings,
+        Block[ { $cacheEmbeddings = False }, getAndCacheEmbeddings @ strings ]
+    ] // LogChatTiming[ "GetEmbeddings" ];
+
+getEmbeddings // endDefinition;
+
+
+getEmbeddings0 // beginDefinition;
+
+getEmbeddings0[ strings: { __String } ] := Enclose[
     Module[ { notCached },
         notCached = Select[ strings, ! KeyExistsQ[ $embeddingCache, # ] & ];
         ConfirmMatch[ getAndCacheEmbeddings @ notCached, { ___NumericArray }, "CacheEmbeddings" ];
         ConfirmMatch[ Lookup[ $embeddingCache, strings ], { __NumericArray }, "Result" ]
-    ] // LogChatTiming[ "GetEmbeddings" ],
+    ],
     throwInternalFailure
 ];
 
-getEmbeddings // endDefinition;
+getEmbeddings0 // endDefinition;
 
 (* ::**************************************************************************************************************:: *)
 (* ::Subsubsection::Closed:: *)
@@ -742,10 +756,7 @@ getAndCacheEmbeddings[ strings: { __String } ] /; $embeddingModel === "SentenceB
 
         ConfirmAssert[ Length @ strings === Length @ vectors, "LengthCheck" ];
 
-        MapThread[
-            ($embeddingCache[ #1 ] = toTinyVector @ #2) &,
-            { strings, vectors }
-        ]
+        MapThread[ cacheEmbedding, { strings, vectors } ]
     ],
     throwInternalFailure
 ];
@@ -771,16 +782,21 @@ getAndCacheEmbeddings[ strings: { __String } ] := Enclose[
 
         ConfirmAssert[ Length @ strings === Length @ vectors, "LengthCheck" ];
 
-        MapThread[
-            ($embeddingCache[ #1 ] = toTinyVector @ #2) &,
-            { strings, vectors }
-        ]
+        MapThread[ cacheEmbedding, { strings, vectors } ]
     ],
     throwInternalFailure
 ];
 
 getAndCacheEmbeddings // endDefinition;
 
+(* ::**************************************************************************************************************:: *)
+(* ::Subsubsection::Closed:: *)
+(*cacheEmbedding*)
+cacheEmbedding // beginDefinition;
+cacheEmbedding[ key_String, vector_ ] /; ! $cacheEmbeddings := toTinyVector @ vector;
+cacheEmbedding[ key_String, vector_ ] := $embeddingCache[ key ] = toTinyVector @ vector;
+cacheEmbedding // endDefinition;
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsubsection::Closed:: *)
 (*sentenceBERTEmbedding*)