Keep data sources selectable only when their index can answer

This commit is contained in:
Thorsten Sommer 2026-09-16 20:14:26 +02:00
parent a8fded4df5
commit 843cca471b
Signed by untrusted user who does not match committer: tsommer
GPG Key ID: 371BBA77A02C0108
11 changed files with 391 additions and 47 deletions

View File

@ -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>
}

View File

@ -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,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);
/// <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)
{

View File

@ -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);

View File

@ -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);

View File

@ -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,

View File

@ -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,

View File

@ -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();

View File

@ -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);

View File

@ -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);

View File

@ -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<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)

View 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.");
}
}