using AIStudio.Tools.Databases.IndexStore; using AIStudio.Tools.Services.Indexing; using Microsoft.Data.Sqlite; namespace AIStudio.Tests.Tools.Databases; /// /// An index store in a directory of its own, migrated the way the app migrates its own, and deleted /// again afterward. /// /// The directory the database lives in. /// The path of the database file. internal sealed class TemporaryIndexStore(string directory, string databasePath) : IAsyncDisposable { public SqliteIndexStoreClientImplementation Client { get; } = new("SQLite", databasePath, directory, string.Empty); public static async Task CreateAsync() { SQLitePCL.Batteries_V2.Init(); var directory = Path.Combine(Path.GetTempPath(), $"ai-studio-index-store-{Guid.NewGuid():N}"); Directory.CreateDirectory(directory); var databasePath = Path.Combine(directory, "rag-index.sqlite3"); await using (var context = new IndexStoreDbContext(IndexStoreDbContext.CreateOptions(databasePath))) await IndexStoreSchemaMigrator.MigrateAsync(context, CancellationToken.None); return new TemporaryIndexStore(directory, databasePath); } /// /// Opens the database directly, to set up or check what the client does not cover. /// public IndexStoreDbContext CreateContext() => new(IndexStoreDbContext.CreateOptions(databasePath)); /// /// Stores a mail the way the sync does: first its document with a single chunk, then the mail itself. /// /// The mailbox, which has to be stored together with the folders of the mail. /// A name for the mail, unique within the mailbox. It makes up the key of the document and its fingerprint. /// The subject, which the document keeps as its file name. /// The text of the single chunk. /// The mail. Its id is replaced by the id of its document. /// The id of the mail. public async Task StoreMailAsync(string dataSourceId, string name, string subject, string text, MailRecord mail) { var key = $"mail:{name}"; var mailId = IndexedDocumentIds.CreateParentId(dataSourceId, key); var now = DateTimeOffset.UtcNow; await this.Client.UpsertFileAsync(dataSourceId, new EmbeddingStateFile(mailId, key, subject, key, "mail", name, 2048, now, now, now, 1), CancellationToken.None); await this.Client.UpsertChunksAsync(dataSourceId, [new EmbeddingStateChunk(IndexedDocumentIds.CreateChunkId(dataSourceId, name, 0), mailId, null, 0, text, now)], CancellationToken.None); await this.Client.UpsertMailAsync(dataSourceId, mail with { MailId = mailId }, CancellationToken.None); return mailId; } public ValueTask DisposeAsync() { // The pooled connections hold the file open, which would keep it from being deleted on Windows: SqliteConnection.ClearAllPools(); Directory.Delete(directory, true); return ValueTask.CompletedTask; } }