using System.Collections.Concurrent; using System.Diagnostics.CodeAnalysis; using System.Threading.Channels; using AIStudio.Provider; using AIStudio.Settings; using AIStudio.Settings.DataModel; using AIStudio.Tools.Databases; using AIStudio.Tools.Databases.EmbeddingState; using AIStudio.Tools.Databases.VectorStore; using AIStudio.Tools.PluginSystem; namespace AIStudio.Tools.Services; public sealed partial class DataSourceEmbeddingService(SettingsManager settingsManager, RustService rustService, DatabaseClientProvider databaseClientProvider, ILogger logger) : BackgroundService { private const int MAX_CHUNK_LENGTH = 3_200; private const int MIN_CHUNK_LENGTH = 800; private const int CHUNK_OVERLAP_LENGTH = 320; private const int EMBEDDING_BATCH_SIZE = 16; private readonly Channel queue = Channel.CreateUnbounded(); private readonly ConcurrentDictionary queuedIds = new(StringComparer.OrdinalIgnoreCase); private readonly ConcurrentDictionary runningIds = new(StringComparer.OrdinalIgnoreCase); private readonly ConcurrentDictionary pendingQueueIds = new(StringComparer.OrdinalIgnoreCase); private readonly ConcurrentDictionary statuses = new(StringComparer.OrdinalIgnoreCase); private readonly object queueStateLock = new(); private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceEmbeddingService).Namespace, nameof(DataSourceEmbeddingService)); private enum DataSourceQueueRequestResult { QUEUED, ALREADY_QUEUED, RUNNING, RUNNING_MARKED_PENDING, } public IReadOnlyList GetStatuses() { return this.statuses.Values .OrderBy(status => status.SortOrder) .ThenBy(status => status.DataSourceName, StringComparer.OrdinalIgnoreCase) .ToList(); } public DataSourceEmbeddingOverview GetOverview() { var orderedStatuses = this.GetStatuses(); var activeStatus = orderedStatuses .FirstOrDefault(status => status.State is DataSourceEmbeddingState.QUEUED or DataSourceEmbeddingState.RUNNING); if (activeStatus is not null) { var total = Math.Max(activeStatus.TotalFiles, 1); return new( true, activeStatus.State, activeStatus.IndexedFiles, total, activeStatus.FailedFiles); } var failedStatus = orderedStatuses .FirstOrDefault(status => status.State is DataSourceEmbeddingState.FAILED || status.FailedFiles > 0); if (failedStatus is not null) return new(true, DataSourceEmbeddingState.FAILED, failedStatus.IndexedFiles, failedStatus.TotalFiles, failedStatus.FailedFiles); return new(false, DataSourceEmbeddingState.COMPLETED, 0, 0, 0); } public Task QueueAllInternalDataSourcesAsync() { return this.QueueAllInternalDataSourcesAsync(true); } private Task QueueAllInternalDataSourcesAsync(bool queueAfterCurrentRun) { this.RefreshWatchers(); var tasks = settingsManager.ConfigurationData.DataSources .Where(this.IsSupportedInternalDataSource) .Select(dataSource => this.QueueDataSourceAsync(dataSource, queueAfterCurrentRun)); return Task.WhenAll(tasks); } public Task QueueAllInternalDataSourcesIfAutomaticRefreshAsync() { if (!settingsManager.ConfigurationData.DataSourceIndexing.AutomaticRefresh) { this.RefreshWatchers(); return Task.CompletedTask; } return this.QueueAllInternalDataSourcesAsync(false); } public void RefreshAutomaticWatchers() { this.RefreshWatchers(); } public Task QueueDataSourceAsync(IDataSource dataSource) { return this.QueueDataSourceAsync(dataSource, true); } private async Task QueueDataSourceAsync(IDataSource dataSource, bool queueAfterCurrentRun) { if (!this.IsSupportedInternalDataSource(dataSource)) return; this.RefreshWatchers(); logger.LogDebug("Ensured watcher for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id); var queueRequestResult = this.TryReserveDataSourceQueueSlot(dataSource.Id, queueAfterCurrentRun); switch (queueRequestResult) { case DataSourceQueueRequestResult.ALREADY_QUEUED: logger.LogDebug("Data source '{DataSourceName}' ({DataSourceId}) is already queued for background embeddings. Ignoring duplicate queue request.", dataSource.Name, dataSource.Id); return; case DataSourceQueueRequestResult.RUNNING: logger.LogDebug("Data source '{DataSourceName}' ({DataSourceId}) is already being embedded. Ignoring duplicate queue request.", dataSource.Name, dataSource.Id); return; case DataSourceQueueRequestResult.RUNNING_MARKED_PENDING: logger.LogDebug("Data source '{DataSourceName}' ({DataSourceId}) is already being embedded. Scheduled one follow-up embedding run.", dataSource.Name, dataSource.Id); return; } logger.LogInformation("Queueing data source '{DataSourceName}' ({DataSourceId}) for background embeddings.", dataSource.Name, dataSource.Id); if (!this.statuses.TryGetValue(dataSource.Id, out var currentStatus) || currentStatus.State is not DataSourceEmbeddingState.RUNNING) this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.QUEUED, currentStatus?.TotalFiles ?? 0, currentStatus?.IndexedFiles ?? 0, currentStatus?.FailedFiles ?? 0)); logger.LogDebug("Upserting status for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id); await this.queue.Writer.WriteAsync(dataSource.Id); logger.LogDebug("Queued data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id); } public async Task RemoveDataSourceAsync(IDataSource dataSource) { if (!this.IsSupportedInternalDataSource(dataSource)) return; this.RemoveWatcher(dataSource.Id); this.statuses.TryRemove(dataSource.Id, out _); await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, null, null, CancellationToken.None); this.PublishStatusChanged(); } protected override async Task ExecuteAsync(CancellationToken stoppingToken) { await this.WaitForInitialSettingsAndBootstrapAsync(stoppingToken); while (!stoppingToken.IsCancellationRequested) { var dataSourceId = await this.queue.Reader.ReadAsync(stoppingToken); this.MarkDataSourceRunStarted(dataSourceId); IDataSource? dataSource = null; try { dataSource = settingsManager.ConfigurationData.DataSources .FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase)); if (dataSource is null || !this.IsSupportedInternalDataSource(dataSource)) continue; await this.ProcessDataSourceAsync(dataSource, stoppingToken); } catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested) { break; } catch (Exception exception) { if (dataSource is null) { logger.LogError(exception, "Background embedding failed for data source '{DataSourceId}'.", dataSourceId); } else { logger.LogError(exception, "Background embedding failed for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id); this.UpsertStatus(this.GetFallbackStatus(dataSource, exception.Message)); } } finally { await this.QueuePendingDataSourceRunAsync(dataSourceId, stoppingToken); } } } public override void Dispose() { this.DisposeWatchers(); base.Dispose(); } private async Task ProcessDataSourceAsync(IDataSource dataSource, CancellationToken token) { logger.LogInformation("Starting background embeddings for data source '{DataSourceName}' ({DataSourceId}).", dataSource.Name, dataSource.Id); var vectorStore = await databaseClientProvider.GetVectorStoreAsync(token); var embeddingState = await databaseClientProvider.GetEmbeddingStateAsync(token); if (!vectorStore.IsAvailable) { logger.LogWarning( "Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because the database client '{DatabaseName}' is unavailable.", dataSource.Name, dataSource.Id, vectorStore.Name); this.UpsertStatus(this.GetFallbackStatus(dataSource, "The vector database is not available.")); return; } if (!embeddingState.IsAvailable) { logger.LogWarning( "Skipping background embeddings for data source '{DataSourceName}' ({DataSourceId}) because the database client '{DatabaseName}' is unavailable.", dataSource.Name, dataSource.Id, embeddingState.Name); this.UpsertStatus(this.GetFallbackStatus(dataSource, "The SQLite embedding state database is not available.")); return; } if (!this.TryResolveEmbeddingProvider(dataSource, out var embeddingProvider)) { this.UpsertStatus(this.GetFallbackStatus(dataSource, "The selected embedding provider is not available.")); return; } logger.LogInformation( "Using embedding provider '{EmbeddingProviderId}' with model '{EmbeddingModelId}' for data source '{DataSourceName}' ({DataSourceId}).", embeddingProvider.Id, embeddingProvider.Model.Id, dataSource.Name, dataSource.Id); var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id); var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, embeddingState, token); var inputFiles = this.GetInputFiles(dataSource); var indexedFiles = inputFiles.Files; var totalFiles = indexedFiles.Count + inputFiles.FailedFiles; logger.LogInformation( "Prepared data source '{DataSourceName}' ({DataSourceId}) for embedding. AccessibleFiles={AccessibleFiles}, FailedFiles={FailedFiles}, Collection='{CollectionName}'.", dataSource.Name, dataSource.Id, indexedFiles.Count, inputFiles.FailedFiles, collectionName); await this.RemoveMissingFileEmbeddingsAsync(vectorStore, embeddingState, dataSource, collectionName, manifest, indexedFiles, token); this.UpsertStatus(this.CreateStatus( dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, 0, inputFiles.FailedFiles, lastError: inputFiles.LastError)); var provider = embeddingProvider.CreateProvider(); var skippedFiles = 0; var completedFiles = 0; var failedFiles = inputFiles.FailedFiles; var lastError = inputFiles.LastError; foreach (var file in indexedFiles) { token.ThrowIfCancellationRequested(); var fingerprint = await this.BuildFingerprintAsync(file, token); if (manifest.Files.TryGetValue(file.FullName, out var existingRecord) && string.Equals(existingRecord.Fingerprint, fingerprint, StringComparison.Ordinal)) { logger.LogDebug( "Skipping unchanged file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}).", file.FullName, dataSource.Name, dataSource.Id); skippedFiles++; this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, skippedFiles + completedFiles, failedFiles, lastError: lastError)); continue; } this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, skippedFiles + completedFiles, failedFiles, file.Name, lastError)); try { logger.LogInformation( "Embedding file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}). Progress={CompletedFiles}/{TotalFiles}.", file.FullName, dataSource.Name, dataSource.Id, skippedFiles + completedFiles + 1, totalFiles); var startedAtUtc = DateTime.UtcNow; var chunkCount = await this.IndexOneFileAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token); var embeddedAtUtc = DateTime.UtcNow; var record = new EmbeddedFileRecord( fingerprint, file.Length, file.LastWriteTimeUtc, embeddedAtUtc, chunkCount); await embeddingState.UpsertFileAsync( dataSource.Id, new EmbeddingStateFile( file.FullName, file.Name, this.TryGetRelativePath(dataSource, file), fingerprint, file.Length, file.LastWriteTimeUtc, embeddedAtUtc, chunkCount), token); manifest.Files[file.FullName] = record; completedFiles++; logger.LogInformation( "Embedded file '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}) successfully. Chunks={ChunkCount}, DurationMs={DurationMs}.", file.FullName, dataSource.Name, dataSource.Id, chunkCount, (DateTime.UtcNow - startedAtUtc).TotalMilliseconds); } catch (Exception exception) { failedFiles++; lastError = exception.Message; manifest.Files.Remove(file.FullName); await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token); await embeddingState.DeleteFileAsync(dataSource.Id, file.FullName, token); logger.LogWarning(exception, "Failed to embed file '{FilePath}' for data source '{DataSourceName}'.", file.FullName, dataSource.Name); this.UpsertStatus(this.CreateStatus(dataSource, DataSourceEmbeddingState.RUNNING, totalFiles, skippedFiles + completedFiles, failedFiles, file.Name, exception.Message)); } } this.UpsertStatus(this.CreateCompletedStatus(dataSource, totalFiles, skippedFiles + completedFiles, failedFiles, lastError)); logger.LogInformation( "Finished background embeddings for data source '{DataSourceName}' ({DataSourceId}). Indexed={IndexedFiles}, Failed={FailedFiles}, Total={TotalFiles}.", dataSource.Name, dataSource.Id, skippedFiles + completedFiles, failedFiles, totalFiles); } private async Task IndexOneFileAsync( EmbeddingStateClient embeddingState, VectorStoreClient vectorStore, IDataSource dataSource, FileInfo file, string fingerprint, EmbeddingProvider embeddingProvider, IProvider provider, DataSourceEmbeddingManifest manifest, CancellationToken token) { var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id); logger.LogDebug( "Resetting stored embeddings for file '{FilePath}' in collection '{CollectionName}' before re-indexing.", file.FullName, collectionName); await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token); var batch = new List<(string Text, int ChunkIndex)>(EMBEDDING_BATCH_SIZE); var totalChunkCount = 0; await foreach (var chunk in this.StreamEmbeddingChunksAsync(file.FullName, embeddingProvider, token)) { batch.Add((chunk, totalChunkCount)); totalChunkCount++; if (batch.Count >= EMBEDDING_BATCH_SIZE) await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token); } if (batch.Count > 0) await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token); if (totalChunkCount == 0) throw new InvalidOperationException($"The file '{file.Name}' did not yield any text chunks."); logger.LogDebug( "Generated {ChunkCount} chunks for file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).", totalChunkCount, file.FullName, dataSource.Name, dataSource.Id); return totalChunkCount; } private async Task FlushBatchAsync( EmbeddingStateClient embeddingState, VectorStoreClient vectorStore, IDataSource dataSource, FileInfo file, string fingerprint, EmbeddingProvider embeddingProvider, IProvider provider, DataSourceEmbeddingManifest manifest, string collectionName, List<(string Text, int ChunkIndex)> batch, CancellationToken token) { logger.LogDebug( "Requesting embeddings for batch of {ChunkCount} chunks from file '{FilePath}' in data source '{DataSourceName}' ({DataSourceId}).", batch.Count, file.FullName, dataSource.Name, dataSource.Id); var texts = batch.Select(item => item.Text).ToList(); var vectors = await provider.EmbedTextAsync(embeddingProvider.Model, settingsManager, token, texts); if (vectors.Count != batch.Count) throw new InvalidOperationException($"The embedding provider returned {vectors.Count} vectors for {batch.Count} text chunks."); var vectorSize = vectors.FirstOrDefault()?.Count ?? 0; if (vectorSize <= 0) throw new InvalidOperationException("The embedding provider returned an empty vector."); if (manifest.VectorSize > 0 && manifest.VectorSize != vectorSize) throw new InvalidOperationException($"The embedding vector size changed from {manifest.VectorSize} to {vectorSize}. Please re-save the data source to trigger a clean re-index."); if (manifest.VectorSize == 0) { manifest.VectorSize = vectorSize; await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token); await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token); logger.LogInformation( "Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).", collectionName, vectorSize, dataSource.Name, dataSource.Id); } await this.UpsertPointsAsync( vectorStore, collectionName, dataSource, file, fingerprint, batch, vectors, this.TryGetRelativePath(dataSource, file), token); logger.LogDebug( "Stored {ChunkCount} embedded chunks for file '{FilePath}' in collection '{CollectionName}'.", batch.Count, file.FullName, collectionName); batch.Clear(); } private async Task EnsureCollectionExistsAsync(VectorStoreClient vectorStore, string collectionName, int vectorSize, CancellationToken token) { await vectorStore.EnsureVectorStoreExists(collectionName, vectorSize, token); } private async Task UpsertPointsAsync( VectorStoreClient vectorStore, string collectionName, IDataSource dataSource, FileInfo file, string fingerprint, IReadOnlyList<(string Text, int ChunkIndex)> batch, IReadOnlyList> vectors, string relativePath, CancellationToken token) { var embeddedAtUtc = DateTime.UtcNow; var points = batch.Select((item, index) => new VectorStoragePoint( this.CreatePointId(dataSource.Id, fingerprint, item.ChunkIndex), vectors[index], dataSource.Id, dataSource.Name, dataSource.Type.ToString(), file.FullName, file.Name, relativePath, item.ChunkIndex, item.Text, fingerprint, file.LastWriteTimeUtc, embeddedAtUtc)).ToList(); await vectorStore.InsertEmbedding(collectionName, points, token); } private async Task DeleteFilePointsAsync(VectorStoreClient vectorStore, string collectionName, string filePath, CancellationToken token) { await vectorStore.DeleteEmbeddingByFile(collectionName, filePath, token); } private async Task DeleteCollectionAsync(string collectionName, VectorStoreClient? vectorStore, CancellationToken token) { vectorStore ??= await databaseClientProvider.GetVectorStoreAsync(token); if (!vectorStore.IsAvailable) { logger.LogWarning("Could not delete embedding collection '{CollectionName}' because the vector store '{VectorStoreName}' is unavailable.", collectionName, vectorStore.Name); return; } await vectorStore.DeleteVectorStore(collectionName, token); } private async Task WaitForInitialSettingsAndBootstrapAsync(CancellationToken token) { while (!token.IsCancellationRequested) { if (settingsManager.HasCompletedInitialSettingsLoad && !string.IsNullOrWhiteSpace(SettingsManager.ConfigDirectory) && !string.IsNullOrWhiteSpace(SettingsManager.DataDirectory)) { break; } await Task.Delay(250, token); } token.ThrowIfCancellationRequested(); logger.LogInformation("Embedding background service is ready. Checking whether automatic data source refresh is enabled."); await this.QueueAllInternalDataSourcesIfAutomaticRefreshAsync(); } private bool IsSupportedInternalDataSource(IDataSource dataSource) { return dataSource is DataSourceLocalDirectory or DataSourceLocalFile; } private bool TryResolveEmbeddingProvider(IDataSource dataSource, [NotNullWhen(true)] out EmbeddingProvider? embeddingProvider) { embeddingProvider = settingsManager.ConfigurationData.EmbeddingProviders.FirstOrDefault(provider => dataSource is IInternalDataSource internalDataSource && provider.Id.Equals(internalDataSource.EmbeddingId, StringComparison.OrdinalIgnoreCase)); return embeddingProvider != default && embeddingProvider.UsedLLMProvider is not LLMProviders.NONE; } private async Task EnsureCompatibleManifestAsync( IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, VectorStoreClient vectorStore, EmbeddingStateClient embeddingState, CancellationToken token) { var embeddingSignature = this.BuildEmbeddingSignature(embeddingProvider); var manifest = await embeddingState.GetManifestAsync(dataSource.Id, token); if (!string.Equals(manifest.EmbeddingSignature, embeddingSignature, StringComparison.Ordinal)) { logger.LogInformation( "Embedding configuration changed for data source '{DataSourceName}' ({DataSourceId}). Resetting persisted state and collection '{CollectionName}'.", dataSource.Name, dataSource.Id, collectionName); await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, vectorStore, embeddingState, token); manifest = await embeddingState.GetManifestAsync(dataSource.Id, token); } if (!string.Equals(manifest.EmbeddingProviderId, embeddingProvider.Id, StringComparison.OrdinalIgnoreCase) || !string.Equals(manifest.EmbeddingSignature, embeddingSignature, StringComparison.Ordinal)) { manifest.EmbeddingProviderId = embeddingProvider.Id; manifest.EmbeddingSignature = embeddingSignature; } await embeddingState.UpsertDataSourceAsync( dataSource.Id, dataSource.Name, dataSource.Type.ToString(), manifest.EmbeddingProviderId, manifest.EmbeddingSignature, manifest.VectorSize, token); return manifest; } private async Task RemoveMissingFileEmbeddingsAsync( VectorStoreClient vectorStore, EmbeddingStateClient embeddingState, IDataSource dataSource, string collectionName, DataSourceEmbeddingManifest manifest, IReadOnlyCollection indexedFiles, CancellationToken token) { var existingPaths = indexedFiles .Select(file => file.FullName) .ToHashSet(StringComparer.OrdinalIgnoreCase); foreach (var removedFilePath in manifest.Files.Keys.Except(existingPaths, StringComparer.OrdinalIgnoreCase).ToList()) { await this.DeleteFilePointsAsync(vectorStore, collectionName, removedFilePath, token); await embeddingState.DeleteFileAsync(dataSource.Id, removedFilePath, token); manifest.Files.Remove(removedFilePath); logger.LogInformation( "Removed stale embeddings for deleted file '{FilePath}' from data source '{DataSourceName}' ({DataSourceId}).", removedFilePath, dataSource.Name, dataSource.Id); } } private DataSourceEmbeddingStatus CreateStatus( IDataSource dataSource, DataSourceEmbeddingState state, int totalFiles, int indexedFiles, int failedFiles, string currentFile = "", string lastError = "") { return new DataSourceEmbeddingStatus( dataSource.Id, dataSource.Name, dataSource.Type, state, totalFiles, indexedFiles, failedFiles, currentFile, lastError); } private DataSourceEmbeddingStatus CreateCompletedStatus(IDataSource dataSource, int totalFiles, int indexedFiles, int failedFiles, string lastError) { return this.CreateStatus( dataSource, failedFiles > 0 ? DataSourceEmbeddingState.FAILED : DataSourceEmbeddingState.COMPLETED, totalFiles, indexedFiles, failedFiles, lastError: failedFiles > 0 ? string.IsNullOrWhiteSpace(lastError) ? "Some files could not be embedded. See the logs for details." : lastError : string.Empty); } private DataSourceEmbeddingStatus GetFallbackStatus(IDataSource dataSource, string errorMessage) { return this.CreateStatus(dataSource, DataSourceEmbeddingState.FAILED, 0, 0, 1, lastError: errorMessage); } private DataSourceQueueRequestResult TryReserveDataSourceQueueSlot(string dataSourceId, bool queueAfterCurrentRun) { lock (this.queueStateLock) { if (this.runningIds.ContainsKey(dataSourceId)) { if (queueAfterCurrentRun && this.pendingQueueIds.TryAdd(dataSourceId, 0)) return DataSourceQueueRequestResult.RUNNING_MARKED_PENDING; return DataSourceQueueRequestResult.RUNNING; } if (!this.queuedIds.TryAdd(dataSourceId, 0)) return DataSourceQueueRequestResult.ALREADY_QUEUED; return DataSourceQueueRequestResult.QUEUED; } } private void MarkDataSourceRunStarted(string dataSourceId) { lock (this.queueStateLock) { this.queuedIds.TryRemove(dataSourceId, out _); this.runningIds.TryAdd(dataSourceId, 0); } } private bool TryCompleteDataSourceRun(string dataSourceId, bool allowPendingRequeue) { lock (this.queueStateLock) { this.runningIds.TryRemove(dataSourceId, out _); if (!this.pendingQueueIds.TryRemove(dataSourceId, out _)) return false; return allowPendingRequeue && this.queuedIds.TryAdd(dataSourceId, 0); } } private void ReleaseQueuedDataSourceRun(string dataSourceId) { lock (this.queueStateLock) { this.queuedIds.TryRemove(dataSourceId, out _); } } private async Task QueuePendingDataSourceRunAsync(string dataSourceId, CancellationToken token) { var dataSource = token.IsCancellationRequested ? null : settingsManager.ConfigurationData.DataSources .FirstOrDefault(source => source.Id.Equals(dataSourceId, StringComparison.OrdinalIgnoreCase)); if (!this.TryCompleteDataSourceRun(dataSourceId, dataSource is not null && this.IsSupportedInternalDataSource(dataSource))) return; if (dataSource is null) { this.ReleaseQueuedDataSourceRun(dataSourceId); return; } logger.LogInformation("Queueing one follow-up embedding run for data source '{DataSourceName}' ({DataSourceId}) after changes arrived during the previous run.", dataSource.Name, dataSource.Id); this.statuses.TryGetValue(dataSource.Id, out var currentStatus); this.UpsertStatus(this.CreateStatus( dataSource, DataSourceEmbeddingState.QUEUED, currentStatus?.TotalFiles ?? 0, currentStatus?.IndexedFiles ?? 0, currentStatus?.FailedFiles ?? 0, lastError: currentStatus?.LastError ?? string.Empty)); try { await this.queue.Writer.WriteAsync(dataSourceId, token); } catch (OperationCanceledException) when (token.IsCancellationRequested) { this.ReleaseQueuedDataSourceRun(dataSourceId); } } private void UpsertStatus(DataSourceEmbeddingStatus status) { this.statuses[status.DataSourceId] = status; this.PublishStatusChanged(); } private void PublishStatusChanged() { _ = MessageBus.INSTANCE.SendMessage(null, Event.RAG_EMBEDDING_STATUS_CHANGED, true); } }