mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-14 12:33:48 +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
@@ -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);
|
||||
|
||||
Reference in new issue
Block a user