using System.Globalization; using AIStudio.Settings; using AIStudio.Tools.PluginSystem; using AIStudio.Tools.Services; using Microsoft.Data.Sqlite; namespace AIStudio.Tools.Databases.EmbeddingState; public sealed class SqliteEmbeddingStateClientImplementation( string name, string databasePath, string basePath, string version) : EmbeddingStateClient(name, basePath) { private const string DATABASE_NAME = "SQLite"; private const string DATABASE_FILENAME = "rag-embedding-state.sqlite3"; private readonly string databasePath = databasePath; private readonly string connectionString = new SqliteConnectionStringBuilder { DataSource = databasePath, Mode = SqliteOpenMode.ReadWriteCreate, Cache = SqliteCacheMode.Shared, ForeignKeys = true, DefaultTimeout = 30, }.ToString(); private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(SqliteEmbeddingStateClientImplementation).Namespace, nameof(SqliteEmbeddingStateClientImplementation)); public override string CacheKey => $"{this.Name}:{this.databasePath}:{version}"; public static async Task CreateAsync( ILogger logger, ILogger databaseClientLogger, CancellationToken cancellationToken) { if (string.IsNullOrWhiteSpace(SettingsManager.DataDirectory)) return CreateNoEmbeddingStateClient(DATABASE_NAME, "The application data directory is not available yet.", DatabaseClientStatus.STARTING, databaseClientLogger); try { SQLitePCL.Batteries_V2.Init(); var basePath = Path.Combine(SettingsManager.DataDirectory, "databases", "sqlite"); Directory.CreateDirectory(basePath); var databasePath = Path.Combine(basePath, DATABASE_FILENAME); var client = new SqliteEmbeddingStateClientImplementation(DATABASE_NAME, databasePath, basePath, string.Empty); await client.InitializeAsync(cancellationToken); var version = await client.GetSqliteVersionAsync(cancellationToken); client = new SqliteEmbeddingStateClientImplementation(DATABASE_NAME, databasePath, basePath, version); client.SetLogger(databaseClientLogger); return client; } catch (Exception exception) { logger.LogWarning(exception, "{DatabaseName} is not available. Indexed file fingerprints are disabled.", DATABASE_NAME); return CreateNoEmbeddingStateClient(DATABASE_NAME, exception.Message, DatabaseClientStatus.UNAVAILABLE, databaseClientLogger); } } public override async IAsyncEnumerable<(string Label, string Value)> GetDisplayInfo() { yield return (TB("Reported version"), version); yield return (TB("Database path"), this.databasePath); yield return (TB("Storage size"), this.GetStorageSize()); yield return (TB("Indexed data sources"), (await this.CountAsync("data_sources", CancellationToken.None)).ToString(CultureInfo.InvariantCulture)); yield return (TB("Indexed files"), (await this.CountAsync("embedded_files", CancellationToken.None)).ToString(CultureInfo.InvariantCulture)); } public override async Task GetManifestAsync(string dataSourceId, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); var manifest = new DataSourceEmbeddingManifest(); await using (var command = CreateCommand(connection, """ SELECT embedding_provider_id, embedding_signature, source_hash, vector_size FROM data_sources WHERE data_source_id = $dataSourceId """)) { command.Parameters.AddWithValue("$dataSourceId", dataSourceId); await using var reader = await command.ExecuteReaderAsync(token); if (await reader.ReadAsync(token)) { manifest.EmbeddingProviderId = reader.GetString(0); manifest.EmbeddingSignature = reader.GetString(1); manifest.SourceHash = reader.GetString(2); manifest.VectorSize = reader.GetInt32(3); } } await using (var command = CreateCommand(connection, """ SELECT file_path, fingerprint, file_size, last_write_utc, embedded_at_utc, chunk_count FROM embedded_files WHERE data_source_id = $dataSourceId """)) { command.Parameters.AddWithValue("$dataSourceId", dataSourceId); await using var reader = await command.ExecuteReaderAsync(token); while (await reader.ReadAsync(token)) { manifest.Files[reader.GetString(0)] = new EmbeddedFileRecord( reader.GetString(1), reader.GetInt64(2), ParseUtc(reader.GetString(3)), ParseUtc(reader.GetString(4)), reader.GetInt32(5)); } } return manifest; } public override async Task UpsertDataSourceAsync( string dataSourceId, string dataSourceName, string dataSourceType, string embeddingProviderId, string embeddingSignature, string sourceHash, int vectorSize, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ INSERT INTO data_sources ( data_source_id, data_source_name, data_source_type, embedding_provider_id, embedding_signature, source_hash, vector_size, updated_at_utc) VALUES ( $dataSourceId, $dataSourceName, $dataSourceType, $embeddingProviderId, $embeddingSignature, $sourceHash, $vectorSize, $updatedAtUtc) ON CONFLICT(data_source_id) DO UPDATE SET data_source_name = excluded.data_source_name, data_source_type = excluded.data_source_type, embedding_provider_id = excluded.embedding_provider_id, embedding_signature = excluded.embedding_signature, source_hash = excluded.source_hash, vector_size = excluded.vector_size, updated_at_utc = excluded.updated_at_utc """, token, ("$dataSourceId", dataSourceId), ("$dataSourceName", dataSourceName), ("$dataSourceType", dataSourceType), ("$embeddingProviderId", embeddingProviderId), ("$embeddingSignature", embeddingSignature), ("$sourceHash", sourceHash), ("$vectorSize", vectorSize), ("$updatedAtUtc", ToUtcText(DateTime.UtcNow))); } public override async Task UpdateVectorSizeAsync(string dataSourceId, int vectorSize, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ UPDATE data_sources SET vector_size = $vectorSize, updated_at_utc = $updatedAtUtc WHERE data_source_id = $dataSourceId """, token, ("$dataSourceId", dataSourceId), ("$vectorSize", vectorSize), ("$updatedAtUtc", ToUtcText(DateTime.UtcNow))); } public override async Task UpdateDataSourceHashAsync(string dataSourceId, string sourceHash, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ UPDATE data_sources SET source_hash = $sourceHash, updated_at_utc = $updatedAtUtc WHERE data_source_id = $dataSourceId """, token, ("$dataSourceId", dataSourceId), ("$sourceHash", sourceHash), ("$updatedAtUtc", ToUtcText(DateTime.UtcNow))); } public override async Task UpsertFileAsync(string dataSourceId, EmbeddingStateFile file, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ INSERT INTO embedded_files ( data_source_id, file_path, file_name, relative_path, fingerprint, file_size, last_write_utc, embedded_at_utc, chunk_count) VALUES ( $dataSourceId, $filePath, $fileName, $relativePath, $fingerprint, $fileSize, $lastWriteUtc, $embeddedAtUtc, $chunkCount) ON CONFLICT(data_source_id, file_path) DO UPDATE SET file_name = excluded.file_name, relative_path = excluded.relative_path, fingerprint = excluded.fingerprint, file_size = excluded.file_size, last_write_utc = excluded.last_write_utc, embedded_at_utc = excluded.embedded_at_utc, chunk_count = excluded.chunk_count """, token, ("$dataSourceId", dataSourceId), ("$filePath", file.FilePath), ("$fileName", file.FileName), ("$relativePath", file.RelativePath), ("$fingerprint", file.Fingerprint), ("$fileSize", file.FileSize), ("$lastWriteUtc", ToUtcText(file.LastWriteUtc)), ("$embeddedAtUtc", ToUtcText(file.EmbeddedAtUtc)), ("$chunkCount", file.ChunkCount)); } public override async Task DeleteFileAsync(string dataSourceId, string filePath, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ DELETE FROM embedded_files WHERE data_source_id = $dataSourceId AND file_path = $filePath """, token, ("$dataSourceId", dataSourceId), ("$filePath", filePath)); } public override async Task DeleteDataSourceAsync(string dataSourceId, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ DELETE FROM data_sources WHERE data_source_id = $dataSourceId """, token, ("$dataSourceId", dataSourceId)); } public override void Dispose() { } private async Task InitializeAsync(CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await ExecuteNonQueryAsync(connection, """ PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON; CREATE TABLE IF NOT EXISTS data_sources ( data_source_id TEXT PRIMARY KEY, data_source_name TEXT NOT NULL, data_source_type TEXT NOT NULL, embedding_provider_id TEXT NOT NULL, embedding_signature TEXT NOT NULL, source_hash TEXT NOT NULL DEFAULT '', vector_size INTEGER NOT NULL DEFAULT 0, updated_at_utc TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS embedded_files ( data_source_id TEXT NOT NULL, file_path TEXT COLLATE NOCASE NOT NULL, file_name TEXT NOT NULL, relative_path TEXT NOT NULL, fingerprint TEXT NOT NULL, file_size INTEGER NOT NULL, last_write_utc TEXT NOT NULL, embedded_at_utc TEXT NOT NULL, chunk_count INTEGER NOT NULL, PRIMARY KEY (data_source_id, file_path), FOREIGN KEY (data_source_id) REFERENCES data_sources(data_source_id) ON DELETE CASCADE ); CREATE INDEX IF NOT EXISTS idx_embedded_files_data_source ON embedded_files(data_source_id); """, token); await EnsureColumnAsync(connection, "data_sources", "source_hash", "TEXT NOT NULL DEFAULT ''", token); } private async Task GetSqliteVersionAsync(CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await using var command = CreateCommand(connection, "SELECT sqlite_version()"); var versionObject = await command.ExecuteScalarAsync(token); return Convert.ToString(versionObject, CultureInfo.InvariantCulture) ?? string.Empty; } private async Task CountAsync(string tableName, CancellationToken token) { await using var connection = await this.OpenConnectionAsync(token); await using var command = CreateCommand(connection, $"SELECT COUNT(*) FROM {tableName}"); var countObject = await command.ExecuteScalarAsync(token); return Convert.ToInt64(countObject, CultureInfo.InvariantCulture); } private static async Task EnsureColumnAsync(SqliteConnection connection, string tableName, string columnName, string columnDefinition, CancellationToken token) { await using (var command = CreateCommand(connection, $"PRAGMA table_info({tableName})")) { await using var reader = await command.ExecuteReaderAsync(token); while (await reader.ReadAsync(token)) { if (string.Equals(reader.GetString(1), columnName, StringComparison.OrdinalIgnoreCase)) return; } } await ExecuteNonQueryAsync(connection, $"ALTER TABLE {tableName} ADD COLUMN {columnName} {columnDefinition}", token); } private async Task OpenConnectionAsync(CancellationToken token) { var connection = new SqliteConnection(this.connectionString); await connection.OpenAsync(token); return connection; } private static SqliteCommand CreateCommand(SqliteConnection connection, string commandText) { var command = connection.CreateCommand(); command.CommandText = commandText; return command; } private static async Task ExecuteNonQueryAsync( SqliteConnection connection, string commandText, CancellationToken token, params (string Name, object? Value)[] parameters) { await using var command = CreateCommand(connection, commandText); foreach (var (name, value) in parameters) command.Parameters.AddWithValue(name, value ?? DBNull.Value); await command.ExecuteNonQueryAsync(token); } private static string ToUtcText(DateTime dateTime) { var utc = dateTime.Kind is DateTimeKind.Utc ? dateTime : dateTime.ToUniversalTime(); return utc.ToString("O", CultureInfo.InvariantCulture); } private static DateTime ParseUtc(string value) { return DateTime.TryParse(value, CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind, out var dateTime) ? dateTime.ToUniversalTime() : DateTime.UnixEpoch; } private static NoEmbeddingStateClient CreateNoEmbeddingStateClient(string name, string? unavailableReason, DatabaseClientStatus status, ILogger databaseClientLogger) { var client = new NoEmbeddingStateClient(name, unavailableReason, status); client.SetLogger(databaseClientLogger); return client; } }