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/Chat/ProvenanceCleanupTests.cs

@@ -0,0 +1,247 @@
+using AIStudio.Chat;
+using AIStudio.Tests.Tools.Databases;
+using AIStudio.Tools.Databases.ProvenanceStore;
+using AIStudio.Tools.RAG;
+
+namespace AIStudio.Tests.Chat;
+
+/// <summary>
+/// Checks that the provenance store keeps what stored chats refer to, and nothing else.
+/// </summary>
+/// <remarks>
+/// What a chat received belongs to it: kept as long as the chat refers to it, removed once it no
+/// longer does. The cleanup at the start must never take what this run stored since, because a
+/// chat in memory may refer to it before its file does.
+/// </remarks>
+[TestFixture]
+public sealed class ProvenanceCleanupTests
+{
+    private static readonly CancellationToken TOKEN = CancellationToken.None;
+
+    private static readonly Func<Guid, Guid, Task<ChatThread?>> UNREADABLE = (_, _) => Task.FromResult<ChatThread?>(null);
+
+    [Test]
+    public void AChatRefersToItsPassagesItsAttachmentsAndItsDraft()
+    {
+        var passageId = Guid.NewGuid();
+        var attachmentId = Guid.NewGuid();
+        var draftAttachmentId = Guid.NewGuid();
+        var previewedTextId = Guid.NewGuid();
+        var chat = ChatWith(passageId, attachmentId);
+
+        // A message being edited brings its attachment along, and a preview read another one:
+        chat.Draft.AddAttachments([new FileAttachment(FileAttachmentType.DOCUMENT, "draft.txt", "/tmp/draft.txt", 12) { ProvenanceId = draftAttachmentId }]);
+        chat.Draft.NotePreviewedText(new("/tmp/notes.txt", previewedTextId, new(DateTime.UtcNow, 12)));
+
+        Assert.That(chat.CollectProvenanceIds(), Is.EquivalentTo(new[] { passageId, attachmentId, draftAttachmentId, previewedTextId }), "The message being written is referred to as well, or the cleanup would take the texts it sends.");
+    }
+
+    [Test]
+    public async Task WhatNoChatRefersToIsRemoved()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var keptChatId = Guid.NewGuid();
+        var deletedChatId = Guid.NewGuid();
+        var referenced = await store.Client.StoreAsync(keptChatId, ProvenanceKind.DATA_SOURCE_PASSAGE, "still referred to", null, TOKEN);
+        var ofRemovedMessage = await store.Client.StoreAsync(keptChatId, ProvenanceKind.DATA_SOURCE_PASSAGE, "of a removed message", null, TOKEN);
+        var ofDeletedChat = await store.Client.StoreAsync(deletedChatId, ProvenanceKind.WEB_PAGE, "of a deleted chat", null, TOKEN);
+        var chat = ChatWith(referenced!.Value, attachmentId: null, keptChatId);
+
+        var result = await CleanUpAsync(store, Locations(keptChatId), LoaderOf(chat), DateTimeOffset.UtcNow.AddMinutes(1));
+        var entries = await store.Client.GetEntriesAsync([referenced.Value, ofRemovedMessage!.Value, ofDeletedChat!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(entries.Keys, Is.EquivalentTo(new[] { referenced.Value }));
+            Assert.That(result.Orphaned, Is.EqualTo(1));
+            Assert.That(result.Unreferenced, Is.EqualTo(1));
+        });
+    }
+
+    [Test]
+    public async Task WhatThisRunStoredStays()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var startedAt = DateTimeOffset.UtcNow.AddMinutes(-1);
+        var openChatId = Guid.NewGuid();
+        var newChatId = Guid.NewGuid();
+        var previewed = await store.Client.StoreAsync(openChatId, ProvenanceKind.ATTACHMENT, "previewed, not yet in the chat file", null, TOKEN);
+        var ofNewChat = await store.Client.StoreAsync(newChatId, ProvenanceKind.ATTACHMENT, "of a chat not stored yet", null, TOKEN);
+
+        await CleanUpAsync(store, Locations(openChatId), LoaderOf(ChatWith(passageId: null, attachmentId: null, openChatId)), startedAt);
+        var entries = await store.Client.GetEntriesAsync([previewed!.Value, ofNewChat!.Value], TOKEN);
+
+        Assert.That(entries.Keys, Is.EquivalentTo(new[] { previewed.Value, ofNewChat.Value }), "A chat in memory may refer to entries its file does not know yet.");
+    }
+
+    [Test]
+    public async Task WhatWasStoredForAChatWithoutAnIdIsRemoved()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var entry = await store.Client.StoreAsync(Guid.Empty, ProvenanceKind.WEB_PAGE, "of an assistant's conversation without an id", null, TOKEN);
+
+        var result = await CleanUpAsync(store, Locations(Guid.NewGuid()), UNREADABLE, DateTimeOffset.UtcNow.AddMinutes(1));
+        var entries = await store.Client.GetEntriesAsync([entry!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(entries, Is.Empty, "No chat on the disk has no id, so nothing would ever remove such an entry.");
+            Assert.That(result.Orphaned, Is.EqualTo(1));
+        });
+    }
+
+    [Test]
+    public async Task AChatWhichCannotBeReadKeepsEverything()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var chatId = Guid.NewGuid();
+        var entry = await store.Client.StoreAsync(chatId, ProvenanceKind.MAIL, "a mail", null, TOKEN);
+
+        var result = await CleanUpAsync(store, Locations(chatId), UNREADABLE, DateTimeOffset.UtcNow.AddMinutes(1));
+        var entries = await store.Client.GetEntriesAsync([entry!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(entries.Keys, Is.EquivalentTo(new[] { entry.Value }), "Nothing tells what an unreadable chat refers to, e.g., one written by a newer version.");
+            Assert.That(result.Unreadable, Is.EqualTo(1), "The next start has to read it again.");
+        });
+    }
+
+    [Test]
+    public async Task AChatWhichDidNotChangeSinceTheLastCleanupIsNotRead()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var chatId = Guid.NewGuid();
+        var entry = await store.Client.StoreAsync(chatId, ProvenanceKind.ATTACHMENT, "kept by the last cleanup", null, TOKEN);
+        var lastCleanup = DateTimeOffset.UtcNow.AddMinutes(1);
+        var reads = new List<Guid>();
+
+        var result = await CleanUpAsync(store, Locations(chatId), LoaderOf(ChatWith(passageId: null, attachmentId: null, chatId), reads), lastCleanup.AddMinutes(1), lastCleanup, writtenAt: lastCleanup.AddHours(-1));
+        var entries = await store.Client.GetEntriesAsync([entry!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(reads, Is.Empty, "Reading every chat at every start costs much on a small computer.");
+            Assert.That(result.Skipped, Is.EqualTo(1));
+            Assert.That(entries.Keys, Is.EquivalentTo(new[] { entry.Value }), "Skipping a chat never removes anything.");
+        });
+    }
+
+    [Test]
+    public async Task AChatWhichReceivedAnEntrySinceTheLastCleanupIsRead()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var lastCleanup = DateTimeOffset.UtcNow.AddMinutes(-1);
+        var chatId = Guid.NewGuid();
+        var previewed = await store.Client.StoreAsync(chatId, ProvenanceKind.ATTACHMENT, "previewed, but never sent", null, TOKEN);
+        var reads = new List<Guid>();
+
+        await CleanUpAsync(store, Locations(chatId), LoaderOf(ChatWith(passageId: null, attachmentId: null, chatId), reads), DateTimeOffset.UtcNow.AddMinutes(1), lastCleanup, writtenAt: lastCleanup.AddHours(-1));
+        var entries = await store.Client.GetEntriesAsync([previewed!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(reads, Is.EquivalentTo(new[] { chatId }), "The last cleanup had to leave alone what was stored after it began.");
+            Assert.That(entries, Is.Empty);
+        });
+    }
+
+    [TestCase(60, TestName = "AChatWrittenSinceTheLastCleanupIsRead")]
+    [TestCase(-2, TestName = "AChatWrittenShortlyBeforeTheLastCleanupIsRead")]
+    public async Task AChatWrittenAroundTheLastCleanupIsRead(int secondsAfterTheLastCleanup)
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var chatId = Guid.NewGuid();
+        var ofRemovedMessage = await store.Client.StoreAsync(chatId, ProvenanceKind.DATA_SOURCE_PASSAGE, "of a message removed after the last cleanup", null, TOKEN);
+        var lastCleanup = DateTimeOffset.UtcNow.AddMinutes(1);
+        var reads = new List<Guid>();
+
+        // Shortly before counts as well: the time of a file may be coarse, e.g., two seconds on FAT.
+        await CleanUpAsync(store, Locations(chatId), LoaderOf(ChatWith(passageId: null, attachmentId: null, chatId), reads), lastCleanup.AddMinutes(1), lastCleanup, writtenAt: lastCleanup.AddSeconds(secondsAfterTheLastCleanup));
+        var entries = await store.Client.GetEntriesAsync([ofRemovedMessage!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(reads, Is.EquivalentTo(new[] { chatId }));
+            Assert.That(entries, Is.Empty);
+        });
+    }
+
+    [Test]
+    public async Task AChatWhoseTimeOfWritingIsNotKnownIsRead()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var chatId = Guid.NewGuid();
+        await store.Client.StoreAsync(chatId, ProvenanceKind.MAIL, "a mail", null, TOKEN);
+        var lastCleanup = DateTimeOffset.UtcNow.AddMinutes(1);
+        var reads = new List<Guid>();
+
+        await CleanUpAsync(store, Locations(chatId), LoaderOf(ChatWith(passageId: null, attachmentId: null, chatId), reads), lastCleanup.AddMinutes(1), lastCleanup, writtenAt: null);
+
+        Assert.That(reads, Is.EquivalentTo(new[] { chatId }), "Only a chat known to be unchanged may be skipped.");
+    }
+
+    [Test]
+    public async Task AChatWhichNoLongerExistsIsRemovedAfterEveryCleanup()
+    {
+        await using var store = await TemporaryProvenanceStore.CreateAsync();
+        var entry = await store.Client.StoreAsync(Guid.NewGuid(), ProvenanceKind.WEB_PAGE, "of a chat deleted while the store was not available", null, TOKEN);
+        var lastCleanup = DateTimeOffset.UtcNow.AddMinutes(1);
+        var reads = new List<Guid>();
+
+        var result = await CleanUpAsync(store, Locations(), LoaderOf(ChatWith(passageId: null, attachmentId: null), reads), lastCleanup.AddMinutes(1), lastCleanup, writtenAt: lastCleanup.AddHours(-1));
+        var entries = await store.Client.GetEntriesAsync([entry!.Value], TOKEN);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(entries, Is.Empty, "Telling that a chat no longer exists costs no reading.");
+            Assert.That(result.Orphaned, Is.EqualTo(1));
+            Assert.That(reads, Is.Empty);
+        });
+    }
+
+    private static ChatThread ChatWith(Guid? passageId, Guid? attachmentId, Guid? chatId = null) => new()
+    {
+        ChatId = chatId ?? Guid.NewGuid(),
+        Blocks =
+        [
+            new ContentBlock
+            {
+                ContentType = ContentType.TEXT,
+                Role = ChatRole.USER,
+                Content = new ContentText
+                {
+                    Text = "Summarize the minutes.",
+                    FileAttachments = attachmentId is null ? [] : [new FileAttachment(FileAttachmentType.DOCUMENT, "minutes.txt", "/tmp/minutes.txt", 24) { ProvenanceId = attachmentId }],
+                },
+            },
+            new ContentBlock
+            {
+                ContentType = ContentType.TEXT,
+                Role = ChatRole.AI,
+                Content = new ContentText
+                {
+                    Text = "The board met on Monday.",
+                    SourcePassages = passageId is null ? [] : [new SourcePassage { SourceUrl = "file:///C:/docs/minutes.pdf", ProvenanceId = passageId }],
+                },
+            },
+        ],
+    };
+
+    private static IReadOnlyDictionary<Guid, Guid> Locations(params Guid[] chatIds) => chatIds.ToDictionary(chatId => chatId, _ => Guid.Empty);
+
+    /// <summary>
+    /// Reads the given chat, and notes every chat it was asked for.
+    /// </summary>
+    private static Func<Guid, Guid, Task<ChatThread?>> LoaderOf(ChatThread chat, List<Guid>? reads = null) => (_, chatId) =>
+    {
+        reads?.Add(chatId);
+        return Task.FromResult<ChatThread?>(chatId == chat.ChatId ? chat : null);
+    };
+
+    /// <summary>
+    /// Runs the cleanup on chats which were all written at the same moment.
+    /// </summary>
+    private static Task<ProvenanceCleanupResult> CleanUpAsync(TemporaryProvenanceStore store, IReadOnlyDictionary<Guid, Guid> chatLocations, Func<Guid, Guid, Task<ChatThread?>> loadChat, DateTimeOffset createdBefore, DateTimeOffset? lastCleanup = null, DateTimeOffset? writtenAt = null) =>
+        ProvenanceCleanup.CleanUpAsync(store.Client, chatLocations, loadChat, (_, _) => writtenAt, createdBefore, lastCleanup, TOKEN);
+}