AI-Studio/app/MindWork AI Studio/Components/DataSourceSelection.razor.cs

using AIStudio.Dialogs.Settings;
using AIStudio.Provider;
using AIStudio.Settings;
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Services;
using AIStudio.Tools.ToolCallingSystem;

using Microsoft.AspNetCore.Components;

using DialogOptions = AIStudio.Dialogs.DialogOptions;

namespace AIStudio.Components;

public partial class DataSourceSelection : MSGComponentBase
{
    [Parameter]
    public DataSourceSelectionMode SelectionMode { get; set; } = DataSourceSelectionMode.SELECTION_MODE;
    
    [Parameter]
    public string PopoverButtonClasses { get; set; } = string.Empty;
    
    [Parameter]
    public required AIStudio.Settings.Provider LLMProvider { get; set; }

    [Parameter]
    public required DataSourceOptions DataSourceOptions { get; set; }
    
    [Parameter]
    public EventCallback<DataSourceOptions> DataSourceOptionsChanged { get; set; }

    [Parameter]
    public IReadOnlyList<DataSourceAgentSelected> DataSourcesAISelected { get; set; } = [];
    
    [Parameter]
    public string ConfigurationHeaderMessage { get; set; } = string.Empty;
    
    [Parameter]
    public bool AutoSaveAppSettings { get; set; }

    /// <summary>
    /// Shows the options without letting the user change them.
    /// </summary>
    /// <remarks>
    /// For options somebody else decided on, such as those of a chat template an organization
    /// rolled out. Seeing which data such a chat will search is the point; changing it here is not.<br/><br/>
    ///
    /// Also for options which must not change right now, such as those of a chat while an answer is
    /// being written for it: the answer reads them, and they are the very options shown here, so a
    /// change would reach it halfway. Every change is turned down before it touches the options.
    /// </remarks>
    [Parameter]
    public bool ReadOnly { get; set; }

    /// <summary>
    /// Tells why the options cannot be changed, shown at the top of the popover while they are read-only.
    /// </summary>
    [Parameter]
    public string ReadOnlyMessage { get; set; } = string.Empty;

    /// <summary>
    /// Whether the options edited here are the data source defaults of the chat.
    /// </summary>
    /// <remarks>
    /// Those defaults can be locked by a configuration plugin, and this component reads the locks
    /// from the chat settings. Wherever the same options belong to something else — to a chat
    /// template, say — the locks of the chat defaults have nothing to say about them.
    /// </remarks>
    [Parameter]
    public bool ConfiguresChatDefaults { get; set; } = true;

    /// <summary>
    /// Where the chats run which search these data sources. It decides, among other things, whether the models may
    /// search the data sources themselves, see ToolRegistry.GetEffectiveRetrievalModeAsync.
    /// </summary>
    [Parameter]
    public Tools.Components Component { get; set; } = Tools.Components.CHAT;

    /// <summary>
    /// Further models which chat with the same data source options, e.g., the second model of a comparison series.
    /// </summary>
    /// <remarks>
    /// The selection then offers only the data sources each model may use, and Semantic Search only when each model
    /// can search itself, see SharedSelectionRules. Without further models, it offers what LLMProvider alone may use.
    /// </remarks>
    [Parameter]
    public IReadOnlyList<AIStudio.Settings.Provider> AdditionalProviders { get; set; } = [];

    /// <summary>
    /// Further models which read what the data sources gave without searching them, e.g., the judge of a comparison
    /// series, which reads both chats. Each of them must be allowed every data source the selection offers, but they
    /// have no say in how the data sources are searched.
    /// </summary>
    [Parameter]
    public IReadOnlyList<AIStudio.Settings.Provider> AdditionalReaders { get; set; } = [];

    /// <summary>
    /// Whether the user may let an AI select the data sources and validate what they gave.
    /// </summary>
    /// <remarks>
    /// Without, the user selects the data sources by hand, and the selection neither shows the switches for the AI nor
    /// listens for the data sources an AI selected elsewhere: e.g., a comparison series which gives both models the
    /// same data sources. Options which switch the AI on count as switched off then.
    /// </remarks>
    [Parameter]
    public bool AllowsAutomaticSelection { get; set; } = true;

    [Inject]
    private DataSourceService DataSourceService { get; init; } = null!;

    [Inject]
    private ToolRegistry ToolRegistry { get; init; } = null!;

    [Inject]
    private IDialogService DialogService { get; init; } = null!;

    private bool internalChange;
    private bool showDataSourceSelection;
    private bool waitingForDataSources = true;
    private IReadOnlyList<IDataSource> availableDataSources = [];
    private IReadOnlyList<IDataSource> dataSourcesAwaitingReindex = [];
    private HashSet<string> dataSourceIdsAwaitingReindex = new(StringComparer.Ordinal);
    private IReadOnlyList<IDataSource> dataSourcesNeedingRepair = [];
    private HashSet<string> dataSourceIdsNeedingRepair = new(StringComparer.Ordinal);
    private IReadOnlyCollection<IDataSource> selectedDataSources = [];
    private bool aiBasedSourceSelection;
    private bool aiBasedValidation;
    private bool areDataSourcesEnabled;
    private DataSourceRetrievalMode retrievalMode;
    private EffectiveRetrievalMode effectiveRetrievalMode;
    private uint loadAndApplyFiltersGeneration;
    
    #region Overrides of ComponentBase

    protected override async Task OnInitializedAsync()
    {
        this.ApplyFilters([], this.AllowsAutomaticSelection ? [ Event.RAG_AUTO_DATA_SOURCES_SELECTED ] : []);

        //
        // Load the settings:
        //
        this.aiBasedSourceSelection = this.IsAutomaticSelectionOn;
        this.aiBasedValidation = this.IsAutomaticValidationOn;
        this.areDataSourcesEnabled = !this.DataSourceOptions.DisableDataSources;
        this.retrievalMode = this.DataSourceOptions.RetrievalMode;
        this.waitingForDataSources = this.areDataSourcesEnabled && this.SelectionMode is not DataSourceSelectionMode.CONFIGURATION_MODE;

        //
        // Preselect the data sources. Right now, we cannot filter
        // the data sources. Later, when the component is shown, we
        // will filter the data sources.
        //
        // Right before the preselection would be used to kick off the
        // RAG process, we will filter the data sources as well.
        //
        this.selectedDataSources = this.GetDataSourcesFromConfiguredIds();
        await base.OnInitializedAsync();
    }

    protected override async Task OnParametersSetAsync()
    {
        if (!this.internalChange)
        {
            this.aiBasedSourceSelection = this.IsAutomaticSelectionOn;
            this.aiBasedValidation = this.IsAutomaticValidationOn;
            this.areDataSourcesEnabled = !this.DataSourceOptions.DisableDataSources;
            this.retrievalMode = this.DataSourceOptions.RetrievalMode;
            this.selectedDataSources = this.GetDataSourcesFromConfiguredIds();
        }

        switch (this.SelectionMode)
        {
            //
            // In selection mode, we have to load & filter the data sources
            // when the component is shown:
            //
            case DataSourceSelectionMode.SELECTION_MODE:
                
                //
                // For external changes, we have to reload & filter
                // the data sources:
                //
                if (this.showDataSourceSelection && !this.internalChange)
                    await this.LoadAndApplyFilters();
                else
                    this.internalChange = false;
                
                break;
            
            //
            // In configuration mode, we have to load all data sources:
            //
            case DataSourceSelectionMode.CONFIGURATION_MODE:
                this.availableDataSources = this.GetConfiguredDataSourcesSnapshot();
                break;
        }
        
        await base.OnParametersSetAsync();
    }

    #endregion
    
    private async Task OpenSettingsDialog()
    {
        this.showDataSourceSelection = false;
        this.StateHasChanged();
        
        var dialogParameters = new DialogParameters();
        var dialogReference = await this.DialogService.ShowAsync<SettingsDialogDataSources>(null, dialogParameters, DialogOptions.FULLSCREEN);
        await dialogReference.Result;
        await this.LoadAndApplyFilters();
        
        this.showDataSourceSelection = true;
        this.StateHasChanged();
    }

    /// <summary>
    /// Whether an AI selects the data sources, see AllowsAutomaticSelection.
    /// </summary>
    private bool IsAutomaticSelectionOn => this.AllowsAutomaticSelection && this.DataSourceOptions.AutomaticDataSourceSelection;

    /// <summary>
    /// Whether an AI validates what the data sources gave, see AllowsAutomaticSelection.
    /// </summary>
    private bool IsAutomaticValidationOn => this.AllowsAutomaticSelection && this.DataSourceOptions.AutomaticValidation;

    /// <summary>
    /// The models which chat with these data sources.
    /// </summary>
    private IReadOnlyList<AIStudio.Settings.Provider> ChattingProviders => [this.LLMProvider, .. this.AdditionalProviders];

    private SelectionMode GetListSelectionMode() => this.aiBasedSourceSelection ? MudBlazor.SelectionMode.SingleSelection : MudBlazor.SelectionMode.MultiSelection;
    
    private IReadOnlyCollection<DataSourceAgentSelected> GetSelectedDataSourcesWithAI() => this.DataSourcesAISelected.Where(n => n.Selected).ToList();

    private string GetAIReasoning(DataSourceAgentSelected source) => $"AI reasoning (confidence {source.AIDecision.Confidence:P0}): {source.AIDecision.Reason}";

    private string GetConfidenceIconStyle(IInternalDataSource source) => $"{source.ConfidenceLevel.SetColorStyle(this.SettingsManager)} flex-shrink: 0;";
    
    public void ChangeOptionWithoutSaving(DataSourceOptions options, IReadOnlyList<DataSourceAgentSelected>? aiSelectedDataSources = null)
    {
        this.DataSourceOptions = options;
        this.DataSourcesAISelected = aiSelectedDataSources ?? [];

        this.aiBasedSourceSelection = this.IsAutomaticSelectionOn;
        this.aiBasedValidation = this.IsAutomaticValidationOn;
        this.areDataSourcesEnabled = !this.DataSourceOptions.DisableDataSources;
        this.retrievalMode = this.DataSourceOptions.RetrievalMode;
        this.selectedDataSources = this.GetDataSourcesFromConfiguredIds();
        this.waitingForDataSources = false;

        //
        // Remark: We do not apply the filters here. This is done later
        // when either the parameters are changed or just before the
        // RAG process is started (outside of this component).
        //
        // In fact, when we apply the filters here, multiple calls
        // to the filter method would be made. We would get conflicts.
        //
    }
    
    public bool IsVisible => this.showDataSourceSelection;

    public void Hide()
    {
        this.showDataSourceSelection = false;
        this.StateHasChanged();
    }

    private IReadOnlyList<IDataSource> GetConfiguredDataSourcesSnapshot() => this.SettingsManager.ConfigurationData.DataSources.ToList();

    private IReadOnlyCollection<IDataSource> GetDataSourcesFromConfiguredIds()
    {
        var preselectedDataSourceIds = this.DataSourceOptions.PreselectedDataSourceIds.ToHashSet(StringComparer.Ordinal);
        return this.GetConfiguredDataSourcesSnapshot().Where(ds => preselectedDataSourceIds.Contains(ds.Id)).ToList();
    }

    /// <summary>
    /// Collects the preselected data sources which the filters removed.
    /// </summary>
    /// <remarks>
    /// The list of available sources shows what survived the filters, while the preselection keeps
    /// what the user asked for. Without this, a preselected source which cannot be used right now
    /// is simply missing from that list, and nothing says so. Preselected ids without a configured
    /// source are left out: that source is gone, not unavailable.
    /// </remarks>
    /// <returns>The unusable preselected data sources, or an empty list when there are none.</returns>
    private IReadOnlyList<IDataSource> GetUnavailablePreselectedDataSources()
    {
        var availableDataSourceIds = this.availableDataSources.Select(ds => ds.Id).ToHashSet(StringComparer.Ordinal);
        return this.GetDataSourcesFromConfiguredIds().Where(ds => !availableDataSourceIds.Contains(ds.Id)).ToList();
    }
    
    private async Task LoadAndApplyFilters()
    {
        if(this.DataSourceOptions.DisableDataSources)
        {
            this.loadAndApplyFiltersGeneration++;
            return;
        }

        if(this.SelectionMode is DataSourceSelectionMode.CONFIGURATION_MODE)
        {
            this.loadAndApplyFiltersGeneration++;
            return;
        }

        var generation = ++this.loadAndApplyFiltersGeneration;
        this.waitingForDataSources = true;
        this.StateHasChanged();
            
        //
        // Load the data sources. We ask with the preselection rather than with the field below:
        // that field holds what was usable the last time we looked, so a source filtered out once
        // would never come back, while the RAG process keeps reading it from the preselection.
        //
        // How the chat searches decides whether the agents of the RAG process see the data, and
        // with that which data sources the provider of the chat may use at all. With several
        // models, all of them search the same way, and each must be allowed every data source:
        //
        var preselectedDataSources = this.GetDataSourcesFromConfiguredIds();
        List<EffectiveRetrievalMode> modes = [];
        foreach (var provider in this.ChattingProviders)
            modes.Add(await this.ToolRegistry.GetEffectiveRetrievalModeAsync(this.DataSourceOptions, provider, this.Component));

        var effectiveMode = SharedSelectionRules.RetrievalModeOfAll(modes);
        List<AllowedSelectedDataSources> results = [];
        foreach (var provider in this.ChattingProviders.Concat(this.AdditionalReaders))
            results.Add(await this.DataSourceService.GetDataSources(provider, this.DataSourceOptions, effectiveMode.Mode, preselectedDataSources));

        if (generation != this.loadAndApplyFiltersGeneration)
            return;

        var sources = SharedSelectionRules.DataSourcesOfAll(results);

        this.effectiveRetrievalMode = effectiveMode;
        this.availableDataSources = sources.AllowedDataSources;
        this.dataSourcesAwaitingReindex = sources.DataSourcesAwaitingReindex;
        this.dataSourceIdsAwaitingReindex = sources.DataSourcesAwaitingReindex.Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
        this.dataSourcesNeedingRepair = sources.DataSourcesNeedingRepair;
        this.dataSourceIdsNeedingRepair = sources.DataSourcesNeedingRepair.Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
        this.selectedDataSources = sources.SelectedDataSources;
        this.waitingForDataSources = false;
        this.StateHasChanged();
    }

    /// <summary>
    /// Why a data source is listed but cannot be picked, if it cannot.
    /// </summary>
    /// <remarks>
    /// The repair is asked about first. The service hands a data source to one of the two lists
    /// only, but should that ever change, the reason the user can act on is the one worth showing.
    /// </remarks>
    private DataSourceBlockReason GetBlockReason(IDataSource dataSource)
    {
        if (this.dataSourceIdsNeedingRepair.Contains(dataSource.Id))
            return DataSourceBlockReason.NEEDS_REPAIR;

        if (this.dataSourceIdsAwaitingReindex.Contains(dataSource.Id))
            return DataSourceBlockReason.AWAITING_REINDEX;

        return DataSourceBlockReason.NONE;
    }

    /// <summary>
    /// The data sources the list shows: the usable ones, plus the ones which cannot be searched.
    /// </summary>
    /// <remarks>
    /// Kept in the order the data sources were configured in, rather than usable ones first. A row
    /// which jumps to another place the moment its data source starts being re-indexed is a row the
    /// user has to find again.
    /// </remarks>
    private IReadOnlyList<IDataSource> GetListedDataSources()
    {
        if (this.dataSourcesAwaitingReindex.Count == 0 && this.dataSourcesNeedingRepair.Count == 0)
            return this.availableDataSources;

        var listedIds = this.availableDataSources.Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
        listedIds.UnionWith(this.dataSourceIdsAwaitingReindex);
        listedIds.UnionWith(this.dataSourceIdsNeedingRepair);
        return this.GetConfiguredDataSourcesSnapshot().Where(source => listedIds.Contains(source.Id)).ToList();
    }

    /// <summary>
    /// The preselected but unusable data sources the warning box lists.
    /// </summary>
    /// <remarks>
    /// The ones which are only blocked are left out: they have a row of their own in the list
    /// above, which says the same thing in the place the user is already looking.
    /// </remarks>
    private IReadOnlyList<IDataSource> GetUnavailablePreselectedDataSourcesToList() =>
        this.GetUnavailablePreselectedDataSources().Where(source => this.GetBlockReason(source) is DataSourceBlockReason.NONE).ToList();
    
    private async Task EnabledChanged(bool state)
    {
        if (this.ReadOnly)
            return;

        this.areDataSourcesEnabled = state;
        this.DataSourceOptions.DisableDataSources = !this.areDataSourcesEnabled;
        
        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
        this.StateHasChanged();
    }
    
    private async Task AutoModeChanged(bool state)
    {
        if (this.ReadOnly)
            return;

        this.aiBasedSourceSelection = state;
        this.DataSourceOptions.AutomaticDataSourceSelection = this.aiBasedSourceSelection;

        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
    }
    
    private async Task SemanticSearchChanged(bool state)
    {
        if (this.ReadOnly)
            return;

        this.retrievalMode = state ? DataSourceRetrievalMode.SEMANTIC_SEARCH : DataSourceRetrievalMode.EVERY_MESSAGE;
        this.DataSourceOptions.RetrievalMode = this.retrievalMode;

        // Which data sources the provider may use depends on it, see LoadAndApplyFilters:
        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
    }

    private bool IsSemanticSearchPreferred => this.retrievalMode is DataSourceRetrievalMode.SEMANTIC_SEARCH;

    /// <summary>
    /// Whether the model of the chat searches the data sources itself.
    /// </summary>
    /// <remarks>
    /// Only the selection mode can tell, since only it knows the provider of a chat. The effective
    /// mode is found out while loading the data sources; asking the preference as well keeps a
    /// switch to every message from showing the other way until that loading is done.
    /// </remarks>
    private bool IsSemanticSearchEffective => this.IsSemanticSearchPreferred && this.effectiveRetrievalMode.Mode is DataSourceRetrievalMode.SEMANTIC_SEARCH;

    /// <summary>
    /// Whether Semantic Search cannot be used in this chat, whatever the user prefers.
    /// </summary>
    /// <remarks>
    /// Then there is nothing to choose: classic RAG searches the data sources. The choice is left
    /// out rather than shown as one which changes nothing, and the preference stays as it is, so it
    /// is back the moment the chat gets a model which can search itself.
    /// </remarks>
    private bool IsSemanticSearchUnavailable => this.effectiveRetrievalMode.SemanticSearchBlockReason is not ToolOfferBlockReason.NONE;

    private string GetSemanticSearchUnavailableMessage() => this.effectiveRetrievalMode.SemanticSearchBlockReason switch
    {
        ToolOfferBlockReason.TOOLS_SWITCHED_OFF => TA("Your organization has switched tools off, so AI Studio uses classic RAG instead and searches your data sources with every message."),
        ToolOfferBlockReason.MODEL_CANNOT_USE_TOOLS when this.AdditionalProviders.Count > 0 => TA("One of the selected models cannot use tools, so AI Studio uses classic RAG for all of them instead and searches your data sources with every message."),
        ToolOfferBlockReason.PROVIDER_CONFIDENCE_TOO_LOW when this.AdditionalProviders.Count > 0 => TA("One of the selected providers is not trusted enough for Semantic Search, so AI Studio uses classic RAG for all of them instead and searches your data sources with every message."),
        ToolOfferBlockReason.MODEL_CANNOT_USE_TOOLS => TA("The selected model cannot use tools, so AI Studio uses classic RAG instead and searches your data sources with every message."),
        ToolOfferBlockReason.TOOL_SWITCHED_OFF => TA("Your organization has switched Semantic Search off, so AI Studio uses classic RAG instead and searches your data sources with every message."),
        ToolOfferBlockReason.TOOL_SWITCHED_OFF_BY_USER => TA("You switched Semantic Search off in the tool settings, so AI Studio uses classic RAG instead and searches your data sources with every message."),
        ToolOfferBlockReason.PROVIDER_CONFIDENCE_TOO_LOW => TA("The selected provider is not trusted enough for Semantic Search, so AI Studio uses classic RAG instead and searches your data sources with every message."),
        _ => TA("Semantic Search cannot be used here, so AI Studio uses classic RAG instead and searches your data sources with every message."),
    };

    private async Task ValidationModeChanged(bool state)
    {
        if (this.ReadOnly)
            return;

        this.aiBasedValidation = state;
        this.DataSourceOptions.AutomaticValidation = this.aiBasedValidation;

        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
    }
    
    private async Task SelectionChanged(IReadOnlyCollection<IDataSource>? chosenDataSources)
    {
        if (this.ReadOnly)
            return;

        this.selectedDataSources = chosenDataSources ?? [];

        //
        // The list offers only the data sources which survived the filters, so what the user picks
        // there says nothing about the preselected ones it could not show. Those are kept: dropping
        // them would undo a choice the user never revisited, and it is these ids -- not this list --
        // which the RAG process reads when an answer is created. The query has to run before the
        // assignment, because it reads what we are about to replace.
        //
        var keptDataSourceIds = this.GetUnavailablePreselectedDataSources().Select(ds => ds.Id).ToList();
        this.DataSourceOptions.PreselectedDataSourceIds = [..keptDataSourceIds, ..this.selectedDataSources.Select(ds => ds.Id)];

        await this.OptionsChanged();
    }

    private bool IsPreselectedDataSourcesDisabledLocked()
    {
        return this.SelectionMode is DataSourceSelectionMode.CONFIGURATION_MODE
               && this.ConfiguresChatDefaults
               && ManagedConfiguration.TryGet(x => x.Chat, x => x.PreselectedDataSourcesDisabled, out var meta)
               && meta.IsLocked;
    }

    private bool IsPreselectedDataSourcesRetrievalModeLocked()
    {
        return this.SelectionMode is DataSourceSelectionMode.CONFIGURATION_MODE
               && this.ConfiguresChatDefaults
               && ManagedConfiguration.TryGet(x => x.Chat, x => x.PreselectedDataSourcesRetrievalMode, out var meta)
               && meta.IsLocked;
    }

    private bool IsPreselectedDataSourcesAutomaticSelectionLocked()
    {
        return this.SelectionMode is DataSourceSelectionMode.CONFIGURATION_MODE
               && this.ConfiguresChatDefaults
               && ManagedConfiguration.TryGet(x => x.Chat, x => x.PreselectedDataSourcesAutomaticSelection, out var meta)
               && meta.IsLocked;
    }

    private bool IsPreselectedDataSourcesAutomaticValidationLocked()
    {
        return this.SelectionMode is DataSourceSelectionMode.CONFIGURATION_MODE
               && this.ConfiguresChatDefaults
               && ManagedConfiguration.TryGet(x => x.Chat, x => x.PreselectedDataSourcesAutomaticValidation, out var meta)
               && meta.IsLocked;
    }

    private bool IsPreselectedDataSourceIdsLocked()
    {
        return this.SelectionMode is DataSourceSelectionMode.CONFIGURATION_MODE
               && this.ConfiguresChatDefaults
               && ManagedConfiguration.TryGet(x => x.Chat, x => x.PreselectedDataSourceIds, out var meta)
               && meta.IsLocked;
    }

    private async Task OptionsChanged()
    {
        this.internalChange = true;
        this.loadAndApplyFiltersGeneration++;
        
        await this.DataSourceOptionsChanged.InvokeAsync(this.DataSourceOptions);
        
        if(this.AutoSaveAppSettings)
            await this.SettingsManager.StoreSettings();
    }
    
    private async Task ToggleDataSourceSelection()
    {
        this.showDataSourceSelection = !this.showDataSourceSelection;
        if (this.showDataSourceSelection)
            await this.LoadAndApplyFilters();
    }
    
    private void HideDataSourceSelection() => this.showDataSourceSelection = false;

    #region Overrides of MSGComponentBase

    protected override Task ProcessIncomingMessage<T>(ComponentBase? sendingComponent, Event triggeredEvent, T? data) where T : default
    {
        switch (triggeredEvent)
        {
            case Event.RAG_AUTO_DATA_SOURCES_SELECTED:
                if(data is IReadOnlyList<DataSourceAgentSelected> aiSelectedDataSources)
                    this.DataSourcesAISelected = aiSelectedDataSources;
                    
                this.StateHasChanged();
                break;
        }

        return Task.CompletedTask;
    }

    #endregion
}