mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-17 15:38:00 +00:00
Report prompt injection findings from the embedding pipeline
This commit is contained in:
2 files changed
+41
-4
No files matched your search
@@ -9,13 +9,12 @@ using AIStudio.Tools.Databases;
|
|||||||
using AIStudio.Tools.Databases.IndexStore;
|
using AIStudio.Tools.Databases.IndexStore;
|
||||||
using AIStudio.Tools.Databases.VectorStore;
|
using AIStudio.Tools.Databases.VectorStore;
|
||||||
using AIStudio.Tools.PluginSystem;
|
using AIStudio.Tools.PluginSystem;
|
||||||
|
using AIStudio.Tools.Security;
|
||||||
|
|
||||||
namespace AIStudio.Tools.Services;
|
namespace AIStudio.Tools.Services;
|
||||||
|
|
||||||
public sealed partial class DataSourceEmbeddingService(
|
public sealed partial class DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider,
|
||||||
SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider,
|
PromptInjectionGuardService guardService, ILogger<DataSourceEmbeddingService> logger) : BackgroundService
|
||||||
ILogger<DataSourceEmbeddingService> logger)
|
|
||||||
: BackgroundService
|
|
||||||
{
|
{
|
||||||
private const int VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD = 100_000;
|
private const int VECTOR_STORE_OPTIMIZATION_CHUNK_THRESHOLD = 100_000;
|
||||||
|
|
||||||
@@ -568,6 +567,14 @@ public sealed partial class DataSourceEmbeddingService(
|
|||||||
var lastError = inputFiles.LastError;
|
var lastError = inputFiles.LastError;
|
||||||
var failureDetails = inputFiles.Failures.ToList();
|
var failureDetails = inputFiles.Failures.ToList();
|
||||||
|
|
||||||
|
//
|
||||||
|
// Everything the runtime filters out of these files is reported once for the whole data
|
||||||
|
// source. A run 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 run is finished rather than in the middle of it.
|
||||||
|
//
|
||||||
|
await using var promptInjectionReportingScope = guardService.BeginAction();
|
||||||
|
|
||||||
foreach (var file in indexedFiles)
|
foreach (var file in indexedFiles)
|
||||||
{
|
{
|
||||||
token.ThrowIfCancellationRequested();
|
token.ThrowIfCancellationRequested();
|
||||||
|
|||||||
@@ -302,6 +302,9 @@ public sealed partial class RustService
|
|||||||
yield break;
|
yield break;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
var promptInjectionFindings = new List<PromptInjectionFinding>();
|
||||||
|
var promptInjectionRedactedCount = 0;
|
||||||
|
|
||||||
string? finalContentChunk;
|
string? finalContentChunk;
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
@@ -372,6 +375,21 @@ public sealed partial class RustService
|
|||||||
throw new InvalidOperationException($"Rust could not extract '{path}': {error.Message}");
|
throw new InvalidOperationException($"Rust could not extract '{path}': {error.Message}");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (processedEvent.PromptInjection is { } promptInjection)
|
||||||
|
{
|
||||||
|
//
|
||||||
|
// Not a failure: the passages were removed and the document around them is
|
||||||
|
// intact, so what remains still belongs into the index. It only has to reach
|
||||||
|
// the user, because from here on the indexed document is no longer the one
|
||||||
|
// sitting on their disk.
|
||||||
|
//
|
||||||
|
promptInjectionRedactedCount += promptInjection.RedactedCount;
|
||||||
|
if (promptInjection.Findings is { } findings)
|
||||||
|
promptInjectionFindings.AddRange(findings);
|
||||||
|
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
if (!string.IsNullOrWhiteSpace(processedEvent.Content))
|
if (!string.IsNullOrWhiteSpace(processedEvent.Content))
|
||||||
yield return (processedEvent.Content, sseEvent.TokenCount);
|
yield return (processedEvent.Content, sseEvent.TokenCount);
|
||||||
}
|
}
|
||||||
@@ -383,6 +401,18 @@ public sealed partial class RustService
|
|||||||
|
|
||||||
if (!string.IsNullOrWhiteSpace(finalContentChunk))
|
if (!string.IsNullOrWhiteSpace(finalContentChunk))
|
||||||
yield return (finalContentChunk, null);
|
yield return (finalContentChunk, null);
|
||||||
|
|
||||||
|
if (promptInjectionRedactedCount is 0)
|
||||||
|
yield break;
|
||||||
|
|
||||||
|
//
|
||||||
|
// Reported from here for the same reason as in ReadArbitraryFileData above: these two
|
||||||
|
// methods together are every way of reading a file, so they are the only two places
|
||||||
|
// where no caller can forget the report. Here it was missing, which is why a whole
|
||||||
|
// indexing run could filter documents without ever saying so.
|
||||||
|
//
|
||||||
|
var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
|
||||||
|
await guardService.ReportAsync(new(PromptInjectionSource.FileContent(path), promptInjectionFindings, promptInjectionRedactedCount));
|
||||||
}
|
}
|
||||||
|
|
||||||
private bool TryLogSseErrorMessage(string jsonContent, string path)
|
private bool TryLogSseErrorMessage(string jsonContent, string path)
|
||||||
|
|||||||
Reference in new issue
Block a user