Merge pull request #1021 from WolframResearch/feature/reranking-speed…

…-improvements Optimized reranking prompt to reduce output token counts
WolframResearch · Jan 13, 2025 · 1b80649 · 1b80649
2 parents 1b6a690 + 86d33cb
commit 1b80649
Show file tree

Hide file tree

Showing 7 changed files with 276 additions and 35 deletions.
diff --git a/Source/Chatbook/Handlers.wl b/Source/Chatbook/Handlers.wl
@@ -27,6 +27,11 @@ addHandlerArguments // beginDefinition;
 addHandlerArguments[ args_ ] :=
     addHandlerArguments[ $ChatHandlerData, Association @ args ];
 
+addHandlerArguments[ current_? AssociationQ, new_? AssociationQ ] /; AnyTrue[ new, AssociationQ ] := Enclose[
+    $ChatHandlerData = ConfirmBy[ combineNestedHandlerData[ current, new ], AssociationQ, "AddHandlerArguments" ],
+    throwInternalFailure
+];
+
 addHandlerArguments[ current_? AssociationQ, new_? AssociationQ ] :=
     $ChatHandlerData = <| current, new |>;
 
@@ -37,6 +42,19 @@ addHandlerArguments[ current_, new_? AssociationQ ] := (
 
 addHandlerArguments // endDefinition;
 
+(* ::**************************************************************************************************************:: *)
+(* ::Subsubsection::Closed:: *)
+(*combineNestedHandlerData*)
+combineNestedHandlerData // beginDefinition;
+combineNestedHandlerData[ as1_Association, as2_Association ] := combineNestedHandlerData0 @ { as1, as2 };
+combineNestedHandlerData // endDefinition;
+
+combineNestedHandlerData0 // beginDefinition;
+combineNestedHandlerData0[ { as1_Association, as2_Association } ] := Merge[ { as1, as2 }, combineNestedHandlerData0 ];
+combineNestedHandlerData0[ { value_ } ] := value;
+combineNestedHandlerData0[ { _, value_ } ] := value;
+combineNestedHandlerData0 // endDefinition;
+
 (* ::**************************************************************************************************************:: *)
 (* ::Subsection::Closed:: *)
 (*applyHandlerFunction*)

diff --git a/Source/Chatbook/LLMUtilities.wl b/Source/Chatbook/LLMUtilities.wl
@@ -34,7 +34,11 @@ llmSynthesize[ prompt: $$llmPrompt ] :=
     llmSynthesize[ prompt, <| |> ];
 
 llmSynthesize[ prompt: $$llmPrompt, evaluator_Association ] := Enclose[
-    ConfirmMatch[ llmSynthesize0[ prompt, evaluator, 1 ], Except[ "", _String ], "Result" ],
+    ConfirmMatch[
+        llmSynthesize0[ prompt, evaluator, 1 ],
+        If[ MatchQ[ Flatten @ { evaluator[ "StopTokens" ] }, { __String } ], _String, Except[ "", _String ] ],
+        "Result"
+    ],
     throwInternalFailure
 ];
 
@@ -70,19 +74,25 @@ llmSynthesizeSubmit[ prompt: $$llmPrompt, callback_ ] :=
     llmSynthesizeSubmit[ prompt, <| |>, callback ];
 
 llmSynthesizeSubmit[ prompt0: $$llmPrompt, evaluator0_Association, callback_ ] := Enclose[
-    Module[ { evaluator, prompt, messages, config, chunks, handlers, keys },
+    Module[ { evaluator, prompt, messages, config, chunks, allowEmpty, handlers, keys, auth },
 
-        evaluator = ConfirmBy[
-            <| $defaultLLMSynthesizeEvaluator, DeleteCases[ evaluator0, Automatic | _Missing ] |>,
-            AssociationQ,
-            "Evaluator"
+        evaluator = Replace[
+            ConfirmBy[
+                <| $defaultLLMSynthesizeEvaluator, DeleteCases[ evaluator0, Automatic | _Missing ] |>,
+                AssociationQ,
+                "Evaluator"
+            ],
+            Verbatim[ Verbatim ][ value_ ] :> value,
+            { 1 }
         ];
 
         prompt   = ConfirmMatch[ truncatePrompt[ prompt0, evaluator ], $$llmPrompt, "Prompt" ];
         messages = { <| "Role" -> "User", "Content" -> prompt |> };
         config   = LLMConfiguration @ evaluator;
         chunks   = Internal`Bag[ ];
 
+        allowEmpty = MatchQ[ Flatten @ { evaluator[ "StopTokens" ] }, { __String } ];
+
         handlers = <|
             "BodyChunkReceived" -> Function[
                 Internal`StuffBag[ chunks, # ]
@@ -93,7 +103,7 @@ llmSynthesizeSubmit[ prompt0: $$llmPrompt, evaluator0_Association, callback_ ] :
                     $lastSynthesizeSubmitLog = data;
                     strings = extractBodyChunks @ data;
                     Which[
-                        MatchQ[ strings, { __String } ],
+                        MatchQ[ strings, { __String } ] || (allowEmpty && strings === { }),
                             With[ { s = StringJoin @ strings }, callback[ s, #1 ] ],
                         FailureQ @ strings,
                             callback[ strings, #1 ],
@@ -106,10 +116,12 @@ llmSynthesizeSubmit[ prompt0: $$llmPrompt, evaluator0_Association, callback_ ] :
 
         keys = { "BodyChunk", "BodyChunkProcessed", "StatusCode", "EventName" };
 
+        auth = Lookup[ evaluator, "Authentication", $llmSynthesizeAuthentication ];
+
         setServiceCaller @ LLMServices`ChatSubmit[
             messages,
             config,
-            Authentication       -> $llmSynthesizeAuthentication,
+            Authentication       -> auth,
             HandlerFunctions     -> handlers,
             HandlerFunctionsKeys -> keys,
             "TestConnection"     -> False

diff --git a/Source/Chatbook/PromptGenerators/EmbeddingContext.wl b/Source/Chatbook/PromptGenerators/EmbeddingContext.wl
@@ -22,7 +22,7 @@ $smallContextStringLength = 8000;
 (*getSmallContextString*)
 getSmallContextString // beginDefinition;
 
-getSmallContextString // Options = { "IncludeSystemMessage" -> False };
+getSmallContextString // Options = { "IncludeSystemMessage" -> False, "SingleMessageTemplate" -> Automatic };
 
 getSmallContextString[ messages0: { ___Association }, opts: OptionsPattern[ ] ] := Enclose[
     Catch @ Module[ { messages, string },