Skip to content

Commit

Permalink
Merge pull request #934 from WolframResearch/feature/improved-rag-rer…
Browse files Browse the repository at this point in the history
…anking-and-filtering

Improved RAG reranking and filtering to optimize token usage
  • Loading branch information
rhennigan authored Nov 19, 2024
2 parents ecaaf94 + 41f92c0 commit 57f7f4a
Show file tree
Hide file tree
Showing 6 changed files with 162 additions and 40 deletions.
1 change: 1 addition & 0 deletions .vscode/settings.json
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
"Qwen",
"Rasterizer",
"Reparse",
"rerank",
"Rescrape",
"Roboto",
"RSNB",
Expand Down
9 changes: 2 additions & 7 deletions LLMConfiguration/Personas/NotebookAssistant/Pre.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,2 @@
You are a helpful Wolfram Language assistant named Notebook Assistant. Your job is to do the following:

* Offer Wolfram Language suggestions based on previous inputs.
* Help fix errors.
* Explain error messages so the user understands what went wrong.
* Explain outputs that may be confusing to inexperienced users.
* Assist the user with other topics as needed. While you specialize in Wolfram Language, you can help with any topic.
You are a helpful Wolfram Language assistant named Notebook Assistant.
You specialize in Wolfram Language, but you can help with any topic.
11 changes: 7 additions & 4 deletions Source/Chatbook/ChatTitle.wl
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,10 @@ Needs[ "Wolfram`Chatbook`Common`" ];
(* ::**************************************************************************************************************:: *)
(* ::Section::Closed:: *)
(*Configuration*)
$maxTitleLength = 30;
$hardMaxTitleLength = 40;
$maxContextLength = 100000; (* characters *)
$maxTitleLength = 30;
$hardMaxTitleLength = 40;
$maxContextLength = 100000; (* characters *)
$multimodalTitleContext = False;

$titlePrompt := "\
Please come up with a meaningful window title for the current chat conversation using no more than \
Expand Down Expand Up @@ -79,7 +80,9 @@ generateChatTitleAsync[ messages: $$chatMessages, callback_, temperature: Automa
instructions = ConfirmBy[ TemplateApply[ $titlePrompt, short ], StringQ, "Prompt" ];

context = ConfirmMatch[
Block[ { $multimodalMessages = True }, expandMultimodalString @ instructions ],
Block[ { $multimodalMessages = TrueQ @ $multimodalTitleContext },
expandMultimodalString @ instructions
],
_String | { __ },
"Context"
];
Expand Down
151 changes: 137 additions & 14 deletions Source/Chatbook/PromptGenerators/RelatedDocumentation.wl
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ $documentationSnippetsCacheDirectory := $documentationSnippetsCacheDirectory =
$resourceSnippetsCacheDirectory := $resourceSnippetsCacheDirectory =
ChatbookFilesDirectory @ { "DocumentationSnippets", "ResourceSystem" };

$rerankMethod := CurrentChatSettings[ "DocumentationRerankMethod" ];

(* ::**************************************************************************************************************:: *)
(* ::Section::Closed:: *)
(*Messages*)
Expand Down Expand Up @@ -192,7 +194,8 @@ prependRelatedDocsHeader // endDefinition;


$relatedDocsStringFilteredHeader =
"IMPORTANT: Here are some Wolfram documentation snippets that you should use to respond.\n\n";
"IMPORTANT: Here are some Wolfram documentation snippets that were retrieved based on semantic similarity to the \
current context. Please use them if they can help answer the user's latest message.\n\n";

$relatedDocsStringUnfilteredHeader =
"Here are some Wolfram documentation snippets that you may find useful.\n\n";
Expand All @@ -202,12 +205,64 @@ $relatedDocsStringUnfilteredHeader =
(*filterSnippets*)
filterSnippets // beginDefinition;

filterSnippets[ messages_, uris: { __String }, filter_, filterCount_Integer? Positive ] := Enclose[
Catch @ Module[ { snippets, inserted, transcript, xml, instructions, response, pages },

filterSnippets[ messages_, uris: { __String }, Except[ True ], filterCount_ ] := Enclose[
ConfirmMatch[ makeDocSnippets @ uris, { ___String }, "Snippets" ],
throwInternalFailure
];


filterSnippets[
messages_,
uris: { __String },
True,
filterCount_Integer? Positive
] /; $rerankMethod === "rerank-english-v3.0" (* EXPERIMENTAL *) := Enclose[
Catch @ Module[ { snippets, inserted, transcript, instructions, resp, results, idx, ranked },

snippets = ConfirmMatch[ makeDocSnippets @ uris, { ___String }, "Snippets" ];
If[ ! TrueQ @ filter, Throw @ snippets ];
setProgressDisplay[ "Choosing relevant documentation" ];
inserted = insertContextPrompt @ messages;
transcript = ConfirmBy[ getSmallContextString @ inserted, StringQ, "Transcript" ];

instructions = ConfirmBy[
TemplateApply[ $documentationRerankPrompt, <| "Transcript" -> transcript |> ],
StringQ,
"Prompt"
];

resp = ServiceExecute[
"Cohere",
"RawRerank",
{
"model" -> "rerank-english-v3.0",
"query" -> instructions,
"documents" -> snippets
}
];

If[ FailureQ @ resp, throwTop @ resp ];

results = ConfirmMatch[ resp[ "results" ], { __Association }, "Results" ];

idx = ConfirmMatch[
Select[ results, #[ "relevance_score" ] > 0.01 & ][[ All, "index" ]] + 1,
{ ___Integer },
"Indices"
];

ranked = ConfirmMatch[ snippets[[ idx ]], { ___String }, "Ranked" ];

Take[ ranked, UpTo[ filterCount ] ]
],
throwInternalFailure
];


filterSnippets[ messages_, uris: { __String }, True, filterCount_Integer? Positive ] := Enclose[
Catch @ Module[ { snippets, inserted, transcript, xml, instructions, response, pages },

snippets = ConfirmMatch[ makeDocSnippets @ uris, { ___String }, "Snippets" ];
setProgressDisplay[ "Choosing relevant documentation" ];
inserted = insertContextPrompt @ messages;
transcript = ConfirmBy[ getSmallContextString @ inserted, StringQ, "Transcript" ];
Expand All @@ -228,31 +283,54 @@ filterSnippets[ messages_, uris: { __String }, filter_, filterCount_Integer? Pos

response = StringTrim @ ConfirmBy[
LogChatTiming[
llmSynthesize @ instructions,
llmSynthesize[ instructions, <| "StopTokens" -> "\"CasualChat\"" |> ],
"WaitForFilterSnippetsTask"
] // withApproximateProgress[ 0.5 ],
StringQ,
"Response"
];

pages = ConfirmMatch[ makeDocSnippets @ StringCases[ response, uris ], { ___String }, "Pages" ];
pages = ConfirmMatch[ makeDocSnippets @ selectSnippetsFromJSON[ response, uris ], { ___String }, "Pages" ];

pages
],
throwInternalFailure
];


filterSnippets // endDefinition;



$bestDocumentationPrompt = StringTemplate[ "\
Your task is to read a chat transcript between a user and assistant, and then select the most relevant \
Wolfram Language documentation snippets that could help the assistant answer the user's latest message. \
Your task is to read a chat transcript between a user and assistant, and then select any relevant Wolfram Language \
documentation snippets that could help the assistant answer the user's latest message.
Each snippet is uniquely identified by a URI (always starts with 'paclet:' or 'https://resources.wolframcloud.com').
Choose up to %%FilteredCount%% documentation snippets that would help answer the user's MOST RECENT message. \
Respond only with the corresponding URIs of the snippets and nothing else. \
If there are no relevant pages, respond with just the string \"none\".
Choose up to %%FilteredCount%% documentation snippets that would help answer the user's MOST RECENT message.
Respond with JSON in the following format:
{
\"AssistantType\": assistantType,
\"Snippets\": [
{\"URI\": uri1, \"Score\": score1},
{\"URI\": uri2, \"Score\": score2},
...
]
}
For \"AssistantType\", specify the type of assistant that should handle the user's message:
\"Computational\": The user's message requires a computational response.
\"Knowledge\": The user's message requires a knowledge-based response.
\"CasualChat\": The user's message is casual and could be answered by a non-specialist. For example, simple greetings or general questions.
Specify a score as any number from 1 to 5 for your chosen snippets using the following rubric:
1: The snippet is completely irrelevant to the user's message or has no usefulness.
2: The snippet is somewhat related, but the assistant could easily answer the user's message without it.
3: The snippet is related and might help the assistant answer the user's message.
4: The snippet is very relevant and would significantly help the assistant answer the user's message.
5: It would be impossible for the assistant to answer the user's message correctly without this snippet.
Here is the chat transcript:
Expand All @@ -266,11 +344,56 @@ Here are the available documentation snippets to choose from:
%%Snippets%%
</snippets>
Reminder: Choose up to %%FilteredCount%% documentation snippets that would help answer the user's MOST RECENT message. \
Respond only with the corresponding URIs of the snippets and nothing else. \
If there are no relevant pages, respond with just the string \"none\".\
Reminder: Choose up to %%FilteredCount%% documentation snippets that would help answer the user's MOST RECENT message.
You can (and should) skip snippets that are not relevant to the user's message or are redundant.
Respond only with the specified JSON and nothing else.
If there are no relevant pages, respond with [].
", Delimiters -> "%%" ];



$documentationRerankPrompt = StringTemplate[ "\
Read the chat transcript between a user and assistant, and then give me the best Wolfram Language documentation \
snippet that could help the assistant answer the user's latest message.
The snippet does not need to exactly answer the user's message if it can be easily generalized.
Prefer built-in system symbols over other resources.
Simpler is better.
Here is the chat transcript:
<transcript>
%%Transcript%%
</transcript>\
", Delimiters -> "%%" ];

(* ::**************************************************************************************************************:: *)
(* ::Subsubsection::Closed:: *)
(*selectSnippetsFromJSON*)
selectSnippetsFromJSON // beginDefinition;

selectSnippetsFromJSON[ response_String, uris_List ] := Enclose[
Catch @ Module[ { jsonString, jsonData, selected },
jsonString = ConfirmBy[ First[ StringCases[ response, Longest[ "{" ~~ __ ~~ "}" ], 1 ], None ], StringQ ];
jsonData = ConfirmBy[ Quiet @ Developer`ReadRawJSONString @ jsonString, AssociationQ ];
selected = ConfirmMatch[ Select[ jsonData[ "Snippets" ], #[ "Score" ] >= 3 & ], { __ } ];
ConfirmMatch[
Intersection[ Cases[ selected, KeyValuePattern[ "URI" -> uri_String ] :> StringTrim @ uri ], uris ],
{ __String }
]
],
selectSnippetsFromString[ response, uris ] &
];

selectSnippetsFromJSON // endDefinition;

(* ::**************************************************************************************************************:: *)
(* ::Subsubsection::Closed:: *)
(*selectSnippetsFromString*)
selectSnippetsFromString // beginDefinition;
selectSnippetsFromString[ response_String, uris: { ___String } ] := StringCases[ response, uris ];
selectSnippetsFromString // endDefinition;

(* ::**************************************************************************************************************:: *)
(* ::Subsubsection::Closed:: *)
(*snippetXML*)
Expand Down
23 changes: 14 additions & 9 deletions Source/Chatbook/Prompting.wl
Original file line number Diff line number Diff line change
Expand Up @@ -114,14 +114,20 @@ $basePromptComponents[ "GeneralInstructionsHeader" ] = "\
# General Instructions
";

$basePromptComponents[ "NotebooksPreamble" ] = "\
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]
Expand Down Expand Up @@ -256,7 +262,6 @@ $basePromptComponents[ "WolframLanguageStyle" ] = "
* Always use proper naming conventions for your variables (e.g. lowerCamelCase)
* Never use single capital letters to represent variables (e.g. use `a Sin[k x + \[Phi]]` instead of `A Sin[k x + \[Phi]]`)
* Prefer modern Wolfram Language symbols and methods
* Many new symbols have been added to WL since your knowledge cutoff date, so check documentation as needed
* When creating plots, add options such as labels and legends to make them easier to understand";

$basePromptComponents[ "WolframLanguageEvaluatorTool" ] = "\
Expand Down
7 changes: 1 addition & 6 deletions Source/Chatbook/Tools/Examples.wl
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,7 @@ $fullExamplesKeys :=
With[ { selected = Keys @ $selectedTools },
Select[
{
"AstroGraphicsDocumentation",
"FileSystemTree",
"FractionalDerivatives",
"NaturalLanguageInput",
"PlotEvaluate",
"TemporaryDirectory"
"NaturalLanguageInput"
},
ContainsAll[ selected, $exampleDependencies[ #1 ] ] &
]
Expand Down

0 comments on commit 57f7f4a

Please sign in to comment.