mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-13 13:53:47 +00:00
Fixed prompt injection warnings waiting for other actions to finish (#1053)
This commit is contained in:
1 parent
a3c6b29313
commit
57dda9ff8b
25 files changed
+589
-146
No files matched your search
@@ -5,6 +5,7 @@ using AIStudio.Chat;
|
||||
using AIStudio.Provider;
|
||||
using AIStudio.Settings;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
namespace AIStudio.Tools.RAG.AugmentationProcesses;
|
||||
|
||||
@@ -37,7 +38,14 @@ public sealed class AugmentationOne : IAugmentationProcess
|
||||
}
|
||||
|
||||
var numTotalRetrievalContexts = retrievalContexts.Count;
|
||||
|
||||
|
||||
//
|
||||
// One report for the whole answer. With the check of which passages fit, every passage is
|
||||
// filtered twice: once for that check, and once more for the model.
|
||||
//
|
||||
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
|
||||
await using var reportingScope = guardService.BeginAction();
|
||||
|
||||
// Want the user to validate all retrieval contexts?
|
||||
if (settings.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation && chatThread.DataSourceOptions.AutomaticValidation)
|
||||
{
|
||||
@@ -48,7 +56,7 @@ public sealed class AugmentationOne : IAugmentationProcess
|
||||
try
|
||||
{
|
||||
// Let's validate all retrieval contexts:
|
||||
var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, retrievalContexts, token);
|
||||
var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, retrievalContexts, reportingScope, token);
|
||||
if (validationResults.Count == 0)
|
||||
LOGGER.LogWarning("Retrieval context validation returned no results. Continuing augmentation with all retrieved contexts.");
|
||||
else
|
||||
@@ -99,7 +107,7 @@ public sealed class AugmentationOne : IAugmentationProcess
|
||||
sb.AppendLine();
|
||||
|
||||
// Let's convert all retrieval contexts to Markdown:
|
||||
await retrievalContexts.AsMarkdown(sb, token);
|
||||
await retrievalContexts.AsMarkdown(sb, reportingScope, token);
|
||||
|
||||
// Add the augmented data to the chat thread:
|
||||
chatThread.AugmentedData = sb.ToString();
|
||||
|
||||
@@ -38,28 +38,46 @@ public static class IRetrievalContextExtensions
|
||||
contextBuilder.AppendLine($"- {link}");
|
||||
}
|
||||
|
||||
public static async Task<string> AsMarkdown(this IReadOnlyList<IRetrievalContext> retrievalContexts, StringBuilder? sb = null, CancellationToken token = default)
|
||||
/// <summary>
|
||||
/// Writes what the AI is told about several retrieval contexts, filtered for prompt injections.
|
||||
/// </summary>
|
||||
/// <param name="retrievalContexts">The contexts to describe.</param>
|
||||
/// <param name="sb">The builder to write into, or null for a new one.</param>
|
||||
/// <param name="reportingScope">The action the contexts are retrieved for, or null when they are retrieved for none.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>Everything written into the builder so far.</returns>
|
||||
public static async Task<string> AsMarkdown(this IReadOnlyList<IRetrievalContext> retrievalContexts, StringBuilder? sb = null, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
|
||||
{
|
||||
sb ??= new StringBuilder();
|
||||
var index = 0;
|
||||
|
||||
|
||||
//
|
||||
// One report for the whole retrieval run: a query may pull in dozens of contexts, and
|
||||
// the user wants to know that something was filtered, not to acknowledge it per context.
|
||||
//
|
||||
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
|
||||
await using var reportingScope = guardService.BeginAction();
|
||||
await using var retrievalScope = guardService.BeginAction(within: reportingScope);
|
||||
|
||||
foreach(var retrievalContext in retrievalContexts)
|
||||
{
|
||||
index++;
|
||||
await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, token);
|
||||
await retrievalContext.AsMarkdown(sb, index, retrievalContexts.Count, retrievalScope, token);
|
||||
}
|
||||
|
||||
return sb.ToString();
|
||||
}
|
||||
|
||||
public static async Task<string> AsMarkdown(this IRetrievalContext retrievalContext, StringBuilder? sb = null, int index = -1, int numTotalRetrievalContexts = -1, CancellationToken token = default)
|
||||
|
||||
/// <summary>
|
||||
/// Writes what the AI is told about one retrieval context, filtered for prompt injections.
|
||||
/// </summary>
|
||||
/// <param name="retrievalContext">The context to describe.</param>
|
||||
/// <param name="sb">The builder to write into, or null for a new one.</param>
|
||||
/// <param name="index">The number of the context among all of them, or -1 when it stands alone.</param>
|
||||
/// <param name="numTotalRetrievalContexts">How many contexts there are, or -1 when that is unknown.</param>
|
||||
/// <param name="reportingScope">The action the context is retrieved for, or null when it is retrieved for none.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>Everything written into the builder so far.</returns>
|
||||
public static async Task<string> AsMarkdown(this IRetrievalContext retrievalContext, StringBuilder? sb = null, int index = -1, int numTotalRetrievalContexts = -1, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
|
||||
{
|
||||
sb ??= new StringBuilder();
|
||||
var contextBuilder = new StringBuilder();
|
||||
@@ -141,7 +159,7 @@ public static class IRetrievalContextExtensions
|
||||
//
|
||||
async Task FilterWhatWeHaveSoFar()
|
||||
{
|
||||
var sanitized = await guardService.SanitizeAsync(contextBuilder.ToString(), source);
|
||||
var sanitized = await guardService.SanitizeAsync(contextBuilder.ToString(), source, reportingScope);
|
||||
contextBuilder.Clear();
|
||||
contextBuilder.Append(sanitized);
|
||||
}
|
||||
|
||||
@@ -11,7 +11,10 @@ namespace AIStudio.Tools.Security;
|
||||
/// The detection itself lives in the Rust runtime. File content is filtered while the runtime
|
||||
/// streams it, so it never passes through here; what this service adds is the path for content
|
||||
/// the runtime does not read itself — web pages and retrieval contexts — and the reporting the
|
||||
/// user sees.
|
||||
/// user sees.<br/><br/>
|
||||
/// The reporting works per user action, cf. BeginAction: each action tells the user once, when
|
||||
/// it ends, about what was filtered for it. One service serves the whole app, so the action has
|
||||
/// to travel with the content instead of being kept here.
|
||||
/// </remarks>
|
||||
public sealed class PromptInjectionGuardService(
|
||||
RustService rustService,
|
||||
@@ -27,8 +30,6 @@ public sealed class PromptInjectionGuardService(
|
||||
|
||||
private readonly ILogger detectionLogger = loggerFactory.CreateLogger(DETECTION_LOG_CATEGORY);
|
||||
private readonly Lock reportLock = new();
|
||||
private readonly List<PromptInjectionScanResult> pendingResults = [];
|
||||
private int openActions;
|
||||
|
||||
/// <summary>
|
||||
/// Filters prompt injections out of a text the runtime did not read itself, such as a web
|
||||
@@ -42,8 +43,9 @@ public sealed class PromptInjectionGuardService(
|
||||
/// </remarks>
|
||||
/// <param name="text">The content to filter.</param>
|
||||
/// <param name="source">Where the content came from, for the report shown to the user.</param>
|
||||
/// <param name="reportingScope">The action the content belongs to, or null when it belongs to none.</param>
|
||||
/// <returns>The content with any suspicious passages removed.</returns>
|
||||
public async Task<string> SanitizeAsync(string text, PromptInjectionSource source)
|
||||
public async Task<string> SanitizeAsync(string text, PromptInjectionSource source, PromptInjectionReportingScope? reportingScope = null)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(text))
|
||||
return text;
|
||||
@@ -59,7 +61,7 @@ public sealed class PromptInjectionGuardService(
|
||||
}
|
||||
|
||||
if (response.RedactedCount > 0)
|
||||
await this.ReportAsync(new(source, response.Findings, response.RedactedCount));
|
||||
await this.ReportAsync(new(source, response.Findings, response.RedactedCount), reportingScope);
|
||||
|
||||
return response.SanitizedText;
|
||||
}
|
||||
@@ -75,8 +77,9 @@ public sealed class PromptInjectionGuardService(
|
||||
/// the check cannot run, every text is passed through unchanged.
|
||||
/// </remarks>
|
||||
/// <param name="texts">The contents to filter, each with its source.</param>
|
||||
/// <param name="reportingScope">The action the contents belong to, or null when they belong to none.</param>
|
||||
/// <returns>The contents with any suspicious passages removed, in the order they came in.</returns>
|
||||
public async Task<IReadOnlyList<string>> SanitizeAsync(IReadOnlyList<PromptInjectionText> texts)
|
||||
public async Task<IReadOnlyList<string>> SanitizeAsync(IReadOnlyList<PromptInjectionText> texts, PromptInjectionReportingScope? reportingScope = null)
|
||||
{
|
||||
if (texts.Count is 0)
|
||||
return [];
|
||||
@@ -113,7 +116,9 @@ public sealed class PromptInjectionGuardService(
|
||||
|
||||
//
|
||||
// Findings are collected per source, not per text: a page whose content and title were
|
||||
// both filtered is one thing that happened to the user, not two.
|
||||
// both filtered is one thing that happened to the user, not two. Doing so here rather than
|
||||
// in the report keeps both counts: two texts with the same findings would look like one
|
||||
// text filtered twice there.
|
||||
//
|
||||
var findingsBySource = new Dictionary<PromptInjectionSource, (List<PromptInjectionFinding> Findings, int RedactedCount)>();
|
||||
for (var responseIndex = 0; responseIndex < indicesToScan.Count; responseIndex++)
|
||||
@@ -139,9 +144,9 @@ public sealed class PromptInjectionGuardService(
|
||||
// One scope around all sources, so a search across five pages reports once instead of
|
||||
// five times:
|
||||
//
|
||||
await using var reportingScope = this.BeginAction();
|
||||
await using var batchScope = this.BeginAction(within: reportingScope);
|
||||
foreach (var (source, aggregate) in findingsBySource)
|
||||
await this.ReportAsync(new(source, aggregate.Findings, aggregate.RedactedCount));
|
||||
await this.ReportAsync(new(source, aggregate.Findings, aggregate.RedactedCount), batchScope);
|
||||
|
||||
return sanitizedTexts;
|
||||
}
|
||||
@@ -150,25 +155,31 @@ public sealed class PromptInjectionGuardService(
|
||||
/// Records what was filtered out of one piece of content and tells the user about it.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Within a BeginAction scope the result is collected and reported together
|
||||
/// with the rest of that action. Outside of one it is reported immediately: a result that
|
||||
/// simply waited for the next scope would either never reach the user, or reach them as
|
||||
/// part of an unrelated action later on.
|
||||
/// Within an action, the result is collected and reported together with the rest of that
|
||||
/// action once it ends. Without one it is reported right away, and so is a result whose
|
||||
/// action has ended already: waiting for the next action would either never reach the user,
|
||||
/// or reach them as part of an unrelated action later on.<br/><br/>
|
||||
/// Only the action passed in counts, never what else is running. A run over a large data
|
||||
/// source takes minutes, and a message sent meanwhile must not wait for it, nor end up in the
|
||||
/// same report.
|
||||
/// </remarks>
|
||||
public async Task ReportAsync(PromptInjectionScanResult result)
|
||||
/// <param name="result">What was filtered out of the content.</param>
|
||||
/// <param name="reportingScope">The action the content belongs to, or null when it belongs to none.</param>
|
||||
public async Task ReportAsync(PromptInjectionScanResult result, PromptInjectionReportingScope? reportingScope = null)
|
||||
{
|
||||
if (!result.WasFiltered)
|
||||
return;
|
||||
|
||||
bool reportNow;
|
||||
lock (this.reportLock)
|
||||
{
|
||||
this.pendingResults.Add(result);
|
||||
reportNow = this.openActions is 0;
|
||||
if (FindCollectingScope(reportingScope) is { } collectingScope)
|
||||
{
|
||||
collectingScope.PendingResults.Add(result);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (reportNow)
|
||||
await this.ReportPendingAsync();
|
||||
await this.ReportNowAsync([result]);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@@ -177,44 +188,104 @@ public sealed class PromptInjectionGuardService(
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Results are collected until the action finishes, so the user gets one report about
|
||||
/// twenty documents instead of twenty reports. Actions may nest: only the outermost one
|
||||
/// reports.
|
||||
/// twenty documents instead of twenty reports. An action begun within another one is part
|
||||
/// of it, and only the outermost one reports. Unrelated actions report on their own, each
|
||||
/// as soon as it ends, however much they overlap.<br/><br/>
|
||||
/// The scope has to travel with the content: hand it to wherever that content is read or
|
||||
/// filtered, such as ReportAsync, SanitizeAsync, or RustService.ReadArbitraryFileData.
|
||||
/// </remarks>
|
||||
/// <param name="within">The action this one is part of, or null when it is an action of its own.</param>
|
||||
/// <returns>A scope that reports what was filtered once it is disposed.</returns>
|
||||
public ReportingScope BeginAction()
|
||||
{
|
||||
lock (this.reportLock)
|
||||
this.openActions++;
|
||||
public PromptInjectionReportingScope BeginAction(PromptInjectionReportingScope? within = null) => new(this, within);
|
||||
|
||||
return new(this);
|
||||
}
|
||||
|
||||
private async Task EndActionAsync()
|
||||
{
|
||||
lock (this.reportLock)
|
||||
{
|
||||
this.openActions--;
|
||||
|
||||
// An inner scope reports nothing: the action the user started is still running.
|
||||
if (this.openActions > 0)
|
||||
return;
|
||||
}
|
||||
|
||||
await this.ReportPendingAsync();
|
||||
}
|
||||
|
||||
private async Task ReportPendingAsync()
|
||||
/// <summary>
|
||||
/// Ends an action and reports what was collected for it.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A scope within an action which still runs has collected nothing, since its results went to
|
||||
/// the outermost scope. It collects only once the outer one has ended before it, e.g., for work
|
||||
/// which outlived the action that started it. Ending a scope a second time does nothing.
|
||||
/// </remarks>
|
||||
/// <param name="scope">The scope which ends.</param>
|
||||
internal async Task EndActionAsync(PromptInjectionReportingScope scope)
|
||||
{
|
||||
List<PromptInjectionScanResult> results;
|
||||
lock (this.reportLock)
|
||||
{
|
||||
if (this.pendingResults.Count is 0)
|
||||
if (!scope.IsOpen)
|
||||
return;
|
||||
|
||||
results = [..this.pendingResults];
|
||||
this.pendingResults.Clear();
|
||||
scope.IsOpen = false;
|
||||
results = [..scope.PendingResults];
|
||||
scope.PendingResults.Clear();
|
||||
}
|
||||
|
||||
if (results.Count > 0)
|
||||
await this.ReportNowAsync(results);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Finds the scope which collects a result of the given one: the outermost scope of its action
|
||||
/// which is still running.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Call under the report lock only.
|
||||
/// </remarks>
|
||||
/// <param name="scope">The scope the result was reported for, or null.</param>
|
||||
/// <returns>The collecting scope, or null when the result has to be reported right away.</returns>
|
||||
private static PromptInjectionReportingScope? FindCollectingScope(PromptInjectionReportingScope? scope)
|
||||
{
|
||||
PromptInjectionReportingScope? collectingScope = null;
|
||||
for (; scope is not null; scope = scope.Parent)
|
||||
{
|
||||
if (scope.IsOpen)
|
||||
collectingScope = scope;
|
||||
}
|
||||
|
||||
return collectingScope;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Combines the results of one report into one result per source, in the order the sources came in.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// An action may filter the same content twice: the classic data source search filters a passage when
|
||||
/// it checks whether the passage fits the question, and once more when it hands the passage to the
|
||||
/// model. Both times, the result is the same, so a result equal to one already seen for its source is
|
||||
/// dropped; the user hears about the passage once. The findings carry no position, though, so two
|
||||
/// passages of one file with exactly the same findings count as one as well. That mostly happens
|
||||
/// where two neighboring chunks overlap, which is the same passage indeed.<br/><br/>
|
||||
/// Other results of one source, such as two passages of one file or the mails of one mailbox,
|
||||
/// become a single result with all their findings and the sum of the removed passages. The report
|
||||
/// counts sources, and one file is one source, however many of its passages were filtered.
|
||||
/// </remarks>
|
||||
/// <param name="results">The results of one report.</param>
|
||||
/// <returns>One result per source.</returns>
|
||||
private static IReadOnlyList<PromptInjectionScanResult> CombineBySource(IReadOnlyList<PromptInjectionScanResult> results)
|
||||
{
|
||||
if (results.Count < 2)
|
||||
return results;
|
||||
|
||||
return results.GroupBy(result => result.Source).Select(CombineOneSource).ToList();
|
||||
|
||||
static PromptInjectionScanResult CombineOneSource(IGrouping<PromptInjectionSource, PromptInjectionScanResult> resultsOfSource)
|
||||
{
|
||||
List<PromptInjectionScanResult> distinctResults = [];
|
||||
foreach (var result in resultsOfSource)
|
||||
{
|
||||
if (!distinctResults.Any(known => known.RedactedCount == result.RedactedCount && known.Findings.SequenceEqual(result.Findings)))
|
||||
distinctResults.Add(result);
|
||||
}
|
||||
|
||||
return distinctResults.Count is 1
|
||||
? distinctResults[0]
|
||||
: new(resultsOfSource.Key, distinctResults.SelectMany(result => result.Findings).ToList(), distinctResults.Sum(result => result.RedactedCount));
|
||||
}
|
||||
}
|
||||
|
||||
private async Task ReportNowAsync(IReadOnlyList<PromptInjectionScanResult> reportedResults)
|
||||
{
|
||||
var results = CombineBySource(reportedResults);
|
||||
var totalCount = results.Sum(result => result.RedactedCount);
|
||||
this.detectionLogger.LogWarning(
|
||||
"Detected and removed {PassageCount} potentially dangerous passage(s) in {SourceCount} content source(s).",
|
||||
@@ -230,12 +301,4 @@ public sealed class PromptInjectionGuardService(
|
||||
if (settingsManager.ConfigurationData.App.ShowPromptInjectionAlert)
|
||||
await MessageBus.INSTANCE.SendMessage<PromptInjectionAlertMessage>(null, Event.SHOW_PROMPT_INJECTION_ALERT, new(results));
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Reports everything filtered during one user action when it goes out of scope.
|
||||
/// </summary>
|
||||
public sealed class ReportingScope(PromptInjectionGuardService guardService) : IAsyncDisposable
|
||||
{
|
||||
public async ValueTask DisposeAsync() => await guardService.EndActionAsync();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
namespace AIStudio.Tools.Security;
|
||||
|
||||
/// <summary>
|
||||
/// One user action, such as sending a message, whose filtered content is reported to the user together.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Created through PromptInjectionGuardService.BeginAction. Content which belongs to the action carries
|
||||
/// the scope down to where it is read or filtered; disposing the scope ends the action.
|
||||
/// </remarks>
|
||||
/// <param name="guardService">The service which reports what was filtered.</param>
|
||||
/// <param name="parent">The action this scope is part of, or null when it begins an action of its own.</param>
|
||||
public sealed class PromptInjectionReportingScope(PromptInjectionGuardService guardService, PromptInjectionReportingScope? parent) : IAsyncDisposable
|
||||
{
|
||||
/// <summary>
|
||||
/// The action this scope is part of, or null when it begins an action of its own.
|
||||
/// </summary>
|
||||
internal PromptInjectionReportingScope? Parent => parent;
|
||||
|
||||
/// <summary>
|
||||
/// What was filtered while this scope was the outermost one of its action still running.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Read and written under the lock of the guard service only.
|
||||
/// </remarks>
|
||||
internal List<PromptInjectionScanResult> PendingResults { get; } = [];
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a value indicating whether the action is still running.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Read and written under the lock of the guard service only.
|
||||
/// </remarks>
|
||||
internal bool IsOpen { get; set; } = true;
|
||||
|
||||
public async ValueTask DisposeAsync() => await guardService.EndActionAsync(this);
|
||||
}
|
||||
@@ -77,7 +77,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
var textChunker = new TextChunker(rustService, logger);
|
||||
return
|
||||
[
|
||||
new FileSourceIndexer(settingsManager, rustService, guardService, textChunker, logger),
|
||||
new FileSourceIndexer(settingsManager, rustService, textChunker, logger),
|
||||
new MailboxIndexer(settingsManager, rustService, guardService, textChunker, logger),
|
||||
];
|
||||
}
|
||||
@@ -702,9 +702,10 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
/// <param name="refreshMode">Why the run was started, for the log.</param>
|
||||
/// <param name="share">How much work the run does at most, as its kind of data source says.</param>
|
||||
/// <param name="runControl">The run, which records the embedding settings it works for.</param>
|
||||
/// <param name="promptInjectionScope">The run as one action of the prompt injection filter.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The context of the run, or null when there is no run.</returns>
|
||||
private async Task<IndexedRunContext?> PrepareIndexedRunAsync(IIndexedDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, RunShare share, DataSourceRunControl runControl, CancellationToken token)
|
||||
private async Task<IndexedRunContext?> PrepareIndexedRunAsync(IIndexedDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, RunShare share, DataSourceRunControl runControl, PromptInjectionReportingScope promptInjectionScope, CancellationToken token)
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Starting background embedding hash check for data source '{DataSourceName}' ({DataSourceId}). RefreshMode={RefreshMode}.",
|
||||
@@ -831,7 +832,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
pass.FailedDocumentCount,
|
||||
refreshMode);
|
||||
|
||||
return new IndexedRunContext(dataSource, embeddingProvider, embeddingProvider.CreateProvider(), vectorStore, indexStore, manifest, share, pass, settingsManager, this.UpsertStatus, logger);
|
||||
return new IndexedRunContext(dataSource, embeddingProvider, embeddingProvider.CreateProvider(), vectorStore, indexStore, manifest, share, pass, promptInjectionScope, settingsManager, this.UpsertStatus, logger);
|
||||
}
|
||||
|
||||
private async Task ProcessDataSourceAsync(IIndexedDataSource dataSource, DataSourceEmbeddingRefreshMode refreshMode, DataSourceRunControl runControl, CancellationToken token)
|
||||
@@ -848,7 +849,13 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
var outcome = IndexedRunOutcome.DONE;
|
||||
try
|
||||
{
|
||||
var context = await this.PrepareIndexedRunAsync(dataSource, refreshMode, indexer.Share, runControl, token);
|
||||
//
|
||||
// Everything the prompt injection filter removes from the documents of this run is
|
||||
// reported once, when the run is over. A step over a few thousand documents which
|
||||
// removes something in forty of them is one thing that happened to the user, not forty.
|
||||
//
|
||||
await using var promptInjectionScope = guardService.BeginAction();
|
||||
var context = await this.PrepareIndexedRunAsync(dataSource, refreshMode, indexer.Share, runControl, promptInjectionScope, token);
|
||||
if (context is null)
|
||||
return;
|
||||
|
||||
|
||||
@@ -16,9 +16,9 @@ The parts:
|
||||
finds and reads the documents of one run, and `TrackChanges` / `StopTracking` notice changes on their own.
|
||||
`FileSourceIndexer` is the reference: a file system watcher per data source, fingerprints over path, size and
|
||||
write time.
|
||||
- **`IndexedRunContext`** - one prepared run: both stores, the embedding provider, the manifest and the
|
||||
collection. `IndexDocumentAsync` embeds and stores one document; the cleanup methods remove what a failed
|
||||
attempt left behind.
|
||||
- **`IndexedRunContext`** - one prepared run: both stores, the embedding provider, the manifest, the
|
||||
collection, and the scope for reporting prompt injections. `IndexDocumentAsync` embeds and stores one
|
||||
document; the cleanup methods remove what a failed attempt left behind.
|
||||
- **`EmbeddingDocument`** - one document: its key, its index row, its display name, and how to read its chunks.
|
||||
- **`DocumentRunProgress`** - counts the documents, records indexed and failed ones in the stores, publishes the
|
||||
status, and completes or pauses the run.
|
||||
@@ -49,6 +49,10 @@ To add a kind of data source:
|
||||
returns: it decides whether the data source carries on in another run, or stops because its embedding
|
||||
provider failed for everything. A run which takes more than twice its share without asking once ends up as
|
||||
a warning in the log, while every other data source waits.
|
||||
7. Filter everything the indexer reads: files through `ExtractedFileText.ReadAsync`, any other content through
|
||||
`PromptInjectionGuardService.SanitizeAsync`, both with `IndexedRunContext.PromptInjectionScope`. The service
|
||||
opens that scope per run, so the user hears about a run once, when it is over. Begin no action of your own
|
||||
in the indexer, or the user hears about one run several times.
|
||||
|
||||
Rules which are easy to break:
|
||||
- **A document key is not a path.** Only files use their full path as the key. Never pass a key through the
|
||||
|
||||
@@ -18,13 +18,14 @@ internal static class ExtractedFileText
|
||||
/// <param name="filePath">The path of the file.</param>
|
||||
/// <param name="embeddingProvider">The embedding provider whose tokenizer counts the pieces.</param>
|
||||
/// <param name="reportAs">The source the user is told about when passages were filtered out, or null for the file itself.</param>
|
||||
/// <param name="reportingScope">The action the file is read for, which tells the user about filtered passages together with the rest of it.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The text, empty when the file has none.</returns>
|
||||
public static async Task<SegmentedText> ReadAsync(RustService rustService, string filePath, EmbeddingProvider embeddingProvider, PromptInjectionSource? reportAs, CancellationToken token)
|
||||
public static async Task<SegmentedText> ReadAsync(RustService rustService, string filePath, EmbeddingProvider embeddingProvider, PromptInjectionSource? reportAs, PromptInjectionReportingScope reportingScope, CancellationToken token)
|
||||
{
|
||||
var segments = new List<TextSegment>();
|
||||
|
||||
await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, reportAs, token))
|
||||
await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, reportAs, reportingScope, token))
|
||||
{
|
||||
var normalized = TextChunker.NormalizeSegment(segment.Content);
|
||||
if (!string.IsNullOrWhiteSpace(normalized))
|
||||
|
||||
@@ -32,10 +32,9 @@ namespace AIStudio.Tools.Services.Indexing;
|
||||
/// </remarks>
|
||||
/// <param name="settingsManager">The settings, read for whether local data sources refresh on their own.</param>
|
||||
/// <param name="rustService">The runtime, which extracts the text of the files.</param>
|
||||
/// <param name="guardService">The prompt injection filter, whose findings are reported once per run.</param>
|
||||
/// <param name="textChunker">Cuts the text of a file into chunks.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed partial class FileSourceIndexer(SettingsManager settingsManager, RustService rustService, PromptInjectionGuardService guardService, TextChunker textChunker, ILogger logger) : IIndexedSourceIndexer
|
||||
internal sealed partial class FileSourceIndexer(SettingsManager settingsManager, RustService rustService, TextChunker textChunker, ILogger logger) : IIndexedSourceIndexer
|
||||
{
|
||||
private const string OFFICE_LOCK_FILE_PREFIX = "~$";
|
||||
private const bool IMAGE_EMBEDDING_ENABLED = false;
|
||||
@@ -166,14 +165,6 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
|
||||
partition.ToRead.Count);
|
||||
progress.Publish();
|
||||
|
||||
//
|
||||
// Everything the runtime filters out of these files is reported once for the whole step.
|
||||
// A step over a few thousand documents which removes something in forty of them is one
|
||||
// thing that happened to the user, not forty. The scope ends with this method, so the
|
||||
// report arrives when the step is over rather than in the middle of it.
|
||||
//
|
||||
await using var promptInjectionReportingScope = guardService.BeginAction();
|
||||
|
||||
for (var fileIndex = 0; fileIndex < partition.ToRead.Count; fileIndex++)
|
||||
{
|
||||
token.ThrowIfCancellationRequested();
|
||||
@@ -440,11 +431,11 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
|
||||
: string.Join("; ", reasons);
|
||||
}
|
||||
|
||||
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
private async IAsyncEnumerable<EmbeddingChunk> StreamEmbeddingChunksAsync(string filePath, IDataSource dataSource, EmbeddingProvider embeddingProvider, PromptInjectionReportingScope reportingScope, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
var options = DataSourceEmbeddingService.GetChunkingOptions(dataSource, embeddingProvider);
|
||||
var strategy = TextChunker.GetStrategyForFile(filePath);
|
||||
var content = await ExtractedFileText.ReadAsync(rustService, filePath, embeddingProvider, null, token);
|
||||
var content = await ExtractedFileText.ReadAsync(rustService, filePath, embeddingProvider, null, reportingScope, token);
|
||||
|
||||
await foreach (var chunk in textChunker.SplitAsync(content, strategy, options, embeddingProvider, token))
|
||||
yield return chunk;
|
||||
@@ -731,7 +722,7 @@ internal sealed partial class FileSourceIndexer(SettingsManager settingsManager,
|
||||
file.FullName,
|
||||
this.CreateEmbeddingStateFile(dataSource, file, fingerprint, 0, DateTimeOffset.UtcNow),
|
||||
file.Name,
|
||||
token => this.StreamEmbeddingChunksAsync(file.FullName, dataSource, context.EmbeddingProvider, token));
|
||||
token => this.StreamEmbeddingChunksAsync(file.FullName, dataSource, context.EmbeddingProvider, context.PromptInjectionScope, token));
|
||||
|
||||
private EmbeddingStateFile CreateEmbeddingStateFile(IDataSource dataSource, FileInfo file, string fingerprint, int chunkCount, DateTimeOffset embeddedAtUtc)
|
||||
{
|
||||
|
||||
@@ -3,6 +3,7 @@ using AIStudio.Settings;
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Databases.VectorStore;
|
||||
using AIStudio.Tools.PluginSystem;
|
||||
using AIStudio.Tools.Security;
|
||||
|
||||
namespace AIStudio.Tools.Services.Indexing;
|
||||
|
||||
@@ -26,10 +27,11 @@ namespace AIStudio.Tools.Services.Indexing;
|
||||
/// <param name="manifest">What the index stores about the data source, made to match the current embedding configuration.</param>
|
||||
/// <param name="share">How much work the run does at most, as its kind of data source says.</param>
|
||||
/// <param name="pass">The pass the run is a step of, which knows what failed during the steps before.</param>
|
||||
/// <param name="promptInjectionScope">The run as one action of the prompt injection filter, which tells the user about everything it removed once the run is over.</param>
|
||||
/// <param name="settingsManager">The settings, which the provider reads when it embeds.</param>
|
||||
/// <param name="publishStatus">Hands a new status of the data source to the embedding service.</param>
|
||||
/// <param name="logger">The logger of the embedding service, so the log reads the same whoever writes it.</param>
|
||||
internal sealed class IndexedRunContext(IIndexedDataSource dataSource, EmbeddingProvider embeddingProvider, IProvider provider, VectorStoreClient vectorStore, IndexStoreClient indexStore, DataSourceEmbeddingManifest manifest, RunShare share, IndexingPass pass, SettingsManager settingsManager, Action<DataSourceEmbeddingStatus> publishStatus, ILogger logger)
|
||||
internal sealed class IndexedRunContext(IIndexedDataSource dataSource, EmbeddingProvider embeddingProvider, IProvider provider, VectorStoreClient vectorStore, IndexStoreClient indexStore, DataSourceEmbeddingManifest manifest, RunShare share, IndexingPass pass, PromptInjectionReportingScope promptInjectionScope, SettingsManager settingsManager, Action<DataSourceEmbeddingStatus> publishStatus, ILogger logger)
|
||||
{
|
||||
/// <summary>
|
||||
/// After how many stored chunks the collection is optimized while a run is still going.
|
||||
@@ -64,6 +66,16 @@ internal sealed class IndexedRunContext(IIndexedDataSource dataSource, Embedding
|
||||
|
||||
public IndexingPass Pass => pass;
|
||||
|
||||
/// <summary>
|
||||
/// The run as one action of the prompt injection filter.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Everything read for the run carries it down to the filter, so a run over a few thousand
|
||||
/// documents which removes something in forty of them tells the user once, when it is over,
|
||||
/// rather than forty times. Every kind of data source gets this from here, so none can forget it.
|
||||
/// </remarks>
|
||||
public PromptInjectionReportingScope PromptInjectionScope => promptInjectionScope;
|
||||
|
||||
/// <summary>
|
||||
/// When the run first asked whether it did its share of the work, or null while it never did.
|
||||
/// </summary>
|
||||
|
||||
@@ -45,7 +45,7 @@ internal sealed partial class MailboxIndexer
|
||||
await connector.FetchAttachmentAsync(uid, part, file, token);
|
||||
|
||||
// Filtered by the runtime while it reads, and reported as content of the mailbox, not of a file the user never saw:
|
||||
var content = await ExtractedFileText.ReadAsync(rustService, path, context.EmbeddingProvider, PromptInjectionSource.MailContent(mailbox.Name), token);
|
||||
var content = await ExtractedFileText.ReadAsync(rustService, path, context.EmbeddingProvider, PromptInjectionSource.MailContent(mailbox.Name), context.PromptInjectionScope, token);
|
||||
return new(part, name, MailPartTextState.EXTRACTED, content, TextChunker.GetStrategyForFile(path));
|
||||
}
|
||||
catch (OperationCanceledException) when (token.IsCancellationRequested)
|
||||
|
||||
@@ -148,7 +148,7 @@ internal sealed partial class MailboxIndexer
|
||||
new(text.Body, source),
|
||||
new(text.Subject, source),
|
||||
..attachmentParts.Select(part => new PromptInjectionText(MailTextNormalization.NormalizeHeaderValue(part.FileName), source)),
|
||||
]);
|
||||
], context.PromptInjectionScope);
|
||||
|
||||
var attachments = new List<MailAttachmentText>(attachmentParts.Count);
|
||||
for (var index = 0; index < attachmentParts.Count; index++)
|
||||
|
||||
@@ -244,22 +244,15 @@ internal sealed partial class MailboxIndexer(SettingsManager settingsManager, Ru
|
||||
progress.RecordUnchanged(plannedFolders.Sum(folder => folder.Plan.KeptUids.Count));
|
||||
progress.Publish();
|
||||
|
||||
//
|
||||
// Whatever the prompt injection filter removes from the mails of this run is reported
|
||||
// once, when the second pass is done, rather than once per mail.
|
||||
//
|
||||
var encounteredKeys = new HashSet<string>(StringComparer.Ordinal);
|
||||
var workedThrough = true;
|
||||
await using (guardService.BeginAction())
|
||||
foreach (var folder in plannedFolders)
|
||||
{
|
||||
foreach (var folder in plannedFolders)
|
||||
{
|
||||
if (await this.SyncNewMailsAsync(context, mailbox, connector, folder, progress, encounteredKeys, token))
|
||||
continue;
|
||||
if (await this.SyncNewMailsAsync(context, mailbox, connector, folder, progress, encounteredKeys, token))
|
||||
continue;
|
||||
|
||||
workedThrough = false;
|
||||
break;
|
||||
}
|
||||
workedThrough = false;
|
||||
break;
|
||||
}
|
||||
|
||||
var sourceHash = BuildSourceHash(plannedFolders);
|
||||
|
||||
@@ -37,7 +37,8 @@ public sealed partial class RustService
|
||||
/// attachment: nothing leaves the app on that path, so there is nothing to warn about, and
|
||||
/// reporting it there would warn a second time when the file is actually sent.
|
||||
/// </param>
|
||||
public async Task<FileExtractionResult> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false, bool reportPromptInjections = true, CancellationToken token = default)
|
||||
/// <param name="reportingScope">The action the file is read for, or null when it is read for none.</param>
|
||||
public async Task<FileExtractionResult> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false, bool reportPromptInjections = true, PromptInjectionReportingScope? reportingScope = null, CancellationToken token = default)
|
||||
{
|
||||
//
|
||||
// The runtime filters prompt injections while it streams the file. Doing it there rather
|
||||
@@ -249,7 +250,7 @@ public sealed partial class RustService
|
||||
// where the content never leaves the app.
|
||||
//
|
||||
if (reportPromptInjections)
|
||||
await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
|
||||
await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount), reportingScope);
|
||||
|
||||
//
|
||||
// Filtering does not change the outcome: the passages were removed and the document
|
||||
@@ -263,12 +264,6 @@ public sealed partial class RustService
|
||||
};
|
||||
}
|
||||
|
||||
public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
|
||||
{
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, string.Empty, null, token))
|
||||
yield return segment.Content;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Reads the content of a file for embedding, each piece together with its token count.
|
||||
/// </summary>
|
||||
@@ -279,15 +274,17 @@ public sealed partial class RustService
|
||||
/// itself. A file which only exists while it is read, e.g. the attachment of a mail, names
|
||||
/// where its content came from instead.
|
||||
/// </param>
|
||||
/// <param name="reportingScope">The action the file is read for, or null when it is read for none.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The pieces of the content, in order.</returns>
|
||||
public async IAsyncEnumerable<ArbitraryFileDataSegment> StreamArbitraryFileDataWithTokenCounts(
|
||||
string path,
|
||||
EmbeddingProvider embeddingProvider,
|
||||
PromptInjectionSource? reportAs = null,
|
||||
PromptInjectionReportingScope? reportingScope = null,
|
||||
[EnumeratorCancellation] CancellationToken token = default)
|
||||
{
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, reportAs, token))
|
||||
await foreach (var segment in this.StreamArbitraryFileDataCore(path, embeddingProvider.TokenizerPath, reportAs, reportingScope, token))
|
||||
{
|
||||
if (segment.TokenCount is { } tokenCount)
|
||||
{
|
||||
@@ -318,14 +315,13 @@ public sealed partial class RustService
|
||||
|
||||
private async IAsyncEnumerable<(string Content, int? TokenCount, int? PageNumber)> StreamArbitraryFileDataCore(
|
||||
string path,
|
||||
bool extractImages,
|
||||
bool includeTokenCount,
|
||||
string tokenizerPath,
|
||||
PromptInjectionSource? reportAs,
|
||||
PromptInjectionReportingScope? reportingScope,
|
||||
[EnumeratorCancellation] CancellationToken token)
|
||||
{
|
||||
var streamId = Guid.NewGuid().ToString();
|
||||
var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}&include_token_count={includeTokenCount}&tokenizer_path={Uri.EscapeDataString(tokenizerPath)}";
|
||||
var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images=false&include_token_count=true&tokenizer_path={Uri.EscapeDataString(tokenizerPath)}";
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
|
||||
using var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, token);
|
||||
|
||||
@@ -339,10 +335,7 @@ public sealed partial class RustService
|
||||
path,
|
||||
responseBody);
|
||||
|
||||
if (includeTokenCount)
|
||||
throw new InvalidOperationException($"Rust could not extract and count '{path}'. HTTP {(int)response.StatusCode} ({response.ReasonPhrase}): {responseBody}");
|
||||
|
||||
yield break;
|
||||
throw new InvalidOperationException($"Rust could not extract and count '{path}'. HTTP {(int)response.StatusCode} ({response.ReasonPhrase}): {responseBody}");
|
||||
}
|
||||
|
||||
var promptInjectionFindings = new List<PromptInjectionFinding>();
|
||||
@@ -372,12 +365,7 @@ public sealed partial class RustService
|
||||
catch (JsonException)
|
||||
{
|
||||
if (this.TryLogSseErrorMessage(jsonContent, path))
|
||||
{
|
||||
if (includeTokenCount)
|
||||
throw new InvalidOperationException($"Rust could not extract and count a segment from '{path}'. See the runtime log for details.");
|
||||
|
||||
continue;
|
||||
}
|
||||
throw new InvalidOperationException($"Rust could not extract and count a segment from '{path}'. See the runtime log for details.");
|
||||
|
||||
this.logger?.LogError("Failed to deserialize SSE event: {JsonContent}", jsonContent);
|
||||
}
|
||||
@@ -385,7 +373,7 @@ public sealed partial class RustService
|
||||
if (sseEvent is null)
|
||||
continue;
|
||||
|
||||
var processedEvent = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages);
|
||||
var processedEvent = ContentStreamSseHandler.ProcessEvent(sseEvent, extractImages: false);
|
||||
if (processedEvent.Error is { } error)
|
||||
{
|
||||
// A notice says something about the file without failing the read, so the
|
||||
@@ -461,7 +449,7 @@ public sealed partial class RustService
|
||||
// indexing run could filter documents without ever saying so.
|
||||
//
|
||||
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
|
||||
await guardService.ReportAsync(new(reportAs ?? PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
|
||||
await guardService.ReportAsync(new(reportAs ?? PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount), reportingScope);
|
||||
}
|
||||
|
||||
private bool TryLogSseErrorMessage(string jsonContent, string path)
|
||||
|
||||
+2
-2
@@ -303,7 +303,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
|
||||
// each, and so on. Otherwise, the data source offered first would take the budget, and the
|
||||
// others would get what it left over.
|
||||
//
|
||||
await using (guardService.BeginAction())
|
||||
await using (var reportingScope = guardService.BeginAction())
|
||||
{
|
||||
var mostPassages = pages.Select(page => page.Contexts.Count).DefaultIfEmpty(0).Max();
|
||||
for (var rank = 0; rank < mostPassages; rank++)
|
||||
@@ -314,7 +314,7 @@ public sealed class SemanticSearchTool(SettingsManager settingsManager, DataSour
|
||||
continue;
|
||||
|
||||
var retrievalContext = pages[index].Contexts[rank];
|
||||
var passage = await retrievalContext.AsMarkdown(index: passageCount + 1, token: token);
|
||||
var passage = await retrievalContext.AsMarkdown(index: passageCount + 1, reportingScope: reportingScope, token: token);
|
||||
|
||||
// A passage too long for what is left makes room for shorter ones after it:
|
||||
if (textContent.Length + passage.Length > MAX_RESULT_CHARACTERS)
|
||||
|
||||
Reference in new issue
Block a user