From 843cca471b503d2eb8b3e87ab2066d0202797d39 Mon Sep 17 00:00:00 2001 From: Thorsten Sommer Date: Wed, 16 Sep 2026 20:14:26 +0200 Subject: [PATCH] Keep data sources selectable only when their index can answer --- .../Components/DataSourceSelection.razor | 73 ++++++------ .../Components/DataSourceSelection.razor.cs | 34 ++++++ .../Tools/AllowedSelectedDataSources.cs | 11 +- .../IndexStore/DataSourceIndexState.cs | 19 ++++ .../Databases/IndexStore/IndexStoreClient.cs | 8 ++ .../IndexStore/NoIndexStoreClient.cs | 5 +- .../SqliteIndexStoreClientImplementation.cs | 10 ++ .../Services/DataSourceEmbeddingService.cs | 106 ++++++++++++++++++ .../DataSourceLocalRetrievalService.cs | 19 +++- .../Tools/Services/DataSourceService.cs | 64 +++++++++-- app/Tests/Tools/ReindexPendingTests.cs | 89 +++++++++++++++ 11 files changed, 391 insertions(+), 47 deletions(-) create mode 100644 app/MindWork AI Studio/Tools/Databases/IndexStore/DataSourceIndexState.cs create mode 100644 app/Tests/Tools/ReindexPendingTests.cs diff --git a/app/MindWork AI Studio/Components/DataSourceSelection.razor b/app/MindWork AI Studio/Components/DataSourceSelection.razor index de7ae518..3fcb1c83 100644 --- a/app/MindWork AI Studio/Components/DataSourceSelection.razor +++ b/app/MindWork AI Studio/Components/DataSourceSelection.razor @@ -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 dataSourceRow = source => + @ + + + + @source.Name + + @if (source is IInternalDataSource internalSource) + { + + @if (this.IsAwaitingReindex(source)) + { + + } + + + + } + + + ; +} @if (this.SelectionMode is DataSourceSelectionMode.SELECTION_MODE) {
@@ -69,7 +100,7 @@ @switch (this.aiBasedSourceSelection) { - case true when this.availableDataSources.Count == 0: + case true when this.GetListedDataSources().Count == 0: @T("Your data sources cannot be used with the selected providers due to data privacy or confidence-level requirements, or they are currently unavailable.") @@ -81,7 +112,7 @@ break; - case false when this.availableDataSources.Count == 0: + case false when this.GetListedDataSources().Count == 0: @T("Your data sources cannot be used with the selected providers due to data privacy or confidence-level requirements, or they are currently unavailable.") @@ -90,22 +121,9 @@ case false: - @foreach (var source in this.availableDataSources) + @foreach (var source in this.GetListedDataSources()) { - - - - @source.Name - - @if (source is IInternalDataSource internalSource) - { - - - - - } - - + @dataSourceRow(source) } @@ -115,22 +133,9 @@ - @foreach (var source in this.availableDataSources) + @foreach (var source in this.GetListedDataSources()) { - - - - @source.Name - - @if (source is IInternalDataSource internalSource) - { - - - - - } - - + @dataSourceRow(source) } @@ -166,13 +171,13 @@ break; } - @if (!this.aiBasedSourceSelection && this.GetUnavailablePreselectedDataSources().Count > 0) + @if (!this.aiBasedSourceSelection && this.GetUnavailablePreselectedDataSourcesToList().Count > 0) { @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:")
    - @foreach (var source in this.GetUnavailablePreselectedDataSources()) + @foreach (var source in this.GetUnavailablePreselectedDataSourcesToList()) {
  • @source.Name
    • } diff --git a/app/MindWork AI Studio/Components/DataSourceSelection.razor.cs b/app/MindWork AI Studio/Components/DataSourceSelection.razor.cs index 18ddf6a8..b2f7b8ce 100644 --- a/app/MindWork AI Studio/Components/DataSourceSelection.razor.cs +++ b/app/MindWork AI Studio/Components/DataSourceSelection.razor.cs @@ -49,6 +49,8 @@ public partial class DataSourceSelection : MSGComponentBase private bool showDataSourceSelection; private bool waitingForDataSources = true; private IReadOnlyList availableDataSources = []; + private IReadOnlyList dataSourcesAwaitingReindex = []; + private HashSet dataSourceIdsAwaitingReindex = new(StringComparer.Ordinal); private IReadOnlyCollection selectedDataSources = []; private bool aiBasedSourceSelection; private bool aiBasedValidation; @@ -226,10 +228,42 @@ 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); + + /// + /// The data sources the list shows: the usable ones, plus the ones waiting for their index. + /// + /// + /// 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. + /// + private IReadOnlyList 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(); + } + + /// + /// The preselected but unusable data sources the warning box lists. + /// + /// + /// 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. + /// + private IReadOnlyList GetUnavailablePreselectedDataSourcesToList() => + this.GetUnavailablePreselectedDataSources().Where(source => !this.IsAwaitingReindex(source)).ToList(); private async Task EnabledChanged(bool state) { diff --git a/app/MindWork AI Studio/Tools/AllowedSelectedDataSources.cs b/app/MindWork AI Studio/Tools/AllowedSelectedDataSources.cs index 1aed9d1c..b2cf5d9c 100644 --- a/app/MindWork AI Studio/Tools/AllowedSelectedDataSources.cs +++ b/app/MindWork AI Studio/Tools/AllowedSelectedDataSources.cs @@ -3,11 +3,18 @@ using AIStudio.Settings; namespace AIStudio.Tools; /// -/// Contains both the allowed and selected data sources. +/// Contains the allowed and selected data sources, plus the ones waiting for their index. /// /// /// 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. /// /// The allowed data sources. /// The selected data sources, which are a subset of the allowed data sources. -public readonly record struct AllowedSelectedDataSources(IReadOnlyList AllowedDataSources, IReadOnlyList SelectedDataSources); \ No newline at end of file +/// The data sources which passed every check but cannot be searched until their index has been rebuilt. +public readonly record struct AllowedSelectedDataSources(IReadOnlyList AllowedDataSources, IReadOnlyList SelectedDataSources, IReadOnlyList DataSourcesAwaitingReindex); \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Databases/IndexStore/DataSourceIndexState.cs b/app/MindWork AI Studio/Tools/Databases/IndexStore/DataSourceIndexState.cs new file mode 100644 index 00000000..aebd29fb --- /dev/null +++ b/app/MindWork AI Studio/Tools/Databases/IndexStore/DataSourceIndexState.cs @@ -0,0 +1,19 @@ +namespace AIStudio.Tools.Databases.IndexStore; + +/// +/// What the index knows about a data source as a whole, without its files. +/// +/// +/// 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. +/// +/// The embedding provider the stored vectors were created with. +/// Identifies the embedding configuration the stored vectors belong to. +/// The hash of the data source as a whole, written when a run completes. +/// The dimension of the stored vectors. +public sealed record DataSourceIndexState(string EmbeddingProviderId, string EmbeddingSignature, string SourceHash, int VectorSize); \ No newline at end of file diff --git a/app/MindWork AI Studio/Tools/Databases/IndexStore/IndexStoreClient.cs b/app/MindWork AI Studio/Tools/Databases/IndexStore/IndexStoreClient.cs index 1cf0ae49..644529f2 100644 --- a/app/MindWork AI Studio/Tools/Databases/IndexStore/IndexStoreClient.cs +++ b/app/MindWork AI Studio/Tools/Databases/IndexStore/IndexStoreClient.cs @@ -6,6 +6,14 @@ public abstract class IndexStoreClient(string name, string path) : DatabaseClien { public abstract Task GetManifestAsync(string dataSourceId, CancellationToken token); + /// + /// Reads what the index knows about a data source as a whole, without its files. + /// + /// The data source to read. + /// The cancellation token. + /// The stored state, or null when the index holds nothing about this data source. + public abstract Task GetDataSourceStateAsync(string dataSourceId, CancellationToken token); + public abstract Task UpsertDataSourceAsync( string dataSourceId, string dataSourceType, diff --git a/app/MindWork AI Studio/Tools/Databases/IndexStore/NoIndexStoreClient.cs b/app/MindWork AI Studio/Tools/Databases/IndexStore/NoIndexStoreClient.cs index 4ab41065..b9373ea0 100644 --- a/app/MindWork AI Studio/Tools/Databases/IndexStore/NoIndexStoreClient.cs +++ b/app/MindWork AI Studio/Tools/Databases/IndexStore/NoIndexStoreClient.cs @@ -23,8 +23,9 @@ public sealed class NoIndexStoreClient(string name, string? unavailableReason, D await Task.CompletedTask; } - public override Task GetManifestAsync(string dataSourceId, CancellationToken token) => - Task.FromResult(new DataSourceEmbeddingManifest()); + public override Task GetManifestAsync(string dataSourceId, CancellationToken token) => Task.FromResult(new DataSourceEmbeddingManifest()); + + public override Task GetDataSourceStateAsync(string dataSourceId, CancellationToken token) => Task.FromResult(null); public override Task UpsertDataSourceAsync( string dataSourceId, diff --git a/app/MindWork AI Studio/Tools/Databases/IndexStore/SqliteIndexStoreClientImplementation.cs b/app/MindWork AI Studio/Tools/Databases/IndexStore/SqliteIndexStoreClientImplementation.cs index 1593d4a1..739b0c4a 100644 --- a/app/MindWork AI Studio/Tools/Databases/IndexStore/SqliteIndexStoreClientImplementation.cs +++ b/app/MindWork AI Studio/Tools/Databases/IndexStore/SqliteIndexStoreClientImplementation.cs @@ -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 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 GetManifestAsync(string dataSourceId, CancellationToken token) { await using var context = this.CreateContext(); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs index 553663cb..1f32eb14 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.cs @@ -23,6 +23,15 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM /// private static readonly TimeSpan BLOCK_PROGRESS_INTERVAL = TimeSpan.FromSeconds(3); + /// + /// How long the re-index check waits for the index database before it gives up. + /// + /// + /// Asked while somebody waits for the data source selection to open, and possibly while a run + /// writes to the same database. + /// + private static readonly TimeSpan REINDEX_CHECK_TIMEOUT = TimeSpan.FromSeconds(2); + private readonly Channel queue = Channel.CreateUnbounded(); private readonly ConcurrentDictionary queuedIds = new(StringComparer.OrdinalIgnoreCase); private readonly ConcurrentDictionary runningIds = new(StringComparer.OrdinalIgnoreCase); @@ -211,6 +220,103 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM || manifest.PermanentFailures.Count > 0; } + /// + /// Whether a data source cannot answer a search right now because its index has to be built anew. + /// + /// + /// 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. + /// + /// The data source to ask about. + /// The cancellation token. + /// True when the data source is waiting for its index to be rebuilt. + public async Task 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; + } + } + + /// + /// Decides from the stored index state alone whether a data source has to be indexed anew. + /// + /// + /// 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. + /// + /// What the index holds about the data source, or null when it holds nothing. + /// The signature the current embedding configuration produces. + /// The state of this data source's last or current run, when one is known. + /// True when the data source is waiting for its index to be rebuilt. + 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); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs index 448b6e70..efa318e7 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs @@ -13,7 +13,7 @@ namespace AIStudio.Tools.Services; public sealed class DataSourceLocalRetrievalService( SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider, - ILogger logger) + DataSourceEmbeddingService embeddingService, ILogger 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); diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceService.cs index 38bdd98d..8c70e82b 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceService.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceService.cs @@ -18,16 +18,18 @@ 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 logger; - public DataSourceService(SettingsManager settingsManager, ILogger logger, RustService rustService) + public DataSourceService(SettingsManager settingsManager, ILogger 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(); } /// @@ -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); + } + + /// + /// Picks out the data sources whose index has to be rebuilt before they can be searched. + /// + /// + /// 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. + /// + /// The data sources which passed every other check. + /// Those of them which are waiting for their index, in the order they came in. + private async Task> GetDataSourcesAwaitingReindex(IReadOnlyList dataSources) + { + var checks = new List>(dataSources.Count); + foreach (var dataSource in dataSources) + checks.Add(this.embeddingService.IsAwaitingReindexAsync(dataSource)); + + var awaitingReindex = new List(); + 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> GetAllowedDataSources(bool usingTrustedProvider, IReadOnlyList participatingProviders, IReadOnlyCollection requestedDataSources) diff --git a/app/Tests/Tools/ReindexPendingTests.cs b/app/Tests/Tools/ReindexPendingTests.cs new file mode 100644 index 00000000..3ffccdaa --- /dev/null +++ b/app/Tests/Tools/ReindexPendingTests.cs @@ -0,0 +1,89 @@ +using AIStudio.Tools.Databases.IndexStore; +using AIStudio.Tools.Services; + +namespace AIStudio.Tests.Tools; + +/// +/// Checks when a data source counts as waiting for its index to be rebuilt. +/// +/// +/// 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. +/// +[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."); + } +} \ No newline at end of file