Fixed classic RAG sending more passages than the model can take (#1168)

app/MindWork AI Studio/Tools/RAG/AugmentationProcesses/AugmentationOne.cs

@@ -6,41 +6,43 @@ using AIStudio.Provider;
 using AIStudio.Settings;
 using AIStudio.Tools.PluginSystem;
 using AIStudio.Tools.Security;
+using AIStudio.Tools.Services;
 
 namespace AIStudio.Tools.RAG.AugmentationProcesses;
 
 public sealed class AugmentationOne : IAugmentationProcess
 {
     private static readonly ILogger<AugmentationOne> LOGGER = Program.LOGGER_FACTORY.CreateLogger<AugmentationOne>();
-    
+
     private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(AugmentationOne).Namespace, nameof(AugmentationOne));
-    
+
     #region Implementation of IAugmentationProcess
 
     /// <inheritdoc />
     public string TechnicalName => "AugmentationOne";
-    
+
     /// <inheritdoc />
     public string UIName => TB("Standard augmentation process");
-    
+
     /// <inheritdoc />
-    public string Description => TB("This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread.");
-    
+    public string Description => TB("This is the standard augmentation process, which augments the chat thread with as many retrieval contexts as the selected model can read.");
+
     /// <inheritdoc />
-    public async Task<AugmentationResult> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default)
+    public async Task<AugmentationResult> ProcessAsync(IProvider provider, Model chatModel, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IReadOnlyList<IRetrievalContext>> retrievalContextsByDataSource, CancellationToken token = default)
     {
         var settings = Program.SERVICE_PROVIDER.GetService<SettingsManager>()!;
-        
+        var retrievalContexts = retrievalContextsByDataSource.SelectMany(contexts => contexts).ToList();
+
         if(retrievalContexts.Count == 0)
         {
             LOGGER.LogWarning("No retrieval contexts were issued. Skipping the augmentation process.");
-            return new(chatThread, []);
+            return new(chatThread, [], [..retrievalContextsByDataSource.Select(_ => 0)]);
         }
 
         //
         // One report for the whole answer. Every passage is filtered once, before the check of
         // which passages fit: the check then judges exactly what the model would receive, and the
-        // model receives exactly what the check judged.
+        // model receives exactly what the check judged. The same holds for the budget below.
         //
         var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
         await using var reportingScope = guardService.BeginAction();
@@ -97,24 +99,116 @@ public sealed class AugmentationOne : IAugmentationProcess
                 LOGGER.LogWarning("Skipping retrieval context validation because no sufficiently trusted validation agent provider is available. Continuing augmentation with all retrieved contexts.");
             }
         }
-        
-        LOGGER.LogInformation($"Starting the augmentation process over {usedContexts.Count:###,###,###,###} of {retrievalContexts.Count:###,###,###,###} retrieved contexts.");
-        
+
         //
-        // We build a huge prompt from all retrieval contexts:
+        // Of what the check left, the model receives as much as its context window holds:
         //
-        var sb = new StringBuilder();
-        sb.AppendLine("The following useful information will help you in processing the user prompt:");
-        sb.AppendLine();
-        
-        // Let's convert all retrieval contexts to Markdown:
-        for (var index = 0; index < usedContexts.Count; index++)
-            sb.Append(usedContexts[index].ToMarkdown(index + 1, usedContexts.Count));
-        
-        // Add the augmented data to the chat thread:
-        chatThread.AugmentedData = sb.ToString();
-        return new(chatThread, usedContexts);
+        var selection = await SelectWithinBudgetAsync(provider.CreateSettingsProvider(chatModel), ByDataSource(retrievalContextsByDataSource, preparedContexts, usedContexts), guardService, reportingScope, token);
+        LOGGER.LogInformation($"Starting the augmentation process over {selection.Contexts.Count:###,###,###,##0} of {retrievalContexts.Count:###,###,###,##0} retrieved contexts; {selection.LeftOutCount:###,###,###,##0} of them did not fit into what the model can read.");
+
+        //
+        // The user learns about it with every answer: each answer rests on the passages it got,
+        // and each one may miss something else.
+        //
+        if (selection.LeftOutCount > 0)
+            await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.ContentCut, string.Format(TB("Not every passage found in your data sources fits into what the selected model can read. This answer had to do without {0} of them. To give the AI every passage, switch on Semantic Search for this chat, lower the token limit of your data sources, or choose fewer data sources."), selection.LeftOutCount.CompactCount())));
+
+        //
+        // Add the augmented data to the chat thread, as the very last step: the caller restricts
+        // the chat for these passages once this process returns, and a request whose RAG process
+        // failed goes on without it, so nothing may fail between the two:
+        //
+        chatThread.AugmentedData = WriteAugmentedData(selection);
+        return new(chatThread, selection.Contexts, selection.PassageCounts);
     }
 
     #endregion
+
+    /// <summary>
+    /// Writes what the model reads along with the user prompt: the passages which fit, and how many did not.
+    /// </summary>
+    /// <remarks>
+    /// The model learns about the passages which were left out, so that it does not take what it
+    /// read for everything the data sources hold on the user prompt.
+    /// </remarks>
+    /// <param name="selection">The passages which fit into the budget, and how many did not.</param>
+    /// <returns>The augmented data, or nothing when there is nothing to tell.</returns>
+    internal static string WriteAugmentedData(PassageSelection selection)
+    {
+        var sb = new StringBuilder();
+        if (selection.Contexts.Count > 0)
+        {
+            sb.AppendLine("The following useful information will help you in processing the user prompt:");
+            sb.AppendLine();
+
+            // Let's convert the retrieval contexts to Markdown:
+            for (var index = 0; index < selection.Contexts.Count; index++)
+                sb.Append(selection.Contexts[index].ToMarkdown(index + 1, selection.Contexts.Count));
+        }
+
+        if (selection.LeftOutCount > 0)
+            sb.AppendLine(selection.LeftOutCount is 1
+                ? "Of the passages which the search found in the data sources, 1 did not fit into this request and was left out. An answer based on the data sources may therefore be incomplete."
+                : $"Of the passages which the search found in the data sources, {selection.LeftOutCount} did not fit into this request and were left out. An answer based on the data sources may therefore be incomplete.");
+
+        return sb.ToString();
+    }
+
+    /// <summary>
+    /// Sorts the prepared contexts which the check of which passages fit left back to their data sources.
+    /// </summary>
+    /// <remarks>
+    /// The contexts were prepared in one batch, in the order of their data sources, and the check
+    /// hands back the very same objects. So each context left is found by reference, and keeps its
+    /// rank within its data source, whichever order the check reported them in.
+    /// </remarks>
+    /// <param name="retrievalContextsByDataSource">The retrieval contexts per data source, as the retrieval issued them.</param>
+    /// <param name="preparedContexts">All contexts, prepared in the same order.</param>
+    /// <param name="usedContexts">The prepared contexts the check left.</param>
+    /// <returns>The prepared contexts left of each data source, each in the order of their rank.</returns>
+    private static IReadOnlyList<IReadOnlyList<PreparedRetrievalContext>> ByDataSource(IReadOnlyList<IReadOnlyList<IRetrievalContext>> retrievalContextsByDataSource, IReadOnlyList<PreparedRetrievalContext> preparedContexts, IReadOnlyList<PreparedRetrievalContext> usedContexts)
+    {
+        var used = new HashSet<PreparedRetrievalContext>(usedContexts, ReferenceEqualityComparer.Instance);
+        var byDataSource = new List<IReadOnlyList<PreparedRetrievalContext>>(retrievalContextsByDataSource.Count);
+        var first = 0;
+        foreach (var contexts in retrievalContextsByDataSource)
+        {
+            byDataSource.Add(preparedContexts.Skip(first).Take(contexts.Count).Where(used.Contains).ToList());
+            first += contexts.Count;
+        }
+
+        return byDataSource;
+    }
+
+    /// <summary>
+    /// Selects the contexts which fit into the share of the model's context window the passages get, see PassageBudget.
+    /// </summary>
+    /// <remarks>
+    /// The tokens are counted with the tokenizer of the provider, or the built-in one where none is
+    /// configured, see ConversationTokenCounter. When the tokenizer does not answer, the characters
+    /// are counted against the same limit instead: a token of prose holds several characters, and
+    /// hardly ever less than one, so the passages then get less room rather than more.
+    /// </remarks>
+    /// <param name="settingsProvider">The settings of the provider and the model which answers.</param>
+    /// <param name="contextsByDataSource">The prepared contexts of each data source, each in the order of their rank.</param>
+    /// <param name="guardService">The prompt injection filter, for contexts which have to be shortened.</param>
+    /// <param name="reportingScope">The action the contexts were retrieved for.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <returns>The contexts which fit, and how many did not.</returns>
+    private static async Task<PassageSelection> SelectWithinBudgetAsync(AIStudio.Settings.Provider settingsProvider, IReadOnlyList<IReadOnlyList<PreparedRetrievalContext>> contextsByDataSource, PromptInjectionGuardService guardService, PromptInjectionReportingScope reportingScope, CancellationToken token)
+    {
+        var contextWindow = settingsProvider.GetModelProfile().Context;
+        var budget = PassageBudget.For(contextWindow, CountTokensAsync);
+        try
+        {
+            return await budget.SelectAsync(contextsByDataSource, guardService, reportingScope, token);
+        }
+        catch (Exception exception) when (contextWindow.IsKnown && exception is not OperationCanceledException)
+        {
+            LOGGER.LogWarning(exception, "Could not count the tokens of the retrieval contexts. Counting their characters against the same limit instead.");
+            return await PassageBudget.InCharacters(budget.Limit).SelectAsync(contextsByDataSource, guardService, reportingScope, token);
+        }
+
+        Task<int> CountTokensAsync(string text, CancellationToken cancellationToken) => Program.SERVICE_PROVIDER.GetRequiredService<ConversationTokenCounter>().CountTextAsync(settingsProvider, text, cancellationToken);
+    }
 }