mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-09-27 03:53:36 +00:00
Keep data sources selectable only when their index can answer
This commit is contained in:
parent
a8fded4df5
commit
843cca471b
@ -1,6 +1,37 @@
|
||||
@using AIStudio.Settings
|
||||
@using AIStudio.Provider
|
||||
@inherits MSGComponentBase
|
||||
@*
|
||||
One row of the data source lists below. A data source waiting for its index stays in the list
|
||||
instead of disappearing from it, but cannot be picked, and the tooltip says why. The tool
|
||||
selection next to this one in the chat answers the same question the same way.
|
||||
|
||||
The tooltip sits around the list item rather than inside it: a disabled item has its pointer
|
||||
events switched off and would swallow the hover.
|
||||
*@
|
||||
@{
|
||||
RenderFragment<IDataSource> dataSourceRow = source =>
|
||||
@<MudTooltip Text="@T("This data source is waiting to be indexed again. Until that is finished, it cannot be searched.")" Disabled="@(!this.IsAwaitingReindex(source))" RootStyle="display: block;" Placement="Placement.Top">
|
||||
<MudListItem Value="@source" Disabled="@this.IsAwaitingReindex(source)">
|
||||
<MudStack Row="true" AlignItems="AlignItems.Center" Spacing="1" Style="min-width: 0; width: 100%;">
|
||||
<MudText Typo="Typo.body2" Style="min-width: 0; white-space: normal; overflow-wrap: anywhere;">
|
||||
@source.Name
|
||||
</MudText>
|
||||
@if (source is IInternalDataSource internalSource)
|
||||
{
|
||||
<MudSpacer/>
|
||||
@if (this.IsAwaitingReindex(source))
|
||||
{
|
||||
<MudIcon Icon="@Icons.Material.Filled.HourglassTop" Size="Size.Small" Color="Color.Warning" Style="flex-shrink: 0;"/>
|
||||
}
|
||||
<MudTooltip Text="@internalSource.ConfidenceLevel.GetName()">
|
||||
<MudIcon Icon="@Icons.Material.Filled.Security" Size="Size.Small" Class="confidence-icon" Style="@this.GetConfidenceIconStyle(internalSource)"/>
|
||||
</MudTooltip>
|
||||
}
|
||||
</MudStack>
|
||||
</MudListItem>
|
||||
</MudTooltip>;
|
||||
}
|
||||
@if (this.SelectionMode is DataSourceSelectionMode.SELECTION_MODE)
|
||||
{
|
||||
<div class="d-flex">
|
||||
@ -69,7 +100,7 @@
|
||||
|
||||
@switch (this.aiBasedSourceSelection)
|
||||
{
|
||||
case true when this.availableDataSources.Count == 0:
|
||||
case true when this.GetListedDataSources().Count == 0:
|
||||
<MudText Typo="Typo.body2" Class="mb-2">
|
||||
@T("Your data sources cannot be used with the selected providers due to data privacy or confidence-level requirements, or they are currently unavailable.")
|
||||
</MudText>
|
||||
@ -81,7 +112,7 @@
|
||||
</MudText>
|
||||
break;
|
||||
|
||||
case false when this.availableDataSources.Count == 0:
|
||||
case false when this.GetListedDataSources().Count == 0:
|
||||
<MudText Typo="Typo.body2" Class="mb-2">
|
||||
@T("Your data sources cannot be used with the selected providers due to data privacy or confidence-level requirements, or they are currently unavailable.")
|
||||
</MudText>
|
||||
@ -90,22 +121,9 @@
|
||||
case false:
|
||||
<MudField Label="@T("Available Data Sources")" Variant="Variant.Outlined" Class="mb-2" Disabled="@this.aiBasedSourceSelection">
|
||||
<MudList T="IDataSource" Dense="@true" Class="data-source-rows" SelectionMode="@this.GetListSelectionMode()" @bind-SelectedValues:get="@this.selectedDataSources" @bind-SelectedValues:set="@this.SelectionChanged" Style="max-height: 14em;">
|
||||
@foreach (var source in this.availableDataSources)
|
||||
@foreach (var source in this.GetListedDataSources())
|
||||
{
|
||||
<MudListItem Value="@source">
|
||||
<MudStack Row="true" AlignItems="AlignItems.Center" Spacing="1" Style="min-width: 0; width: 100%;">
|
||||
<MudText Typo="Typo.body2" Style="min-width: 0; white-space: normal; overflow-wrap: anywhere;">
|
||||
@source.Name
|
||||
</MudText>
|
||||
@if (source is IInternalDataSource internalSource)
|
||||
{
|
||||
<MudSpacer/>
|
||||
<MudTooltip Text="@internalSource.ConfidenceLevel.GetName()">
|
||||
<MudIcon Icon="@Icons.Material.Filled.Security" Size="Size.Small" Class="confidence-icon" Style="@this.GetConfidenceIconStyle(internalSource)"/>
|
||||
</MudTooltip>
|
||||
}
|
||||
</MudStack>
|
||||
</MudListItem>
|
||||
@dataSourceRow(source)
|
||||
}
|
||||
</MudList>
|
||||
</MudField>
|
||||
@ -115,22 +133,9 @@
|
||||
<MudExpansionPanels MultiExpansion="@false" Class="mt-3" Style="max-height: 14em;">
|
||||
<ExpansionPanel HeaderIcon="@Icons.Material.Filled.TouchApp" IconSize="Size.Small" HeaderTypo="Typo.subtitle1" HeaderClass="expansion-panel-header-compact" HeaderText="@T("Available Data Sources")">
|
||||
<MudList T="IDataSource" Dense="@true" Class="data-source-rows" SelectionMode="MudBlazor.SelectionMode.SingleSelection" SelectedValues="@this.selectedDataSources" Style="max-height: 14em;">
|
||||
@foreach (var source in this.availableDataSources)
|
||||
@foreach (var source in this.GetListedDataSources())
|
||||
{
|
||||
<MudListItem Value="@source">
|
||||
<MudStack Row="true" AlignItems="AlignItems.Center" Spacing="1" Style="min-width: 0; width: 100%;">
|
||||
<MudText Typo="Typo.body2" Style="min-width: 0; white-space: normal; overflow-wrap: anywhere;">
|
||||
@source.Name
|
||||
</MudText>
|
||||
@if (source is IInternalDataSource internalSource)
|
||||
{
|
||||
<MudSpacer/>
|
||||
<MudTooltip Text="@internalSource.ConfidenceLevel.GetName()">
|
||||
<MudIcon Icon="@Icons.Material.Filled.Security" Size="Size.Small" Class="confidence-icon" Style="@this.GetConfidenceIconStyle(internalSource)"/>
|
||||
</MudTooltip>
|
||||
}
|
||||
</MudStack>
|
||||
</MudListItem>
|
||||
@dataSourceRow(source)
|
||||
}
|
||||
</MudList>
|
||||
</ExpansionPanel>
|
||||
@ -166,13 +171,13 @@
|
||||
break;
|
||||
}
|
||||
|
||||
@if (!this.aiBasedSourceSelection && this.GetUnavailablePreselectedDataSources().Count > 0)
|
||||
@if (!this.aiBasedSourceSelection && this.GetUnavailablePreselectedDataSourcesToList().Count > 0)
|
||||
{
|
||||
<MudJustifiedText Typo="Typo.body2" Color="Color.Warning">
|
||||
@T("These data sources are preselected, but cannot be used right now, either due to data privacy or confidence-level requirements, or because they are unavailable:")
|
||||
</MudJustifiedText>
|
||||
<ul class="unavailable-data-sources mb-3 mt-1">
|
||||
@foreach (var source in this.GetUnavailablePreselectedDataSources())
|
||||
@foreach (var source in this.GetUnavailablePreselectedDataSourcesToList())
|
||||
{
|
||||
<li>@source.Name</li>
|
||||
}
|
||||
|
||||
@ -49,6 +49,8 @@ public partial class DataSourceSelection : MSGComponentBase
|
||||
private bool showDataSourceSelection;
|
||||
private bool waitingForDataSources = true;
|
||||
private IReadOnlyList<IDataSource> availableDataSources = [];
|
||||
private IReadOnlyList<IDataSource> dataSourcesAwaitingReindex = [];
|
||||
private HashSet<string> dataSourceIdsAwaitingReindex = new(StringComparer.Ordinal);
|
||||
private IReadOnlyCollection<IDataSource> selectedDataSources = [];
|
||||
private bool aiBasedSourceSelection;
|
||||
private bool aiBasedValidation;
|
||||
@ -226,11 +228,43 @@ public partial class DataSourceSelection : MSGComponentBase
|
||||
return;
|
||||
|
||||
this.availableDataSources = sources.AllowedDataSources;
|
||||
this.dataSourcesAwaitingReindex = sources.DataSourcesAwaitingReindex;
|
||||
this.dataSourceIdsAwaitingReindex = sources.DataSourcesAwaitingReindex.Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
|
||||
this.selectedDataSources = sources.SelectedDataSources;
|
||||
this.waitingForDataSources = false;
|
||||
this.StateHasChanged();
|
||||
}
|
||||
|
||||
private bool IsAwaitingReindex(IDataSource dataSource) => this.dataSourceIdsAwaitingReindex.Contains(dataSource.Id);
|
||||
|
||||
/// <summary>
|
||||
/// The data sources the list shows: the usable ones, plus the ones waiting for their index.
|
||||
/// </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)
|
||||
return this.availableDataSources;
|
||||
|
||||
var listedIds = this.availableDataSources.Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
|
||||
listedIds.UnionWith(this.dataSourceIdsAwaitingReindex);
|
||||
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 waiting for their index 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.IsAwaitingReindex(source)).ToList();
|
||||
|
||||
private async Task EnabledChanged(bool state)
|
||||
{
|
||||
this.areDataSourcesEnabled = state;
|
||||
|
||||
@ -3,11 +3,18 @@ using AIStudio.Settings;
|
||||
namespace AIStudio.Tools;
|
||||
|
||||
/// <summary>
|
||||
/// Contains both the allowed and selected data sources.
|
||||
/// Contains the allowed and selected data sources, plus the ones waiting for their index.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The selected data sources are a subset of the allowed data sources.
|
||||
///
|
||||
/// The data sources waiting for a re-index are deliberately kept apart from the allowed ones rather
|
||||
/// than mixed in. Everything reading the allowed list -- the data source selection agent above all
|
||||
/// -- takes it to mean "may be used to answer with", and a source whose index is being rebuilt
|
||||
/// cannot answer anything. It is listed separately so the user interface can still show it and say
|
||||
/// why it is greyed out, instead of letting it vanish without a word.
|
||||
/// </remarks>
|
||||
/// <param name="AllowedDataSources">The allowed data sources.</param>
|
||||
/// <param name="SelectedDataSources">The selected data sources, which are a subset of the allowed data sources.</param>
|
||||
public readonly record struct AllowedSelectedDataSources(IReadOnlyList<IDataSource> AllowedDataSources, IReadOnlyList<IDataSource> SelectedDataSources);
|
||||
/// <param name="DataSourcesAwaitingReindex">The data sources which passed every check but cannot be searched until their index has been rebuilt.</param>
|
||||
public readonly record struct AllowedSelectedDataSources(IReadOnlyList<IDataSource> AllowedDataSources, IReadOnlyList<IDataSource> SelectedDataSources, IReadOnlyList<IDataSource> DataSourcesAwaitingReindex);
|
||||
@ -0,0 +1,19 @@
|
||||
namespace AIStudio.Tools.Databases.IndexStore;
|
||||
|
||||
/// <summary>
|
||||
/// What the index knows about a data source as a whole, without its files.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The manifest answers the same question, but reads every file and every stored failure of the
|
||||
/// data source to do so. That is the right thing before a run, and far too much for a question
|
||||
/// asked about several data sources every time somebody opens the data source selection.
|
||||
///
|
||||
/// SourceHash is the telling one: it is written once a run has worked through the whole data
|
||||
/// source, and resetting the index deletes the row it lives in. So an empty hash means no run has
|
||||
/// finished since the index was last discarded.
|
||||
/// </remarks>
|
||||
/// <param name="EmbeddingProviderId">The embedding provider the stored vectors were created with.</param>
|
||||
/// <param name="EmbeddingSignature">Identifies the embedding configuration the stored vectors belong to.</param>
|
||||
/// <param name="SourceHash">The hash of the data source as a whole, written when a run completes.</param>
|
||||
/// <param name="VectorSize">The dimension of the stored vectors.</param>
|
||||
public sealed record DataSourceIndexState(string EmbeddingProviderId, string EmbeddingSignature, string SourceHash, int VectorSize);
|
||||
@ -6,6 +6,14 @@ public abstract class IndexStoreClient(string name, string path) : DatabaseClien
|
||||
{
|
||||
public abstract Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token);
|
||||
|
||||
/// <summary>
|
||||
/// Reads what the index knows about a data source as a whole, without its files.
|
||||
/// </summary>
|
||||
/// <param name="dataSourceId">The data source to read.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>The stored state, or null when the index holds nothing about this data source.</returns>
|
||||
public abstract Task<DataSourceIndexState?> GetDataSourceStateAsync(string dataSourceId, CancellationToken token);
|
||||
|
||||
public abstract Task UpsertDataSourceAsync(
|
||||
string dataSourceId,
|
||||
string dataSourceType,
|
||||
|
||||
@ -23,8 +23,9 @@ public sealed class NoIndexStoreClient(string name, string? unavailableReason, D
|
||||
await Task.CompletedTask;
|
||||
}
|
||||
|
||||
public override Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token) =>
|
||||
Task.FromResult(new DataSourceEmbeddingManifest());
|
||||
public override Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token) => Task.FromResult(new DataSourceEmbeddingManifest());
|
||||
|
||||
public override Task<DataSourceIndexState?> GetDataSourceStateAsync(string dataSourceId, CancellationToken token) => Task.FromResult<DataSourceIndexState?>(null);
|
||||
|
||||
public override Task UpsertDataSourceAsync(
|
||||
string dataSourceId,
|
||||
|
||||
@ -69,6 +69,16 @@ public sealed class SqliteIndexStoreClientImplementation(string name, string dat
|
||||
yield return (TB("Permanently skipped files"), (await context.PermanentIndexingFailures.CountAsync(CancellationToken.None)).ToString(CultureInfo.InvariantCulture));
|
||||
}
|
||||
|
||||
public override async Task<DataSourceIndexState?> GetDataSourceStateAsync(string dataSourceId, CancellationToken token)
|
||||
{
|
||||
await using var context = this.CreateContext();
|
||||
return await context.DataSources
|
||||
.AsNoTracking()
|
||||
.Where(source => source.DataSourceId == dataSourceId)
|
||||
.Select(source => new DataSourceIndexState(source.EmbeddingProviderId, source.EmbeddingSignature, source.SourceHash, source.VectorSize))
|
||||
.FirstOrDefaultAsync(token);
|
||||
}
|
||||
|
||||
public override async Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token)
|
||||
{
|
||||
await using var context = this.CreateContext();
|
||||
|
||||
@ -23,6 +23,15 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
/// </summary>
|
||||
private static readonly TimeSpan BLOCK_PROGRESS_INTERVAL = TimeSpan.FromSeconds(3);
|
||||
|
||||
/// <summary>
|
||||
/// How long the re-index check waits for the index database before it gives up.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Asked while somebody waits for the data source selection to open, and possibly while a run
|
||||
/// writes to the same database.
|
||||
/// </remarks>
|
||||
private static readonly TimeSpan REINDEX_CHECK_TIMEOUT = TimeSpan.FromSeconds(2);
|
||||
|
||||
private readonly Channel<DataSourceEmbeddingQueueItem> queue = Channel.CreateUnbounded<DataSourceEmbeddingQueueItem>();
|
||||
private readonly ConcurrentDictionary<string, byte> queuedIds = new(StringComparer.OrdinalIgnoreCase);
|
||||
private readonly ConcurrentDictionary<string, byte> runningIds = new(StringComparer.OrdinalIgnoreCase);
|
||||
@ -211,6 +220,103 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
|
||||
|| manifest.PermanentFailures.Count > 0;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Whether a data source cannot answer a search right now because its index has to be built anew.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Says nothing about a data source which is only catching up with a handful of changed files:
|
||||
/// everything indexed before is still there and still searchable. What this catches is the case
|
||||
/// where the whole index was thrown away, or is about to be, because the embedding configuration
|
||||
/// changed under it. Between discarding the old vectors and finishing the new ones, the data
|
||||
/// source looks perfectly fine and finds nothing.
|
||||
///
|
||||
/// Two things are asked, in this order. The stored signature tells whether the vectors still
|
||||
/// belong to the current configuration; it is written back right after the reset, so on its own
|
||||
/// it would call a rebuild in progress finished. The stored hash of the data source closes that
|
||||
/// gap: it survives an ordinary run but not a reset, so an empty one means no run has completed
|
||||
/// since the index was discarded.
|
||||
///
|
||||
/// Anything unclear counts as not waiting. Whoever asks does so to grey out a row, and a data
|
||||
/// source wrongly greyed out for good is worse than one which turns out to have nothing to say.
|
||||
/// </remarks>
|
||||
/// <param name="dataSource">The data source to ask about.</param>
|
||||
/// <param name="token">The cancellation token.</param>
|
||||
/// <returns>True when the data source is waiting for its index to be rebuilt.</returns>
|
||||
public async Task<bool> IsAwaitingReindexAsync(IDataSource dataSource, CancellationToken token = default)
|
||||
{
|
||||
//
|
||||
// This guard also keeps the index database out of the picture while local RAG is switched
|
||||
// off: asking for the store creates the database and runs its migrations on the first call,
|
||||
// which must not happen because somebody opened the data source selection.
|
||||
//
|
||||
if (!this.IsSupportedInternalDataSource(dataSource))
|
||||
return false;
|
||||
|
||||
if (!this.TryResolveEmbeddingProvider(dataSource, out var embeddingProvider))
|
||||
return false;
|
||||
|
||||
try
|
||||
{
|
||||
//
|
||||
// A timeout of its own: this runs while the user waits for a popover to open, and the
|
||||
// embedding service may be writing to the same database at the time.
|
||||
//
|
||||
using var timeout = CancellationTokenSource.CreateLinkedTokenSource(token);
|
||||
timeout.CancelAfter(REINDEX_CHECK_TIMEOUT);
|
||||
|
||||
var indexStore = await databaseClientProvider.GetIndexStoreAsync(timeout.Token);
|
||||
if (!indexStore.IsAvailable)
|
||||
return false;
|
||||
|
||||
var indexState = await indexStore.GetDataSourceStateAsync(dataSource.Id, timeout.Token);
|
||||
var chunkingOptions = this.GetChunkingOptions(dataSource, embeddingProvider);
|
||||
var embeddingSignature = BuildEmbeddingSignature(dataSource, embeddingProvider, chunkingOptions);
|
||||
var runState = this.statuses.TryGetValue(dataSource.Id, out var status) ? status.State : (DataSourceEmbeddingState?)null;
|
||||
|
||||
return IsIndexAwaitingRebuild(indexState, embeddingSignature, runState);
|
||||
}
|
||||
catch (Exception exception)
|
||||
{
|
||||
logger.LogWarning(exception, "Could not tell whether data source '{DataSourceName}' ({DataSourceId}) is waiting for a re-index. Treating it as usable.", dataSource.Name, dataSource.Id);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Decides from the stored index state alone whether a data source has to be indexed anew.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Kept apart from reading the database so the decision itself can be pinned down in a test.
|
||||
/// The order of the three questions is what makes it correct, see IsAwaitingReindexAsync.
|
||||
/// </remarks>
|
||||
/// <param name="indexState">What the index holds about the data source, or null when it holds nothing.</param>
|
||||
/// <param name="currentEmbeddingSignature">The signature the current embedding configuration produces.</param>
|
||||
/// <param name="runState">The state of this data source's last or current run, when one is known.</param>
|
||||
/// <returns>True when the data source is waiting for its index to be rebuilt.</returns>
|
||||
internal static bool IsIndexAwaitingRebuild(DataSourceIndexState? indexState, string currentEmbeddingSignature, DataSourceEmbeddingState? runState)
|
||||
{
|
||||
// Nothing stored at all: this data source has never been indexed, so there is nothing to
|
||||
// search in it yet.
|
||||
if (indexState is null)
|
||||
return true;
|
||||
|
||||
// The stored vectors belong to another embedding configuration. They will be thrown away
|
||||
// as soon as the next run starts, and they are of no use before that either.
|
||||
if (!string.Equals(indexState.EmbeddingSignature, currentEmbeddingSignature, StringComparison.Ordinal))
|
||||
return true;
|
||||
|
||||
// A run has worked through the whole data source since the index was last discarded.
|
||||
if (!string.IsNullOrWhiteSpace(indexState.SourceHash))
|
||||
return false;
|
||||
|
||||
//
|
||||
// The index was discarded and nothing has finished since. A failed run is the exception:
|
||||
// whatever it managed to index is searchable, and the embeddings page already names the
|
||||
// problem, so there is nothing to be gained from locking the row as well.
|
||||
//
|
||||
return runState is not DataSourceEmbeddingState.FAILED;
|
||||
}
|
||||
|
||||
public Task QueueDataSourceAsync(IDataSource dataSource)
|
||||
{
|
||||
return this.QueueDataSourceAsync(dataSource, true, DataSourceEmbeddingRefreshMode.HASH_CHECK);
|
||||
|
||||
@ -13,7 +13,7 @@ namespace AIStudio.Tools.Services;
|
||||
|
||||
public sealed class DataSourceLocalRetrievalService(
|
||||
SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider,
|
||||
ILogger<DataSourceLocalRetrievalService> logger)
|
||||
DataSourceEmbeddingService embeddingService, ILogger<DataSourceLocalRetrievalService> logger)
|
||||
{
|
||||
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceLocalRetrievalService).Namespace, nameof(DataSourceLocalRetrievalService));
|
||||
|
||||
@ -73,6 +73,23 @@ public sealed class DataSourceLocalRetrievalService(
|
||||
if (maxMatches == 0)
|
||||
return [];
|
||||
|
||||
//
|
||||
// A data source waiting for its index is kept out of the selection before the RAG process
|
||||
// starts. This catches whatever reaches retrieval another way, and turns an answer quietly
|
||||
// put together without the data into a sentence saying so.
|
||||
//
|
||||
// Asked here rather than inside one of the two channels below, because both of them read
|
||||
// what the rebuild is about to discard: with only the embedding signature changed, the old
|
||||
// chunks are still in place and the keyword search would happily answer from them while
|
||||
// the vector search finds nothing.
|
||||
//
|
||||
if (await embeddingService.IsAwaitingReindexAsync(dataSource, token))
|
||||
{
|
||||
logger.LogWarning("Skipping local retrieval for data source '{DataSourceName}' ({DataSourceId}) because its index has to be built anew.", dataSource.Name, dataSource.Id);
|
||||
await this.ReportRetrievalGapAsync(dataSource, "index-rebuilding", string.Format(TB("The data source '{0}' was left out of the answer: it is being indexed again and cannot be searched until that is finished."), dataSource.Name));
|
||||
return [];
|
||||
}
|
||||
|
||||
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
|
||||
var vectorTask = this.SearchVectorAsync(dataSource, query, maxMatches, collectionName, token);
|
||||
var bm25Task = this.SearchBm25Async(dataSource, query, maxMatches, token);
|
||||
|
||||
@ -18,15 +18,17 @@ public sealed class DataSourceService
|
||||
// ReSharper disable once NotAccessedPositionalProperty.Local
|
||||
private readonly record struct ParticipatingProvider(string Role, bool IsTrusted, ConfidenceLevel ConfidenceLevel);
|
||||
|
||||
private readonly DataSourceEmbeddingService embeddingService;
|
||||
private readonly RustService rustService;
|
||||
private readonly SettingsManager settingsManager;
|
||||
private readonly ILogger<DataSourceService> logger;
|
||||
|
||||
public DataSourceService(SettingsManager settingsManager, ILogger<DataSourceService> logger, RustService rustService)
|
||||
public DataSourceService(SettingsManager settingsManager, ILogger<DataSourceService> logger, RustService rustService, DataSourceEmbeddingService embeddingService)
|
||||
{
|
||||
this.logger = logger;
|
||||
this.rustService = rustService;
|
||||
this.settingsManager = settingsManager;
|
||||
this.embeddingService = embeddingService;
|
||||
|
||||
this.logger.LogInformation("The data source service has been initialized.");
|
||||
}
|
||||
@ -49,7 +51,7 @@ public sealed class DataSourceService
|
||||
if (selectedLLMProvider == Settings.Provider.NONE)
|
||||
{
|
||||
this.logger.LogWarning("The selected LLM provider is not set. We cannot filter the data sources by any means.");
|
||||
return new([], []);
|
||||
return new([], [], []);
|
||||
}
|
||||
|
||||
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
|
||||
@ -78,7 +80,15 @@ public sealed class DataSourceService
|
||||
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
|
||||
var participatingProviders = this.GetParticipatingProviders(selectedLLMProvider.Id, dataSourceOptions,
|
||||
new("chat provider", usingTrustedProvider, selectedLLMProvider.GetConfidenceLevel(this.settingsManager)));
|
||||
return await this.GetAllowedDataSources(usingTrustedProvider, participatingProviders, requestedDataSources);
|
||||
var allowedDataSources = await this.GetAllowedDataSources(usingTrustedProvider, participatingProviders, requestedDataSources);
|
||||
|
||||
//
|
||||
// Whoever asks this way has no list to show, so a data source waiting for its index is
|
||||
// dropped rather than marked. Handing it back would start a chat with a data source which
|
||||
// finds nothing -- the very thing being greyed out elsewhere is meant to prevent.
|
||||
//
|
||||
var awaitingReindexIds = (await this.GetDataSourcesAwaitingReindex(allowedDataSources)).Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
|
||||
return allowedDataSources.Where(source => !awaitingReindexIds.Contains(source.Id)).ToList();
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@ -99,7 +109,7 @@ public sealed class DataSourceService
|
||||
if (selectedLLMProvider is NoProvider)
|
||||
{
|
||||
this.logger.LogWarning("The selected LLM provider is the default provider. We cannot filter the data sources by any means.");
|
||||
return new([], []);
|
||||
return new([], [], []);
|
||||
}
|
||||
|
||||
var usingTrustedProvider = selectedLLMProvider.IsTrustedForDataSourceSecurityChecks(this.settingsManager);
|
||||
@ -142,9 +152,47 @@ public sealed class DataSourceService
|
||||
var allDataSources = this.settingsManager.ConfigurationData.DataSources.ToList();
|
||||
var previousSelectedDataSourceIds = previousSelectedDataSources?.Select(source => source.Id).ToHashSet(StringComparer.Ordinal) ?? [];
|
||||
var filteredDataSources = await this.GetAllowedDataSources(usingTrustedProvider, participatingProviders, allDataSources);
|
||||
var filteredSelectedDataSources = filteredDataSources.Where(source => previousSelectedDataSourceIds.Contains(source.Id)).ToList();
|
||||
|
||||
return new(filteredDataSources, filteredSelectedDataSources);
|
||||
//
|
||||
// Which of the sources that passed every check cannot answer a search right now. They are
|
||||
// held back from both lists below rather than removed altogether: a source whose index is
|
||||
// being rebuilt is usable again in a while, and saying so on its own row beats letting it
|
||||
// disappear from the selection without a word.
|
||||
//
|
||||
var awaitingReindex = await this.GetDataSourcesAwaitingReindex(filteredDataSources);
|
||||
var awaitingReindexIds = awaitingReindex.Select(source => source.Id).ToHashSet(StringComparer.Ordinal);
|
||||
var usableDataSources = filteredDataSources.Where(source => !awaitingReindexIds.Contains(source.Id)).ToList();
|
||||
var filteredSelectedDataSources = usableDataSources.Where(source => previousSelectedDataSourceIds.Contains(source.Id)).ToList();
|
||||
|
||||
return new(usableDataSources, filteredSelectedDataSources, awaitingReindex);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Picks out the data sources whose index has to be rebuilt before they can be searched.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Asked for every data source at once, the same way the checks above run in parallel. Each
|
||||
/// answer is a single row read from the index database, and anything unclear counts as usable.
|
||||
/// </remarks>
|
||||
/// <param name="dataSources">The data sources which passed every other check.</param>
|
||||
/// <returns>Those of them which are waiting for their index, in the order they came in.</returns>
|
||||
private async Task<IReadOnlyList<IDataSource>> GetDataSourcesAwaitingReindex(IReadOnlyList<IDataSource> dataSources)
|
||||
{
|
||||
var checks = new List<Task<bool>>(dataSources.Count);
|
||||
foreach (var dataSource in dataSources)
|
||||
checks.Add(this.embeddingService.IsAwaitingReindexAsync(dataSource));
|
||||
|
||||
var awaitingReindex = new List<IDataSource>();
|
||||
for (var index = 0; index < dataSources.Count; index++)
|
||||
{
|
||||
if (await checks[index])
|
||||
{
|
||||
this.logger.LogInformation("The data source '{DataSourceName}' ({DataSourceId}) is waiting for its index to be rebuilt. It is shown, but cannot be selected.", dataSources[index].Name, dataSources[index].Id);
|
||||
awaitingReindex.Add(dataSources[index]);
|
||||
}
|
||||
}
|
||||
|
||||
return awaitingReindex;
|
||||
}
|
||||
|
||||
private async Task<IReadOnlyList<IDataSource>> GetAllowedDataSources(bool usingTrustedProvider, IReadOnlyList<ParticipatingProvider> participatingProviders, IReadOnlyCollection<IDataSource> requestedDataSources)
|
||||
|
||||
89
app/Tests/Tools/ReindexPendingTests.cs
Normal file
89
app/Tests/Tools/ReindexPendingTests.cs
Normal file
@ -0,0 +1,89 @@
|
||||
using AIStudio.Tools.Databases.IndexStore;
|
||||
using AIStudio.Tools.Services;
|
||||
|
||||
namespace AIStudio.Tests.Tools;
|
||||
|
||||
/// <summary>
|
||||
/// Checks when a data source counts as waiting for its index to be rebuilt.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This decides whether the data source selection greys a row out. Two mistakes are possible and
|
||||
/// both are bad in their own way: calling a rebuild finished lets the user pick a data source which
|
||||
/// finds nothing and answers without their data, while calling a healthy data source unusable locks
|
||||
/// a row for good. The stored signature alone cannot tell the two apart, because it is written back
|
||||
/// the moment the old index is discarded -- the stored hash of the data source is what closes that
|
||||
/// gap, since it only appears once a run has worked through everything.
|
||||
/// </remarks>
|
||||
[TestFixture]
|
||||
public sealed class ReindexPendingTests
|
||||
{
|
||||
private const string CURRENT_SIGNATURE = "v1|openai|text-embedding-3-small|512|100";
|
||||
|
||||
[Test]
|
||||
public void ADataSourceWhichWasNeverIndexedIsWaiting()
|
||||
{
|
||||
Assert.That(
|
||||
DataSourceEmbeddingService.IsIndexAwaitingRebuild(null, CURRENT_SIGNATURE, null),
|
||||
Is.True,
|
||||
"Nothing is stored about this data source, so there is nothing to search in it.");
|
||||
}
|
||||
|
||||
[Test]
|
||||
public void AnotherEmbeddingConfigurationMeansWaiting()
|
||||
{
|
||||
var indexState = new DataSourceIndexState("openai", "v1|openai|text-embedding-3-large|512|100", "source-hash", 1536);
|
||||
|
||||
Assert.That(
|
||||
DataSourceEmbeddingService.IsIndexAwaitingRebuild(indexState, CURRENT_SIGNATURE, DataSourceEmbeddingState.COMPLETED),
|
||||
Is.True,
|
||||
"The stored vectors belong to another embedding configuration and are discarded by the next run, so they are of no use now either.");
|
||||
}
|
||||
|
||||
[Test]
|
||||
public void AFinishedIndexIsNotWaiting()
|
||||
{
|
||||
var indexState = new DataSourceIndexState("openai", CURRENT_SIGNATURE, "source-hash", 1536);
|
||||
|
||||
Assert.That(
|
||||
DataSourceEmbeddingService.IsIndexAwaitingRebuild(indexState, CURRENT_SIGNATURE, DataSourceEmbeddingState.COMPLETED),
|
||||
Is.False,
|
||||
"A run has worked through the whole data source since the index was last discarded.");
|
||||
}
|
||||
|
||||
[Test]
|
||||
public void CatchingUpWithChangedFilesIsNotWaiting()
|
||||
{
|
||||
var indexState = new DataSourceIndexState("openai", CURRENT_SIGNATURE, "source-hash", 1536);
|
||||
|
||||
Assert.That(
|
||||
DataSourceEmbeddingService.IsIndexAwaitingRebuild(indexState, CURRENT_SIGNATURE, DataSourceEmbeddingState.RUNNING),
|
||||
Is.False,
|
||||
"An ordinary run leaves the stored hash in place: everything indexed before is still there and still searchable.");
|
||||
}
|
||||
|
||||
[Test]
|
||||
public void ARebuildInProgressIsWaiting()
|
||||
{
|
||||
//
|
||||
// What a reset leaves behind: the row was written anew with the current signature, and the
|
||||
// hash of the data source is empty until a run has been through all of it.
|
||||
//
|
||||
var indexState = new DataSourceIndexState("openai", CURRENT_SIGNATURE, string.Empty, 0);
|
||||
|
||||
Assert.That(
|
||||
DataSourceEmbeddingService.IsIndexAwaitingRebuild(indexState, CURRENT_SIGNATURE, DataSourceEmbeddingState.RUNNING),
|
||||
Is.True,
|
||||
"The signature matches again, but no run has finished since the vectors were thrown away.");
|
||||
}
|
||||
|
||||
[Test]
|
||||
public void AFailedRunIsNotWaiting()
|
||||
{
|
||||
var indexState = new DataSourceIndexState("openai", CURRENT_SIGNATURE, string.Empty, 1536);
|
||||
|
||||
Assert.That(
|
||||
DataSourceEmbeddingService.IsIndexAwaitingRebuild(indexState, CURRENT_SIGNATURE, DataSourceEmbeddingState.FAILED),
|
||||
Is.False,
|
||||
"Whatever the failed run managed to index is searchable, and the embeddings page already names the problem.");
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user