Added a way to read what the AI received from your sources (#1100)

Co-authored-by: Thorsten Sommer <SommerEngineering@users.noreply.github.com>

app/Tests/Tools/DataBackup/BackupChatImportTests.cs

@@ -4,10 +4,14 @@ using System.Text.Json;
 
 using AIStudio.Chat;
 using AIStudio.Settings;
+using AIStudio.Tests.Tools.Databases;
 using AIStudio.Tools;
 using AIStudio.Tools.ChatArchive;
+using AIStudio.Tools.Databases.ProvenanceStore;
 using AIStudio.Tools.DataBackup;
 
+using Microsoft.EntityFrameworkCore;
+
 namespace AIStudio.Tests.Tools.DataBackup;
 
 /// <summary>
@@ -61,16 +65,19 @@ public sealed class BackupChatImportTests
     [Test]
     public async Task TheChatsOfABackupArriveWithTheirTranscripts()
     {
+        await using var provenance = await TemporaryProvenanceStore.CreateAsync();
         var storedChat = await this.WriteBackupWithAChatAsync();
         await UseDataDirectoryAsync(this.TargetDataDirectory);
 
         var preview = await ChatArchiveImporter.ReadPreviewAsync(this.ArchivePath, CancellationToken.None);
-        var result = await ChatArchiveImporter.ImportAsync(this.ArchivePath, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
+        var result = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
         var chat = await WorkspaceBehaviour.LoadChatAsync(new(storedChat.WorkspaceId, storedChat.ChatId));
         var transcript = ((ContentText)chat!.Blocks.Single().Content!).FileAttachments.Single();
+        var chatsWithEntries = await provenance.Client.GetChatIdsAsync(CancellationToken.None);
 
         Assert.Multiple(() =>
         {
+            Assert.That(chatsWithEntries, Is.Empty, "A backup without a provenance store, e.g., of an older version, brings its chats all the same.");
             Assert.That(preview.Success, Is.True, preview.Issue);
             Assert.That(preview.Manifest.Workspaces.Select(workspace => workspace.Name), Is.EqualTo(new[] { "Projects" }));
             Assert.That(preview.Manifest.Workspaces.Single().Chats.Select(listedChat => listedChat.Name), Is.EqualTo(new[] { "Meeting notes" }));
@@ -84,11 +91,12 @@ public sealed class BackupChatImportTests
     [Test]
     public async Task AChatWhichExistsAlreadyIsKept()
     {
+        await using var provenance = await TemporaryProvenanceStore.CreateAsync();
         await this.WriteBackupWithAChatAsync();
         await UseDataDirectoryAsync(this.TargetDataDirectory);
 
-        await ChatArchiveImporter.ImportAsync(this.ArchivePath, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
-        var secondResult = await ChatArchiveImporter.ImportAsync(this.ArchivePath, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
+        await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
+        var secondResult = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
 
         Assert.Multiple(() =>
         {
@@ -97,6 +105,54 @@ public sealed class BackupChatImportTests
         });
     }
 
+    [Test]
+    public async Task TheChatsOfABackupArriveWithWhatTheyReceived()
+    {
+        await using var sourceProvenance = await TemporaryProvenanceStore.CreateAsync();
+        await using var provenance = await TemporaryProvenanceStore.CreateAsync();
+        var storedChat = await this.WriteBackupWithAChatAsync(sourceProvenance);
+        await UseDataDirectoryAsync(this.TargetDataDirectory);
+
+        var result = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
+        var entries = await provenance.Client.GetEntriesAsync([storedChat.ProvenanceId!.Value], CancellationToken.None);
+        var chatsWithEntries = await provenance.Client.GetChatIdsAsync(CancellationToken.None);
+        var stagingDirectory = Path.Combine(this.TargetDataDirectory, ".backup-staging");
+        var stagingLeftovers = Directory.Exists(stagingDirectory) ? Directory.GetFileSystemEntries(stagingDirectory) : Array.Empty<string>();
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(result.ImportedChats, Is.EqualTo(1));
+            Assert.That(entries[storedChat.ProvenanceId.Value].Content, Is.EqualTo("Hello and welcome."), "The chat refers to the entry by its id, which the import keeps.");
+            Assert.That(chatsWithEntries, Is.EqualTo(new[] { storedChat.ChatId }), "Another chat's entries of the source computer stay behind.");
+            Assert.That(stagingLeftovers, Is.Empty, "The unpacked store of the backup is removed afterward.");
+        });
+    }
+
+    [Test]
+    public async Task AChatAddedOnceMoreTakesItsEntriesAlong()
+    {
+        await using var sourceProvenance = await TemporaryProvenanceStore.CreateAsync();
+        await using var provenance = await TemporaryProvenanceStore.CreateAsync();
+        var storedChat = await this.WriteBackupWithAChatAsync(sourceProvenance);
+        await UseDataDirectoryAsync(this.TargetDataDirectory);
+
+        await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.IMPORT_AS_COPY, null, CancellationToken.None);
+        var secondResult = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.IMPORT_AS_COPY, null, CancellationToken.None);
+        var chatsWithEntries = await provenance.Client.GetChatIdsAsync(CancellationToken.None);
+
+        var entryKey = storedChat.ProvenanceId!.Value.ToString("D");
+        await using var context = provenance.CreateContext();
+        var copiesOfTheEntry = await context.Entries.CountAsync(entry => entry.Id == entryKey);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(secondResult.ImportedChats, Is.EqualTo(1));
+            Assert.That(chatsWithEntries, Has.Count.EqualTo(2), "The copy under its new id gets entries of its own, so deleting one of the two chats leaves the other whole.");
+            Assert.That(chatsWithEntries, Does.Contain(storedChat.ChatId));
+            Assert.That(copiesOfTheEntry, Is.EqualTo(2), "Both chats refer to the entry by the same id.");
+        });
+    }
+
     [Test]
     public async Task ABackupOfANewerFormatIsRefused()
     {
@@ -135,7 +191,8 @@ public sealed class BackupChatImportTests
     /// <summary>
     /// Stores a chat with a transcript in a workspace of the source computer, and backs up its data folder.
     /// </summary>
-    private async Task<StoredChat> WriteBackupWithAChatAsync()
+    /// <param name="sourceProvenance">The provenance store of the source computer, which then holds what the chat received, or null when the backup holds no store.</param>
+    private async Task<StoredChat> WriteBackupWithAChatAsync(TemporaryProvenanceStore? sourceProvenance = null)
     {
         await UseDataDirectoryAsync(this.SourceDataDirectory);
 
@@ -169,12 +226,31 @@ public sealed class BackupChatImportTests
 
         Assert.That(await WorkspaceBehaviour.StoreChatAsync(chat), Is.True, "The chat could not be stored.");
 
-        var sources = new DataBackupSources(this.SourceDataDirectory, Path.Combine(this.testDirectory, "config"), "{}", string.Empty, []);
+        //
+        // What the chat received, next to an entry of a chat the backup does not hold, copied the way
+        // a backup copies the store:
+        //
+        Guid? provenanceId = null;
+        Func<CancellationToken, Task<string>>? copyProvenanceStoreAsync = null;
+        if (sourceProvenance is not null)
+        {
+            provenanceId = await sourceProvenance.Client.StoreAsync(chatId, ProvenanceKind.ATTACHMENT, "Hello and welcome.", null, CancellationToken.None);
+            await sourceProvenance.Client.StoreAsync(Guid.NewGuid(), ProvenanceKind.WEB_PAGE, "A page of another chat.", null, CancellationToken.None);
+
+            var provenanceCopyDirectory = Path.Combine(this.testDirectory, "provenance-copy");
+            copyProvenanceStoreAsync = async token =>
+            {
+                await sourceProvenance.Client.BackupToDirectoryAsync(Path.Combine(provenanceCopyDirectory, "databases", "provenance"), token);
+                return provenanceCopyDirectory;
+            };
+        }
+
+        var sources = new DataBackupSources(this.SourceDataDirectory, Path.Combine(this.testDirectory, "config"), "{}", string.Empty, [], copyProvenanceStoreAsync);
         var manifest = new DataBackupManifest(DataBackupFormat.FORMAT_VERSION, "26.10.1", "V6", DateTimeOffset.UtcNow, this.SourceDataDirectory, string.Empty, false, []);
         await using (var archiveStream = File.Create(this.ArchivePath))
             await DataBackupWriter.WriteAsync(archiveStream, sources, manifest, null, CancellationToken.None);
 
-        return new(workspaceId, chatId);
+        return new(workspaceId, chatId, provenanceId);
     }
 
     /// <summary>
@@ -196,5 +272,8 @@ public sealed class BackupChatImportTests
     /// <summary>
     /// A chat as the source computer stored it.
     /// </summary>
-    private sealed record StoredChat(Guid WorkspaceId, Guid ChatId);
+    /// <param name="WorkspaceId">The workspace of the chat.</param>
+    /// <param name="ChatId">The chat.</param>
+    /// <param name="ProvenanceId">The entry of what the chat received, or null when the backup holds no provenance store.</param>
+    private sealed record StoredChat(Guid WorkspaceId, Guid ChatId, Guid? ProvenanceId = null);
 }

app/Tests/Tools/DataBackup/DataBackupWriterTests.cs

@@ -35,6 +35,8 @@ public sealed class DataBackupWriterTests
 
     private string IndexCopiesDirectory => Path.Combine(this.testDirectory, "index-copies");
 
+    private string ProvenanceCopyDirectory => Path.Combine(this.testDirectory, "provenance-copy");
+
     [SetUp]
     public void CreateFolders()
     {
@@ -54,6 +56,8 @@ public sealed class DataBackupWriterTests
         this.WriteDataFile("visualBriefings/briefing/manifest.json", "{}");
         this.WriteDataFile("media-staging/operation/transcript.md", "temporary");
         this.WriteDataFile("databases/sqlite/rag-index.sqlite3", "the index while AI Studio writes to it");
+        this.WriteDataFile("databases/provenance/provenance.sqlite3", "the provenance store while AI Studio writes to it");
+        this.WriteDataFile("databases/provenance/provenance.sqlite3-wal", "the write-ahead log of the provenance store");
         this.WriteDataFile("databases/vector_database/stores/store_1/segments/segment", "the store while AI Studio writes to it");
         this.WriteDataFile(MEDIA_TRANSCRIPT, "the transcript while AI Studio writes it");
 
@@ -61,6 +65,7 @@ public sealed class DataBackupWriterTests
         WriteFile(Path.Combine(this.IndexCopiesDirectory, "databases", "sqlite", "rag-index.sqlite3"), "the copy of the index");
         WriteFile(Path.Combine(this.IndexCopiesDirectory, "databases", "vector_database", "stores", "store_1", "segments", "segment"), "the copy of the store");
         WriteFile(Path.Combine(this.IndexCopiesDirectory, MEDIA_TRANSCRIPT), "the copy of the transcript");
+        WriteFile(Path.Combine(this.ProvenanceCopyDirectory, "databases", "provenance", "provenance.sqlite3"), "the copy of the provenance store");
     }
 
     [TearDown]
@@ -86,6 +91,7 @@ public sealed class DataBackupWriterTests
             "data/plugins/my-assistant/.hidden.lua",
             "data/tokenizers/openai-gpt-6/tokenizer.json",
             "data/visualBriefings/briefing/manifest.json",
+            "data/databases/provenance/provenance.sqlite3",
             "data/databases/sqlite/rag-index.sqlite3",
             "data/databases/vector_database/stores/store_1/segments/segment",
             $"data/{MEDIA_TRANSCRIPT}",
@@ -95,6 +101,7 @@ public sealed class DataBackupWriterTests
         Assert.Multiple(() =>
         {
             Assert.That(archive["config/settings.json"], Is.EqualTo(SETTINGS_JSON), "The settings come from memory, without the objects of configuration plugins.");
+            Assert.That(archive["data/databases/provenance/provenance.sqlite3"], Is.EqualTo("the copy of the provenance store"), "The chats write to the store at any time, so it comes from its copy, without the write-ahead log of the live one.");
             Assert.That(archive["data/databases/sqlite/rag-index.sqlite3"], Is.EqualTo("the copy of the index"));
             Assert.That(archive["data/databases/vector_database/stores/store_1/segments/segment"], Is.EqualTo("the copy of the store"));
             Assert.That(archive[$"data/{MEDIA_TRANSCRIPT}"], Is.EqualTo("the copy of the transcript"), "The transcripts are copied together with the index they belong to.");
@@ -106,10 +113,45 @@ public sealed class DataBackupWriterTests
     {
         var archive = await this.WriteBackupAsync(string.Empty);
 
-        Assert.That(archive.Keys.Where(name => name.StartsWith("data/databases/", StringComparison.Ordinal)), Is.Empty, "The indexes in the data folder change while AI Studio runs.");
+        Assert.That(archive.Keys.Where(name => name.StartsWith("data/databases/", StringComparison.Ordinal) && !name.StartsWith("data/databases/provenance/", StringComparison.Ordinal)), Is.Empty, "The indexes in the data folder change while AI Studio runs.");
         Assert.That(archive.Keys.Where(name => name.StartsWith("data/mediaTranscripts/", StringComparison.Ordinal)), Is.Empty, "The transcripts belong to the indexes and come only together with them.");
     }
 
+    [Test]
+    public async Task TheProvenanceStoreComesAlongAlsoWithoutTheIndexes()
+    {
+        var archive = await this.WriteBackupAsync(string.Empty);
+
+        Assert.That(archive["data/databases/provenance/provenance.sqlite3"], Is.EqualTo("the copy of the provenance store"), "The store belongs to the chats, which every backup holds, not to the indexes.");
+    }
+
+    [Test]
+    public async Task WithoutItsCopyTheProvenanceStoreComesFromTheDataFolder()
+    {
+        var archive = await this.WriteBackupAsync(this.IndexCopiesDirectory, _ => Task.FromResult(string.Empty));
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(archive["data/databases/provenance/provenance.sqlite3"], Is.EqualTo("the provenance store while AI Studio writes to it"), "A store AI Studio could not open is written by nothing, and the backup keeps whatever there is.");
+            Assert.That(archive["data/databases/provenance/provenance.sqlite3-wal"], Is.EqualTo("the write-ahead log of the provenance store"), "Its write-ahead log may hold what the store itself lacks yet.");
+        });
+    }
+
+    [Test]
+    public async Task TheProvenanceStoreIsCopiedOnceTheChatsAreWritten()
+    {
+        var chatFile = Path.Combine(this.DataDirectory, "workspaces", WORKSPACE_ID.ToString(), CHAT_ID.ToString(), "thread.json");
+
+        // A chat deleted while the store is copied was written into the backup before:
+        var archive = await this.WriteBackupAsync(this.IndexCopiesDirectory, _ =>
+        {
+            File.Delete(chatFile);
+            return Task.FromResult(this.ProvenanceCopyDirectory);
+        });
+
+        Assert.That(archive.Keys, Does.Contain($"data/workspaces/{WORKSPACE_ID}/{CHAT_ID}/thread.json"), "Copied after the chats, the store holds everything they refer to, also what an answer completed meanwhile received.");
+    }
+
     [Test]
     public async Task TheManifestComesLastAndListsEveryOtherFile()
     {
@@ -182,10 +224,10 @@ public sealed class DataBackupWriterTests
         });
     }
 
-    private async Task<Dictionary<string, string>> WriteBackupAsync(string indexCopiesDirectory)
+    private async Task<Dictionary<string, string>> WriteBackupAsync(string indexCopiesDirectory, Func<CancellationToken, Task<string>>? copyProvenanceStoreAsync = null)
     {
         await using var target = new MemoryStream();
-        await DataBackupWriter.WriteAsync(target, this.CreateSources(indexCopiesDirectory), CreateManifest(), null, CancellationToken.None);
+        await DataBackupWriter.WriteAsync(target, this.CreateSources(indexCopiesDirectory, copyProvenanceStoreAsync), CreateManifest(), null, CancellationToken.None);
         target.Position = 0;
 
         using var zip = new ZipArchive(target, ZipArchiveMode.Read);
@@ -199,7 +241,7 @@ public sealed class DataBackupWriterTests
         return contents;
     }
 
-    private DataBackupSources CreateSources(string indexCopiesDirectory) => new(this.DataDirectory, this.ConfigDirectory, SETTINGS_JSON, indexCopiesDirectory, ["plugins/company-assistant/"]);
+    private DataBackupSources CreateSources(string indexCopiesDirectory, Func<CancellationToken, Task<string>>? copyProvenanceStoreAsync = null) => new(this.DataDirectory, this.ConfigDirectory, SETTINGS_JSON, indexCopiesDirectory, ["plugins/company-assistant/"], copyProvenanceStoreAsync ?? (_ => Task.FromResult(this.ProvenanceCopyDirectory)));
 
     private static DataBackupManifest CreateManifest() => new(DataBackupFormat.FORMAT_VERSION, "26.10.1", "V6", DateTimeOffset.UtcNow, "/data", "/config", true, []);
 

app/Tests/Tools/DataBackup/DataRestorePreparerTests.cs

@@ -64,6 +64,23 @@ public sealed class DataRestorePreparerTests
         });
     }
 
+    [Test]
+    public async Task TheProvenanceStoreIsRestoredWithTheChatsAlsoWithoutTheIndexes()
+    {
+        await this.WriteBackupAsync("26.10.1", "V6", withProvenanceStore: true);
+
+        var result = await DataRestorePreparer.PrepareAsync(this.ArchivePath, this.Target, null, CancellationToken.None);
+        var dataPending = Path.Combine(this.Target.DataDirectory, ".restore-pending");
+        var plan = JsonSerializer.Deserialize<DataRestorePlan>(await File.ReadAllTextAsync(Path.Combine(dataPending, DataRestoreFormat.PLAN_FILE_NAME)), DataBackupFormat.JSON_OPTIONS)!;
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(result.Success, Is.True, result.Issue);
+            Assert.That(plan.Elements.Select(element => element.RelativePath), Does.Contain("databases/provenance/"), "The chats of the backup refer to the contents in its store, so both are replaced together, whether the backup holds the indexes or not.");
+            Assert.That(File.ReadAllText(Path.Combine(dataPending, "databases", "provenance", "provenance.sqlite3")), Is.EqualTo("the provenance store"));
+        });
+    }
+
     [Test]
     public async Task ADamagedBackupLeavesNothingBehind()
     {
@@ -122,7 +139,8 @@ public sealed class DataRestorePreparerTests
     /// <param name="appVersion">The version of AI Studio the manifest names.</param>
     /// <param name="settingsVersion">The version of the settings the manifest names.</param>
     /// <param name="damageThread">Whether to change the chat thread after the manifest noted its hash.</param>
-    private async Task WriteBackupAsync(string appVersion, string settingsVersion, bool damageThread = false)
+    /// <param name="withProvenanceStore">Whether the data folder holds a provenance store.</param>
+    private async Task WriteBackupAsync(string appVersion, string settingsVersion, bool damageThread = false, bool withProvenanceStore = false)
     {
         var sourceDataDirectory = Path.Combine(this.testDirectory, "source-data");
         var threadPath = Path.Combine(sourceDataDirectory, "workspaces", WORKSPACE_ID.ToString(), CHAT_ID.ToString(), "thread.json");
@@ -130,6 +148,8 @@ public sealed class DataRestorePreparerTests
         WriteFile(threadPath, "{}");
         File.SetLastWriteTime(threadPath, CHAT_LAST_CHANGED);
         WriteFile(Path.Combine(sourceDataDirectory, "plugins", "my-assistant", "plugin.lua"), "-- mine");
+        if (withProvenanceStore)
+            WriteFile(Path.Combine(sourceDataDirectory, "databases", "provenance", "provenance.sqlite3"), "the provenance store");
 
         var sources = new DataBackupSources(sourceDataDirectory, Path.Combine(this.testDirectory, "source-config"), """{ "Version": "V6" }""", string.Empty, []);
         var manifest = new DataBackupManifest(DataBackupFormat.FORMAT_VERSION, appVersion, settingsVersion, DateTimeOffset.UtcNow, sourceDataDirectory, string.Empty, false, []);