Fixed prompt injection warnings waiting for other actions to finish (#1053)

app/MindWork AI Studio/Tools/RAG/AugmentationProcesses/AugmentationOne.cs

@@ -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();

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

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