added sql integration instead of json files

This commit is contained in:
PaulKoudelka committed 2026-07-28 15:25:10 +02:00
1 parent 2e06af13e0
commit c59e33e1b1
22 files changed
+700 -156

No files matched your search

@@ -62,11 +62,6 @@ public sealed class FileEnumerationResult
public string LastError { get; set; } = string.Empty;
}
public sealed class PersistedEmbeddingState
{
public Dictionary<string, DataSourceEmbeddingManifest> DataSources { get; init; } = new(StringComparer.OrdinalIgnoreCase);
}
public sealed class DataSourceEmbeddingManifest
{
public string EmbeddingProviderId { get; set; } = string.Empty;
@@ -10,9 +10,20 @@ namespace AIStudio.Tools.Services;
public sealed partial class DataSourceEmbeddingService
{
private static readonly string[] ADDITIONAL_RAG_FILE_EXTENSIONS = ["csv", "tsv", "ods", "xlsm", "xlsb", "xla", "xlam"];
private const string OFFICE_LOCK_FILE_PREFIX = "~$";
private static readonly string[] RAG_DELIMITED_TABLE_FILE_EXTENSIONS = ["csv", "tsv"];
private static readonly string[] RAG_SPREADSHEET_FILE_EXTENSIONS = ["ods", "xlsm", "xlsb"];
private static readonly string[] RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS = ["xla", "xlam"];
private static readonly string[] SKIPPED_RAG_FILE_EXTENSIONS = ["lnk"];
private enum RagFileIndexingDecision
{
INDEXABLE,
EXCLUDED,
UNSUPPORTED,
}
private async IAsyncEnumerable<string> StreamEmbeddingChunksAsync(string filePath, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken token)
{
if (this.IsImageFilePath(filePath))
@@ -67,14 +78,21 @@ public sealed partial class DataSourceEmbeddingService
switch (dataSource)
{
case DataSourceLocalFile localFile when File.Exists(localFile.FilePath):
if (this.IsSupportedRagFilePath(localFile.FilePath))
var file = new FileInfo(localFile.FilePath);
switch (this.GetRagFileIndexingDecision(file))
{
result.Files.Add(new FileInfo(localFile.FilePath));
}
else
{
result.FailedFiles = 1;
result.LastError = $"The selected file '{localFile.FilePath}' is not supported for background embeddings.";
case RagFileIndexingDecision.INDEXABLE:
result.Files.Add(file);
break;
case RagFileIndexingDecision.EXCLUDED:
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", file.FullName);
break;
default:
result.FailedFiles = 1;
result.LastError = $"The selected file '{localFile.FilePath}' is not supported for background embeddings.";
break;
}
return result;
@@ -141,14 +159,25 @@ public sealed partial class DataSourceEmbeddingService
continue;
}
if (!this.IsSupportedRagFilePath(fileInfo.FullName))
continue;
switch (this.GetRagFileIndexingDecision(fileInfo))
{
case RagFileIndexingDecision.INDEXABLE:
result.Files.Add(fileInfo);
break;
result.Files.Add(fileInfo);
case RagFileIndexingDecision.EXCLUDED:
logger.LogDebug("Skipping excluded file '{FilePath}' while indexing.", fileInfo.FullName);
break;
}
}
foreach (var subDirectory in subDirectories)
{
if (this.IsSkippedRagDirectory(subDirectory))
continue;
pendingDirectories.Push(subDirectory);
}
}
}
@@ -174,11 +203,59 @@ public sealed partial class DataSourceEmbeddingService
private bool IsSupportedRagFilePath(string filePath)
{
var extension = Path.GetExtension(filePath).TrimStart('.');
if (SKIPPED_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase))
return false;
return FileTypes.IsAllowedPath(filePath, FileTypes.DOCUMENT, FileTypes.IMAGE)
|| ADDITIONAL_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase);
|| RAG_DELIMITED_TABLE_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase)
|| RAG_SPREADSHEET_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase)
|| RAG_SPREADSHEET_ADD_IN_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase);
}
private RagFileIndexingDecision GetRagFileIndexingDecision(FileInfo file)
{
if (this.IsSkippedRagFile(file))
return RagFileIndexingDecision.EXCLUDED;
return this.IsSupportedRagFilePath(file.FullName)
? RagFileIndexingDecision.INDEXABLE
: RagFileIndexingDecision.UNSUPPORTED;
}
private bool IsSkippedRagFile(FileInfo file)
{
var extension = file.Extension.TrimStart('.');
if (SKIPPED_RAG_FILE_EXTENSIONS.Contains(extension, StringComparer.OrdinalIgnoreCase))
return true;
if (file.Name.StartsWith(OFFICE_LOCK_FILE_PREFIX, StringComparison.Ordinal))
return true;
try
{
return file.Attributes.HasFlag(FileAttributes.ReparsePoint)
|| file.Attributes.HasFlag(FileAttributes.Offline)
|| file.Attributes.HasFlag(FileAttributes.Temporary)
|| file.Attributes.HasFlag(FileAttributes.System);
}
catch (Exception exception)
{
logger.LogWarning(exception, "Cannot inspect file '{FilePath}' while indexing.", file.FullName);
return true;
}
}
private bool IsSkippedRagDirectory(string path)
{
try
{
var directory = new DirectoryInfo(path);
return directory.Attributes.HasFlag(FileAttributes.ReparsePoint)
|| directory.Attributes.HasFlag(FileAttributes.Offline)
|| directory.Attributes.HasFlag(FileAttributes.System);
}
catch (Exception exception)
{
logger.LogWarning(exception, "Cannot inspect directory '{DirectoryPath}' while indexing.", path);
return true;
}
}
private string BuildImageIndexText(string filePath)
@@ -212,9 +289,17 @@ public sealed partial class DataSourceEmbeddingService
embeddingProvider.TokenizerPath);
}
private string BuildFingerprint(FileInfo file)
private async Task<string> BuildFingerprintAsync(FileInfo file, CancellationToken token)
{
var fingerprintSource = $"{file.FullName}|{file.Length}|{file.LastWriteTimeUtc.Ticks}";
await using var stream = new FileStream(
file.FullName,
FileMode.Open,
FileAccess.Read,
FileShare.ReadWrite | FileShare.Delete,
1024 * 128,
FileOptions.Asynchronous | FileOptions.SequentialScan);
var contentHash = await SHA256.HashDataAsync(stream, token);
var fingerprintSource = $"{file.FullName}|{Convert.ToHexString(contentHash)}";
var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(fingerprintSource));
return Convert.ToHexString(bytes);
}
@@ -1,90 +1,27 @@
using System.Text.Json;
using AIStudio.Settings;
using AIStudio.Tools.Databases.EmbeddingState;
using AIStudio.Tools.Databases.VectorStore;
namespace AIStudio.Tools.Services;
public sealed partial class DataSourceEmbeddingService
{
private const string STATE_FILENAME = "rag-embedding-state.json";
private readonly JsonSerializerOptions jsonOptions = new()
private async Task ResetPersistedStateAsync(
string dataSourceName,
string dataSourceId,
VectorStoreClient? vectorStore,
EmbeddingStateClient? embeddingState,
CancellationToken token)
{
WriteIndented = true,
};
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceName, dataSourceId), vectorStore, token);
private async Task EnsureStateLoadedAsync(CancellationToken token)
{
if (this.stateLoaded)
embeddingState ??= await databaseClientProvider.GetEmbeddingStateAsync(token);
if (!embeddingState.IsAvailable)
{
logger.LogWarning("Could not delete SQLite embedding state for data source '{DataSourceId}' because the embedding state database '{DatabaseName}' is unavailable.", dataSourceId, embeddingState.Name);
return;
await this.stateLock.WaitAsync(token);
try
{
if (this.stateLoaded)
return;
var statePath = this.GetStatePath();
if (!string.IsNullOrWhiteSpace(statePath) && File.Exists(statePath))
{
var json = await File.ReadAllTextAsync(statePath, token);
var persistedState = JsonSerializer.Deserialize<PersistedEmbeddingState>(json, this.jsonOptions);
this.manifests = persistedState?.DataSources ?? new Dictionary<string, DataSourceEmbeddingManifest>(StringComparer.OrdinalIgnoreCase);
}
this.stateLoaded = true;
}
finally
{
this.stateLock.Release();
}
}
private async Task<DataSourceEmbeddingManifest> GetManifestAsync(string dataSourceId, CancellationToken token)
{
await this.EnsureStateLoadedAsync(token);
if (this.manifests.TryGetValue(dataSourceId, out var manifest))
return manifest;
manifest = new DataSourceEmbeddingManifest();
this.manifests[dataSourceId] = manifest;
return manifest;
}
private async Task SaveStateAsync(CancellationToken token)
{
var statePath = this.GetStatePath();
if (string.IsNullOrWhiteSpace(statePath))
return;
var directory = Path.GetDirectoryName(statePath);
if (!string.IsNullOrWhiteSpace(directory))
Directory.CreateDirectory(directory);
var persistedState = new PersistedEmbeddingState
{
DataSources = this.manifests
};
var json = JsonSerializer.Serialize(persistedState, this.jsonOptions);
await File.WriteAllTextAsync(statePath, json, token);
}
private async Task ResetPersistedStateAsync(string dataSourcename, string dataSourceId, VectorStoreClient? vectorStore, CancellationToken token)
{
await this.EnsureStateLoadedAsync(token);
this.manifests.Remove(dataSourceId);
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourcename, dataSourceId), vectorStore, token);
await this.SaveStateAsync(token);
await embeddingState.DeleteDataSourceAsync(dataSourceId, token);
logger.LogInformation("Reset persisted embedding state for data source '{DataSourceId}'.", dataSourceId);
}
private string GetStatePath()
{
if (string.IsNullOrWhiteSpace(SettingsManager.ConfigDirectory))
return string.Empty;
return Path.Combine(SettingsManager.ConfigDirectory, STATE_FILENAME);
}
}
@@ -6,6 +6,7 @@ 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;
@@ -22,13 +23,9 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
private readonly Channel<string> queue = Channel.CreateUnbounded<string>();
private readonly ConcurrentDictionary<string, byte> queuedIds = new(StringComparer.OrdinalIgnoreCase);
private readonly ConcurrentDictionary<string, DataSourceEmbeddingStatus> statuses = new(StringComparer.OrdinalIgnoreCase);
private readonly SemaphoreSlim stateLock = new(1, 1);
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataSourceEmbeddingService).Namespace, nameof(DataSourceEmbeddingService));
private Dictionary<string, DataSourceEmbeddingManifest> manifests = new(StringComparer.OrdinalIgnoreCase);
private bool stateLoaded;
public IReadOnlyList<DataSourceEmbeddingStatus> GetStatuses()
{
return this.statuses.Values
@@ -114,7 +111,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
this.RemoveWatcher(dataSource.Id);
this.statuses.TryRemove(dataSource.Id, out _);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, null, CancellationToken.None);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, null, null, CancellationToken.None);
this.PublishStatusChanged();
}
@@ -152,16 +149,15 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
public override void Dispose()
{
this.DisposeWatchers();
this.stateLock.Dispose();
base.Dispose();
}
private async Task ProcessDataSourceAsync(IDataSource dataSource, CancellationToken token)
{
await this.EnsureStateLoadedAsync(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)
{
@@ -174,6 +170,17 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
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."));
@@ -189,7 +196,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Id);
var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id);
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, token);
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;
@@ -202,8 +209,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
inputFiles.FailedFiles,
collectionName);
await this.RemoveMissingFileEmbeddingsAsync(vectorStore, dataSource, collectionName, manifest, indexedFiles, token);
await this.SaveStateAsync(token);
await this.RemoveMissingFileEmbeddingsAsync(vectorStore, embeddingState, dataSource, collectionName, manifest, indexedFiles, token);
this.UpsertStatus(this.CreateStatus(
dataSource,
@@ -223,7 +229,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
{
token.ThrowIfCancellationRequested();
var fingerprint = this.BuildFingerprint(file);
var fingerprint = await this.BuildFingerprintAsync(file, token);
if (manifest.Files.TryGetValue(file.FullName, out var existingRecord) &&
string.Equals(existingRecord.Fingerprint, fingerprint, StringComparison.Ordinal))
{
@@ -249,14 +255,27 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
skippedFiles + completedFiles + 1,
totalFiles);
var startedAtUtc = DateTime.UtcNow;
var chunkCount = await this.IndexOneFileAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, token);
manifest.Files[file.FullName] = new EmbeddedFileRecord(
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,
DateTime.UtcNow,
embeddedAtUtc,
chunkCount);
await this.SaveStateAsync(token);
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}.",
@@ -272,7 +291,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
lastError = exception.Message;
manifest.Files.Remove(file.FullName);
await this.DeleteFilePointsAsync(vectorStore, collectionName, file.FullName, token);
await this.SaveStateAsync(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));
@@ -290,6 +309,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
}
private async Task<int> IndexOneFileAsync(
EmbeddingStateClient embeddingState,
VectorStoreClient vectorStore,
IDataSource dataSource,
FileInfo file,
@@ -315,11 +335,11 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
totalChunkCount++;
if (batch.Count >= EMBEDDING_BATCH_SIZE)
await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
await this.FlushBatchAsync(embeddingState, vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
}
if (batch.Count > 0)
await this.FlushBatchAsync(vectorStore, dataSource, file, fingerprint, embeddingProvider, provider, manifest, collectionName, batch, token);
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.");
@@ -335,6 +355,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
}
private async Task FlushBatchAsync(
EmbeddingStateClient embeddingState,
VectorStoreClient vectorStore,
IDataSource dataSource,
FileInfo file,
@@ -369,7 +390,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
{
manifest.VectorSize = vectorSize;
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token);
await this.SaveStateAsync(token);
await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token);
logger.LogInformation(
"Created embedding collection '{CollectionName}' with vector size {VectorSize} for data source '{DataSourceName}' ({DataSourceId}).",
collectionName,
@@ -484,10 +505,16 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
return embeddingProvider != default && embeddingProvider.UsedLLMProvider is not LLMProviders.NONE;
}
private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(IDataSource dataSource, EmbeddingProvider embeddingProvider, string collectionName, VectorStoreClient vectorStore, CancellationToken token)
private async Task<DataSourceEmbeddingManifest> EnsureCompatibleManifestAsync(
IDataSource dataSource,
EmbeddingProvider embeddingProvider,
string collectionName,
VectorStoreClient vectorStore,
EmbeddingStateClient embeddingState,
CancellationToken token)
{
var embeddingSignature = this.BuildEmbeddingSignature(embeddingProvider);
var manifest = await this.GetManifestAsync(dataSource.Id, token);
var manifest = await embeddingState.GetManifestAsync(dataSource.Id, token);
if (!string.Equals(manifest.EmbeddingSignature, embeddingSignature, StringComparison.Ordinal))
{
@@ -496,8 +523,8 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Name,
dataSource.Id,
collectionName);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, vectorStore, token);
manifest = await this.GetManifestAsync(dataSource.Id, token);
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) ||
@@ -505,14 +532,23 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
{
manifest.EmbeddingProviderId = embeddingProvider.Id;
manifest.EmbeddingSignature = embeddingSignature;
await this.SaveStateAsync(token);
}
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,
@@ -526,6 +562,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
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}).",