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

using System.Text;

using AIStudio.Agents;
using AIStudio.Chat;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.Security;

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.");
    
    /// <inheritdoc />
    public async Task<AugmentationResult> ProcessAsync(IProvider provider, IContent lastUserPrompt, ChatThread chatThread, IReadOnlyList<IRetrievalContext> retrievalContexts, CancellationToken token = default)
    {
        var settings = Program.SERVICE_PROVIDER.GetService<SettingsManager>()!;
        
        if(retrievalContexts.Count == 0)
        {
            LOGGER.LogWarning("No retrieval contexts were issued. Skipping the augmentation process.");
            return new(chatThread, []);
        }

        //
        // 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.
        //
        var guardService = Program.SERVICE_PROVIDER.GetRequiredService<PromptInjectionGuardService>();
        await using var reportingScope = guardService.BeginAction();
        var preparedContexts = await retrievalContexts.PrepareForModelAsync(guardService, reportingScope, token);
        var usedContexts = preparedContexts;

        // Want the user to validate all retrieval contexts?
        if (settings.ConfigurationData.AgentRetrievalContextValidation.EnableRetrievalContextValidation && chatThread.DataSourceOptions.AutomaticValidation)
        {
            // Let's get the validation agent & set up its provider:
            var validationAgent = Program.SERVICE_PROVIDER.GetService<AgentRetrievalContextValidation>()!;
            if (validationAgent.SetLLMProvider(provider, chatThread.DataSecurity, chatThread.RequiredProviderConfidence))
            {
                try
                {
                    // Let's validate all retrieval contexts:
                    var validationResults = await validationAgent.ValidateRetrievalContextsAsync(lastUserPrompt, chatThread, preparedContexts, token);
                    if (validationResults.Count == 0)
                        LOGGER.LogWarning("Retrieval context validation returned no results. Continuing augmentation with all retrieved contexts.");
                    else
                    {
                        //
                        // Now, filter the retrieval contexts to the most relevant ones:
                        //
                        var targetWindow = validationResults.DetermineTargetWindow(TargetWindowStrategy.TOP10_BETTER_THAN_GUESSING);
                        var threshold = validationResults.GetConfidenceThreshold(targetWindow);

                        // Filter the retrieval contexts:
                        usedContexts = validationResults.Where(x => x.RetrievalContext is not null && x.Confidence >= threshold).Select(x => x.RetrievalContext!).ToList();
                    }
                }
                catch (OperationCanceledException) when (token.IsCancellationRequested)
                {
                    throw;
                }
                catch (Exception exception)
                {
                    LOGGER.LogError(exception, "Retrieval context validation failed. Continuing augmentation with all retrieved contexts.");

                    //
                    // The user switched this check on. Continuing without it silently would hide
                    // that the answer rests on unfiltered passages:
                    //
                    await MessageBus.INSTANCE.SendWarning(new(Icons.Material.Filled.FactCheck, TB("The check of which passages fit your question failed. This answer uses all passages that were found.")));
                }
            }
            else
            {
                //
                // No message to the user here: which providers are trusted enough is a setting, not
                // an event. It does not change between two answers, so a message would repeat itself
                // with every single one until the setting changes.
                //
                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:
        //
        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);
    }

    #endregion
}