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);
+    }
 }

app/MindWork AI Studio/Tools/RAG/AugmentationResult.cs

@@ -7,9 +7,13 @@ namespace AIStudio.Tools.RAG;
 /// </summary>
 /// <remarks>
 /// The contexts the AI received can be fewer than the ones retrieved, e.g., when a check of which
-/// passages fit the question left some out. Whoever lists the sources and passages of the answer
-/// takes them from here, so the user sees what the AI read and nothing else.
+/// passages fit the question left some out, or when not all of them fit into what the model can
+/// read. Whoever lists the sources and passages of the answer takes them from here, so the user
+/// sees what the AI read and nothing else. Whoever restricts the chat for the data sources which
+/// brought passages into it counts them here as well, so that a data source none of whose passages
+/// reached the AI restricts nothing.
 /// </remarks>
 /// <param name="ChatThread">The augmented chat thread.</param>
 /// <param name="UsedContexts">The retrieval contexts handed to the AI, in the order it received them.</param>
-public sealed record AugmentationResult(ChatThread ChatThread, IReadOnlyList<PreparedRetrievalContext> UsedContexts);
+/// <param name="PassageCounts">How many retrieval contexts of each data source the AI received, in the order the data sources were handed in.</param>
+public sealed record AugmentationResult(ChatThread ChatThread, IReadOnlyList<PreparedRetrievalContext> UsedContexts, IReadOnlyList<int> PassageCounts);

app/MindWork AI Studio/Tools/RAG/IAugmentationProcess.cs

@@ -24,10 +24,11 @@ public interface IAugmentationProcess
     /// Starts the augmentation process.
     /// </summary>
     /// <param name="provider">The LLM provider. Gets used, e.g., for automatic retrieval context validation.</param>
+    /// <param name="chatModel">The model which answers. Its context window decides how much room the retrieval contexts get.</param>
     /// <param name="lastUserPrompt">The last user prompt that was issued by the user.</param>
     /// <param name="chatThread">The chat thread.</param>
-    /// <param name="retrievalContexts">The retrieval contexts that were issued by the retrieval process.</param>
+    /// <param name="retrievalContextsByDataSource">The retrieval contexts that were issued by the retrieval process, per data source, each in the order of their rank.</param>
     /// <param name="token">The cancellation token.</param>
     /// <returns>The altered chat thread, together with the retrieval contexts handed to the AI.</returns>
-    public Task<AugmentationResult> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default);
+    public Task<AugmentationResult> ProcessAsync(IProvider provider, Model chatModel, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IReadOnlyList<IRetrievalContext>> retrievalContextsByDataSource, CancellationToken token = default);
 }

app/MindWork AI Studio/Tools/RAG/PassageBudget.cs

@@ -0,0 +1,173 @@
+using AIStudio.Models;
+using AIStudio.Tools.Security;
+using AIStudio.Tools.ToolCallingSystem.ToolCallingImplementations.SemanticSearch;
+
+namespace AIStudio.Tools.RAG;
+
+/// <summary>
+/// How much room the passages of the classic RAG process may take in a request, and which of them fit into it.
+/// </summary>
+/// <remarks>
+/// <para>
+/// The classic RAG process sends its passages along with every message, in the system prompt. A
+/// chunk is as long as the embedding model takes at once, by default 8,192 tokens, so the default
+/// of ten passages of a single data source may already exceed what many models read, and the
+/// provider then refuses the whole message. The budget keeps the passages within a share of the
+/// model's context window, see For.
+/// </para>
+/// <para>
+/// The size of a passage is that of its Markdown as the model reads it, measured after the prompt
+/// injection filter: the filter may make a text longer, and the budget has to judge what the model
+/// receives. Which number a passage gets depends on which passages fit, so it is measured with the
+/// longest header it can get.
+/// </para>
+/// </remarks>
+/// <param name="limit">How much the passages may take together, in the unit of measure; at least 1.</param>
+/// <param name="measure">Measures a text in the unit of the budget, e.g., in tokens.</param>
+public sealed class PassageBudget(int limit, Func<string, CancellationToken, Task<int>> measure)
+{
+    /// <summary>
+    /// The share of a known context window the passages may take.
+    /// </summary>
+    /// <remarks>
+    /// The rest of the window stays for everything else of the request: the conversation, its
+    /// attachments, the system prompt, the tools, and the answer. Half of it leaves a long
+    /// conversation room, while the passages still get as much as the window of most models holds
+    /// of them.
+    /// </remarks>
+    internal const double SHARE_OF_THE_CONTEXT_WINDOW = 0.5;
+
+    /// <summary>
+    /// How many characters the passages may take when nothing is known about the context window of the model.
+    /// </summary>
+    /// <remarks>
+    /// As many as one page of Semantic Search returns, so that both ways of searching hand the
+    /// model at most the same amount of text when nothing more is known.
+    /// </remarks>
+    internal const int CHARACTERS_WITHOUT_A_KNOWN_CONTEXT_WINDOW = SemanticSearchTool.MAX_RESULT_CHARACTERS;
+
+    /// <summary>
+    /// How much the passages may take together, in the unit of the budget.
+    /// </summary>
+    public int Limit { get; } = limit >= 1 ? limit : throw new ArgumentOutOfRangeException(nameof(limit), limit, "A budget has to hold at least 1.");
+
+    /// <summary>
+    /// A budget which counts characters.
+    /// </summary>
+    /// <param name="limit">How many characters the passages may take together; at least 1.</param>
+    /// <returns>The budget.</returns>
+    public static PassageBudget InCharacters(int limit) => new(limit, (text, _) => Task.FromResult(text.Length));
+
+    /// <summary>
+    /// The budget for a model with the given context window.
+    /// </summary>
+    /// <remarks>
+    /// A known window is counted in tokens, as the model counts it. Of an unknown one, nothing can
+    /// be said, so the passages get as many characters as one page of Semantic Search.
+    /// </remarks>
+    /// <param name="contextWindow">The context window of the model which answers.</param>
+    /// <param name="countTokensAsync">Counts the tokens of a text for that model.</param>
+    /// <returns>The budget.</returns>
+    public static PassageBudget For(ContextWindow contextWindow, Func<string, CancellationToken, Task<int>> countTokensAsync) => contextWindow.IsKnown
+        ? new(Math.Max(1, (int)(contextWindow.DefaultTokens * SHARE_OF_THE_CONTEXT_WINDOW)), countTokensAsync)
+        : InCharacters(CHARACTERS_WITHOUT_A_KNOWN_CONTEXT_WINDOW);
+
+    /// <summary>
+    /// Selects the passages which fit into the budget.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// The data sources take turns, rank by rank: the best passage of each data source comes before
+    /// the second best of any, so that a data source with many long passages does not crowd out the
+    /// others. A passage which does not fit into what is left is skipped, and the next ones still
+    /// get their chance: the classic RAG process has no further pages which could bring it later.
+    /// </para>
+    /// <para>
+    /// A passage longer than the whole budget is shortened until it fits, and filtered again, see
+    /// IRetrievalContextExtensions.ShortenToFitAsync. It then has to fit into what is left like any
+    /// other passage.
+    /// </para>
+    /// </remarks>
+    /// <param name="contextsByDataSource">The prepared passages of each data source, each in the order of their rank.</param>
+    /// <param name="guardService">The prompt injection filter, for passages which have to be shortened.</param>
+    /// <param name="reportingScope">The action the passages were retrieved for, or null when they were retrieved for none.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <returns>The passages which fit, grouped by data source, how many of each data source fit, and how many did not.</returns>
+    public async Task<PassageSelection> SelectAsync(IReadOnlyList<IReadOnlyList<PreparedRetrievalContext>> contextsByDataSource, PromptInjectionGuardService guardService, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
+    {
+        var candidateCount = contextsByDataSource.Sum(contexts => contexts.Count);
+        var mostPassages = contextsByDataSource.Select(contexts => contexts.Count).DefaultIfEmpty(0).Max();
+        var selected = contextsByDataSource.Select(_ => new List<PreparedRetrievalContext>()).ToList();
+        var remaining = this.Limit;
+        var leftOutCount = 0;
+
+        for (var rank = 0; rank < mostPassages; rank++)
+        {
+            for (var index = 0; index < contextsByDataSource.Count; index++)
+            {
+                if (rank >= contextsByDataSource[index].Count)
+                    continue;
+
+                // A budget which is used up holds nothing more, so there is nothing to measure:
+                if (remaining is 0)
+                {
+                    leftOutCount++;
+                    continue;
+                }
+
+                var context = contextsByDataSource[index][rank];
+                var size = await this.MeasureAsync(context, candidateCount, token);
+                while (size > this.Limit && await this.ShortenAsync(context, size, candidateCount, guardService, reportingScope, token) is { } shorterContext)
+                {
+                    context = shorterContext;
+                    size = await this.MeasureAsync(context, candidateCount, token);
+                }
+
+                if (size > remaining)
+                {
+                    leftOutCount++;
+                    continue;
+                }
+
+                selected[index].Add(context);
+                remaining -= size;
+            }
+        }
+
+        return new([..selected.SelectMany(contexts => contexts)], [..selected.Select(contexts => contexts.Count)], leftOutCount);
+    }
+
+    /// <summary>
+    /// Measures a passage as the model reads it, in the unit of the budget.
+    /// </summary>
+    private Task<int> MeasureAsync(PreparedRetrievalContext context, int candidateCount, CancellationToken token) => measure(WithLongestHeader(context, candidateCount), token);
+
+    /// <summary>
+    /// Shortens a passage which is longer than the whole budget by about as much as it exceeds it.
+    /// </summary>
+    /// <remarks>
+    /// The budget may count something else than characters, e.g., tokens, while a passage is
+    /// shortened by characters. Both grow roughly in proportion, so the passage keeps the share of
+    /// its characters which the budget holds of its size. Where the proportion is off, the caller
+    /// measures again and shortens once more.
+    /// </remarks>
+    /// <param name="context">The passage.</param>
+    /// <param name="size">The size of the passage in the unit of the budget, more than its limit.</param>
+    /// <param name="candidateCount">How many passages there are to select from.</param>
+    /// <param name="guardService">The prompt injection filter.</param>
+    /// <param name="reportingScope">The action the passage was retrieved for, or null.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <returns>The shorter passage, or null when nothing of it is left to cut.</returns>
+    private async Task<PreparedRetrievalContext?> ShortenAsync(PreparedRetrievalContext context, int size, int candidateCount, PromptInjectionGuardService guardService, PromptInjectionReportingScope? reportingScope, CancellationToken token)
+    {
+        var characters = WithLongestHeader(context, candidateCount).Length;
+        var maxCharacters = (int)((long)characters * this.Limit / size);
+        var shorterContext = await context.ShortenToFitAsync(maxCharacters, shorter => WithLongestHeader(shorter, candidateCount).Length, guardService, reportingScope, token);
+        return WithLongestHeader(shorterContext, candidateCount).Length < characters ? shorterContext : null;
+    }
+
+    /// <summary>
+    /// Writes a passage with the longest header it can get: neither its number nor the count of the passages the model receives exceeds the count of the passages to select from.
+    /// </summary>
+    private static string WithLongestHeader(PreparedRetrievalContext context, int candidateCount) => context.ToMarkdown(candidateCount, candidateCount);
+}

app/MindWork AI Studio/Tools/RAG/PassageSelection.cs

@@ -0,0 +1,9 @@
+namespace AIStudio.Tools.RAG;
+
+/// <summary>
+/// Which passages of the classic RAG process fit into its budget, see PassageBudget.SelectAsync.
+/// </summary>
+/// <param name="Contexts">The passages which fit, grouped by data source in the order the data sources came in, and within a data source in the order of their rank. A passage which had to be shortened is here in its shorter form.</param>
+/// <param name="PassageCounts">How many passages of each data source fit, in the order the data sources came in.</param>
+/// <param name="LeftOutCount">How many passages did not fit.</param>
+public sealed record PassageSelection(IReadOnlyList<PreparedRetrievalContext> Contexts, IReadOnlyList<int> PassageCounts, int LeftOutCount);

app/MindWork AI Studio/Tools/RAG/RAGProcesses/AISrcSelWithRetCtxVal.cs

@@ -181,7 +181,7 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
             //
             // Trigger the retrieval part of the (R)AG process:
             //
-            var dataContexts = new List<IRetrievalContext>();
+            var dataContexts = new List<IReadOnlyList<IRetrievalContext>>();
             if (proceedWithRAG)
             {
                 //
@@ -192,17 +192,17 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
                     retrievalTasks.Add(dataSource.RetrieveDataAsync(lastUserPrompt, chatThreadWithoutWaitingAIBlock, token));
             
                 //
-                // Wait for all retrieval tasks to finish, and count the passages of each data source.
-                // The tasks are in the order of the selected data sources:
+                // Wait for all retrieval tasks to finish. The tasks are in the order of the selected
+                // data sources, and so are the passages: each data source keeps its place, also when
+                // it failed, so that the augmentation can let the data sources take turns when not all
+                // passages fit, and its counts match the data sources:
                 //
-                var passageCounts = new int[retrievalTasks.Count];
                 for (var index = 0; index < retrievalTasks.Count; index++)
                 {
+                    IReadOnlyList<IRetrievalContext> retrievedContexts = [];
                     try
                     {
-                        var retrievedContexts = await retrievalTasks[index];
-                        dataContexts.AddRange(retrievedContexts);
-                        passageCounts[index] = retrievedContexts.Count;
+                        retrievedContexts = await retrievalTasks[index];
                     }
                     catch (OperationCanceledException) when (token.IsCancellationRequested)
                     {
@@ -213,13 +213,9 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
                     {
                         LOGGER.LogError(e, "An error occurred during the retrieval process.");
                     }
-                }
 
-                //
-                // Right after the retrieval, before the augmentation: the passages are in the chat
-                // from here on, and a step below which ends early must not leave it unrestricted:
-                //
-                RequireOutboundDataRestriction(chatThread, selectedDataSources, passageCounts, settings.ConfigurationData.DataSourceSecurity.MinimumOutboundDataRestriction);
+                    dataContexts.Add(retrievedContexts);
+                }
             }
 
             //
@@ -229,9 +225,16 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
             if (proceedWithRAG)
             {
                 var augmentationProcess = new AugmentationOne();
-                var augmentation = await augmentationProcess.ProcessAsync(provider, lastUserPrompt, chatThread, dataContexts, token);
+                var augmentation = await augmentationProcess.ProcessAsync(provider, chatModel, lastUserPrompt, chatThread, dataContexts, token);
                 chatThread = augmentation.ChatThread;
                 usedContexts = augmentation.UsedContexts;
+
+                //
+                // Right after the augmentation, before anything else: the passages are in the chat
+                // from here on, and a step below which ends early must not leave it unrestricted.
+                // Only the passages the AI receives count:
+                //
+                RequireOutboundDataRestriction(chatThread, selectedDataSources, augmentation.PassageCounts, settings.ConfigurationData.DataSourceSecurity.MinimumOutboundDataRestriction);
             }
 
             //
@@ -268,8 +271,8 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
     /// contexts handed to it.
     /// </summary>
     /// <remarks>
-    /// Only those: a passage which the check of which passages fit left out never reached the AI,
-    /// and the answer does not rest on it. Only a text answer keeps passages; every answer keeps
+    /// Only those: a passage which the check of which passages fit left out, or which did not fit
+    /// into what the model can read, never reached the AI, and the answer does not rest on it. Only a text answer keeps passages; every answer keeps
     /// its sources.
     /// </remarks>
     /// <param name="aiAnswerContent">The content of the answer.</param>
@@ -284,13 +287,17 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
     }
 
     /// <summary>
-    /// Restricts where the chat may still send data, because of the passages the retrieval brought in.
+    /// Restricts where the chat may still send data, because of the passages the AI received.
     /// </summary>
     /// <remarks>
     /// Semantic Search asks the same place, DataSourceRequirements, so that both ways of searching
-    /// demand the same of the same data sources. Only the data sources with passages count: one
-    /// which found nothing leaves the chat as it was. A passage the validation agent drops later
-    /// still counts, which is the safe side. The data security and the provider confidence were
+    /// demand the same of the same data sources. Only the data sources whose passages reached the
+    /// AI count: one which found nothing leaves the chat as it was, and so does one whose passages
+    /// the validation agent dropped, or which did not fit into what the model can read, see
+    /// PassageBudget. Neither the model nor any of its tools ever saw those passages; the
+    /// validation agent did, but it has no tools. Otherwise, a data source of public papers could
+    /// fill the budget while the passages of a confidential one stayed out, and the chat would
+    /// still lose the web search for good. The data security and the provider confidence were
     /// raised before the retrieval already, for every selected data source, because the agents see
     /// the passages as well; the data sources with passages are among them, so asking again would
     /// change nothing.<br/><br/>
@@ -300,7 +307,7 @@ public sealed class AISrcSelWithRetCtxVal : IRagProcess
     /// </remarks>
     /// <param name="chatThread">The chat thread.</param>
     /// <param name="selectedDataSources">The data sources searched.</param>
-    /// <param name="passageCounts">How many passages each data source returned, in the same order.</param>
+    /// <param name="passageCounts">How many passages of each data source the AI received, in the same order.</param>
     /// <param name="organizationMinimum">The least strict outbound data restriction the organization allows.</param>
     internal static void RequireOutboundDataRestriction(ChatThread chatThread, IReadOnlyList<IDataSource> selectedDataSources, IReadOnlyList<int> passageCounts, OutboundDataRestriction organizationMinimum)
     {

app/MindWork AI Studio/Tools/Services/ConversationTokenCounter.cs

@@ -206,8 +206,17 @@ public sealed class ConversationTokenCounter(RustService rustService, DatabaseCl
     /// second copy of the whole conversation in memory. Hashing is not free, but it is two orders of
     /// magnitude cheaper than tokenizing the same bytes, so the trade pays for itself on the first
     /// repeat.
+    ///
+    /// The classic RAG process counts its passages here as well, see PassageBudget, so that they are
+    /// counted with the same tokenizer as the rest of the conversation, and a passage which comes up
+    /// with every message of a chat is counted once.
     /// </remarks>
-    private async Task<int> CountTextAsync(Provider provider, string text, CancellationToken token)
+    /// <param name="provider">The configured provider, which decides the tokenizer.</param>
+    /// <param name="text">The text to count.</param>
+    /// <param name="token">Ends the counting when nobody needs the answer anymore.</param>
+    /// <returns>The tokens of the text.</returns>
+    /// <exception cref="InvalidOperationException">The tokenizer did not answer.</exception>
+    public async Task<int> CountTextAsync(Provider provider, string text, CancellationToken token)
     {
         if (string.IsNullOrWhiteSpace(text))
             return 0;