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 PopoverTriggerMode PopoverTriggerMode { get; set; } = PopoverTriggerMode.BUTTON;
    
    [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.
    /// </remarks>
    [Parameter]
    public bool ReadOnly { get; set; }

    /// <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;

    [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([], [ Event.RAG_AUTO_DATA_SOURCES_SELECTED ]);
        
        //
        // Load the settings:
        //
        this.aiBasedSourceSelection = this.DataSourceOptions.AutomaticDataSourceSelection;
        this.aiBasedValidation = this.DataSourceOptions.AutomaticValidation;
        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.DataSourceOptions.AutomaticDataSourceSelection;
            this.aiBasedValidation = this.DataSourceOptions.AutomaticValidation;
            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();
    }

    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.DataSourceOptions.AutomaticDataSourceSelection;
        this.aiBasedValidation = this.DataSourceOptions.AutomaticValidation;
        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. This component
        // only ever selects for a chat:
        //
        var effectiveMode = await this.ToolRegistry.GetEffectiveRetrievalModeAsync(this.DataSourceOptions, this.LLMProvider, Tools.Components.CHAT);
        var sources = await this.DataSourceService.GetDataSources(this.LLMProvider, this.DataSourceOptions, effectiveMode.Mode, this.GetDataSourcesFromConfiguredIds());
        if (generation != this.loadAndApplyFiltersGeneration)
            return;

        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)
    {
        this.areDataSourcesEnabled = state;
        this.DataSourceOptions.DisableDataSources = !this.areDataSourcesEnabled;
        
        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
        this.StateHasChanged();
    }
    
    private async Task AutoModeChanged(bool state)
    {
        this.aiBasedSourceSelection = state;
        this.DataSourceOptions.AutomaticDataSourceSelection = this.aiBasedSourceSelection;

        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
    }
    
    private async Task SemanticSearchChanged(bool state)
    {
        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 => T("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 => T("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 => T("Your organization has switched Semantic Search off, so AI Studio uses classic RAG instead and searches your data sources with every message."),
        ToolOfferBlockReason.PROVIDER_CONFIDENCE_TOO_LOW => T("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."),
        _ => T("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)
    {
        this.aiBasedValidation = state;
        this.DataSourceOptions.AutomaticValidation = this.aiBasedValidation;

        await this.LoadAndApplyFilters();
        await this.OptionsChanged();
    }
    
    private async Task SelectionChanged(IReadOnlyCollection<IDataSource>? chosenDataSources)
    {
        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
}