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

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-11-05 19:27:08 +01:00
1 parent 7342c436dd
commit 188d17688a
16 files changed
+847 -78

No files matched your search

@@ -12382,6 +12382,9 @@ UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3016479965"] = "The text was
-- We tried to communicate with the LLM provider '{0}' (type={1}). Even after {2} retries, there were some problems with the request. The provider message is: '{3}'.
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3049689432"] = "We tried to communicate with the LLM provider '{0}' (type={1}). Even after {2} retries, there were some problems with the request. The provider message is: '{3}'."
-- We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments and the passages of your data sources, is probably too large for the selected model and provider. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T349154708"] = "We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments and the passages of your data sources, is probably too large for the selected model and provider. The provider message is: '{2}'"
-- Tried to communicate with the LLM provider '{0}'. There were some problems with the request. The provider message is: '{1}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3573577433"] = "Tried to communicate with the LLM provider '{0}'. There were some problems with the request. The provider message is: '{1}'"
@@ -12391,9 +12394,6 @@ UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T364882899"] = "The provider
-- We tried to communicate with the LLM provider '{0}' (type={1}). The required message format might be changed. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3759732886"] = "We tried to communicate with the LLM provider '{0}' (type={1}). The required message format might be changed. The provider message is: '{2}'"
-- We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments, is probably too large for the selected model and provider. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T4049517041"] = "We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments, is probably too large for the selected model and provider. The provider message is: '{2}'"
-- The provider '{0}' does not know the selected model. Please select another model.
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T411393889"] = "The provider '{0}' does not know the selected model. Please select another model."
@@ -15520,8 +15520,11 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::POWERPOINT::SLIDEDECKSPEC::T3520350217"] = "Pi
-- Standard augmentation process
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T1072508429"] = "Standard augmentation process"
-- This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T3240406069"] = "This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread."
-- 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.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T2379210749"] = "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."
-- This is the standard augmentation process, which augments the chat thread with as many retrieval contexts as the selected model can read.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T312973328"] = "This is the standard augmentation process, which augments the chat thread with as many retrieval contexts as the selected model can read."
-- The check of which passages fit your question failed. This answer uses all passages that were found.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T392269104"] = "The check of which passages fit your question failed. This answer uses all passages that were found."
@@ -12384,6 +12384,9 @@ UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3016479965"] = "Der Text war
-- We tried to communicate with the LLM provider '{0}' (type={1}). Even after {2} retries, there were some problems with the request. The provider message is: '{3}'.
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3049689432"] = "Wir haben versucht, mit dem LLM-Anbieter „{0}“ (Typ={1}) zu kommunizieren. Selbst nach {2} erneuten Versuchen gab es weiterhin Probleme mit der Anfrage. Die Meldung des Anbieters lautet: „{3}“."
-- We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments and the passages of your data sources, is probably too large for the selected model and provider. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T349154708"] = "Wir haben versucht, mit dem LLM-Anbieter „{0}“ (Typ={1}) zu kommunizieren. Die Daten des Chats, einschließlich aller Dateianhänge und der Textstellen aus Ihren Datenquellen, sind vermutlich zu groß für das ausgewählte Modell und den Anbieter. Die Nachricht des Anbieters lautet: „{2}“"
-- Tried to communicate with the LLM provider '{0}'. There were some problems with the request. The provider message is: '{1}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3573577433"] = "Es wurde versucht, mit dem LLM-Anbieter „{0}“ zu kommunizieren. Dabei sind Probleme bei der Anfrage aufgetreten. Die Meldung des Anbieters lautet: „{1}“"
@@ -12393,9 +12396,6 @@ UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T364882899"] = "Der Anbieter
-- We tried to communicate with the LLM provider '{0}' (type={1}). The required message format might be changed. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3759732886"] = "Wir haben versucht, mit dem LLM-Anbieter „{0}“ (Typ={1}) zu kommunizieren. Das erforderliche Nachrichtenformat hat sich möglicherweise geändert. Die Nachricht des Anbieters lautet: „{2}“"
-- We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments, is probably too large for the selected model and provider. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T4049517041"] = "Wir haben versucht, mit dem LLM-Anbieter „{0}“ (Typ={1}) zu kommunizieren. Die Daten des Chats, einschließlich aller Dateianhänge, sind vermutlich zu groß für das ausgewählte Modell und den Anbieter. Die Nachricht des Anbieters lautet: „{2}“"
-- The provider '{0}' does not know the selected model. Please select another model.
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T411393889"] = "Der Anbieter „{0}“ kennt das ausgewählte Modell nicht. Bitte wählen Sie ein anderes Modell aus."
@@ -15522,8 +15522,11 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::POWERPOINT::SLIDEDECKSPEC::T3520350217"] = "Bi
-- Standard augmentation process
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T1072508429"] = "Standard-Erweiterungsprozess"
-- This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T3240406069"] = "Dies ist der Standard-Erweiterungsprozess, bei dem alle abgerufenen Kontexte verwendet werden, um den Chatverlauf zu ergänzen."
-- 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.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T2379210749"] = "Nicht alle Textstellen aus Ihren Datenquellen passen in das, was das ausgewählte Modell verarbeiten kann. Diese Antwort musste auf {0} davon verzichten. Damit die KI alle Textstellen erhält, schalten Sie die semantische Suche für diesen Chat ein, senken Sie das Token-Limit Ihrer Datenquellen oder wählen Sie weniger Datenquellen."
-- This is the standard augmentation process, which augments the chat thread with as many retrieval contexts as the selected model can read.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T312973328"] = "Dies ist der Standard-Erweiterungsprozess, der den Chatverlauf um so viele abgerufene Kontexte ergänzt, wie das ausgewählte Modell verarbeiten kann."
-- The check of which passages fit your question failed. This answer uses all passages that were found.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T392269104"] = "Die Prüfung, welche Textstellen zu Ihrer Frage passen, ist fehlgeschlagen. Diese Antwort verwendet alle gefundenen Textstellen."
@@ -12384,6 +12384,9 @@ UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3016479965"] = "The text was
-- We tried to communicate with the LLM provider '{0}' (type={1}). Even after {2} retries, there were some problems with the request. The provider message is: '{3}'.
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3049689432"] = "We tried to communicate with the LLM provider '{0}' (type={1}). Even after {2} retries, there were some problems with the request. The provider message is: '{3}'."
-- We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments and the passages of your data sources, is probably too large for the selected model and provider. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T349154708"] = "We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments and the passages of your data sources, is probably too large for the selected model and provider. The provider message is: '{2}'"
-- Tried to communicate with the LLM provider '{0}'. There were some problems with the request. The provider message is: '{1}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3573577433"] = "Tried to communicate with the LLM provider '{0}'. There were some problems with the request. The provider message is: '{1}'"
@@ -12393,9 +12396,6 @@ UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T364882899"] = "The provider
-- We tried to communicate with the LLM provider '{0}' (type={1}). The required message format might be changed. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T3759732886"] = "We tried to communicate with the LLM provider '{0}' (type={1}). The required message format might be changed. The provider message is: '{2}'"
-- We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments, is probably too large for the selected model and provider. The provider message is: '{2}'
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T4049517041"] = "We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments, is probably too large for the selected model and provider. The provider message is: '{2}'"
-- The provider '{0}' does not know the selected model. Please select another model.
UI_TEXT_CONTENT["AISTUDIO::PROVIDER::BASEPROVIDER::T411393889"] = "The provider '{0}' does not know the selected model. Please select another model."
@@ -15522,8 +15522,11 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::POWERPOINT::SLIDEDECKSPEC::T3520350217"] = "Pi
-- Standard augmentation process
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T1072508429"] = "Standard augmentation process"
-- This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T3240406069"] = "This is the standard augmentation process, which uses all retrieval contexts to augment the chat thread."
-- 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.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T2379210749"] = "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."
-- This is the standard augmentation process, which augments the chat thread with as many retrieval contexts as the selected model can read.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T312973328"] = "This is the standard augmentation process, which augments the chat thread with as many retrieval contexts as the selected model can read."
-- The check of which passages fit your question failed. This answer uses all passages that were found.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::RAG::AUGMENTATIONPROCESSES::AUGMENTATIONONE::T392269104"] = "The check of which passages fit your question failed. This answer uses all passages that were found."
@@ -693,11 +693,14 @@ public abstract class BaseProvider : IProvider, ISecretId
/// <remarks>
/// Two answers mean this: one provider says so in the body of a bad request, another turns the
/// request down with 413 instead. For the user they are the same thing, and saying it in one
/// place is also what keeps both on one I18N key.
/// place is also what keeps both on one I18N key.<br/><br/>
/// The message names the passages of the data sources along with the attachments: the classic
/// RAG process sends them with every message, and while their budget leaves the rest of the
/// window to everything else, a long chat may still fill it, see PassageBudget.
/// </remarks>
/// <param name="providerMessage">What the provider itself said about the failure.</param>
/// <returns>The message to show.</returns>
private string GetContextTooLargeUserMessage(string? providerMessage) => string.Format(TB("We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments, is probably too large for the selected model and provider. The provider message is: '{2}'"), this.InstanceName, this.Provider, providerMessage);
private string GetContextTooLargeUserMessage(string? providerMessage) => string.Format(TB("We tried to communicate with the LLM provider '{0}' (type={1}). The data of the chat, including all file attachments and the passages of your data sources, is probably too large for the selected model and provider. The provider message is: '{2}'"), this.InstanceName, this.Provider, providerMessage);
/// <summary>
/// Sends a request and handles rate limiting by exponential backoff.
@@ -1445,7 +1448,9 @@ public abstract class BaseProvider : IProvider, ISecretId
/// With the ID, the settings also know what the provider's own model list reported, e.g., the
/// context window of a self-hosted model, see ProviderExtensions.GetAutomaticModelProfile. So
/// the tools of a request and the way images are passed on see the model as the configured
/// provider does.
/// provider does.<br/><br/>
/// The tokenizer comes along as well, so that the classic RAG process counts its passages for
/// the model as the chat counts the conversation, see PassageBudget.
/// </remarks>
public AIStudio.Settings.Provider CreateSettingsProvider(Model chatModel) => new()
{
@@ -1455,6 +1460,7 @@ public abstract class BaseProvider : IProvider, ISecretId
Model = chatModel,
InstanceName = this.InstanceName,
CapabilityOverrides = this.CapabilityOverrides,
TokenizerPath = this.TokenizerPath,
};
/// <summary>
@@ -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);
}
}
@@ -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);
@@ -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);
}
@@ -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);
}
@@ -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);
@@ -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)
{
@@ -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;
@@ -36,6 +36,7 @@
- Fixed AI Studio asking you to install Pandoc when you load a file as the input of an assistant, although PDF, Word, Markdown, and most other files need no Pandoc.
- Fixed dropping a file onto your custom prompt guide in the Prompt Optimizer. Once you switch on "Use custom prompt guide", a marked area appears where you can drop your guide or select it with a button.
- Fixed Semantic Search leaving out passages of your data sources which did not fit into a single search. The AI now gets every passage it finds, page by page, and a passage too long for a whole page is shortened instead of left out.
- Fixed chats with data sources failing when you switched off Semantic Search, because AI Studio sent more passages along with your message than the AI model can read. AI Studio now sends only as many passages as fit, taking turns among your data sources, and tells you with each answer when some passages had to be left out.
- Upgraded several libraries to improve security.
- Upgraded Tauri to v2.12.5.
- Upgraded MudBlazor to v9.12.0.