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