using System.Text.Json;
using AIStudio.Chat;
using AIStudio.Tests.Tools.Databases;
using AIStudio.Tools;
using AIStudio.Tools.Databases.ProvenanceStore;
using AIStudio.Tools.RAG;
namespace AIStudio.Tests.Tools;
///
/// Checks how the passages of a complete answer move into the provenance store.
///
///
/// The chat keeps where a passage was found, and the store what the AI read there. Whatever the
/// store cannot take stays in the chat itself, so the user can read it either way.
///
[TestFixture]
public sealed class PassageProvenanceTests
{
private const string HANDBOOK_URL = "file:///C:/docs/handbook.pdf";
private static readonly CancellationToken TOKEN = CancellationToken.None;
[Test]
public async Task TheTextsOfAPassageMoveIntoTheStore()
{
await using var store = await TemporaryProvenanceStore.CreateAsync();
var chatId = Guid.NewGuid();
var answer = new ContentText { SourcePassages = [Passage($"{HANDBOOK_URL}#page=3", "The mixing console.", pageNumber: 3) with { SurroundingContent = ["Before.", "After."] }] };
var allStored = await answer.StoreReceivedContentAsync(store.Client, chatId);
var passage = answer.SourcePassages.Single();
var entries = await store.Client.GetEntriesAsync([passage.ProvenanceId!.Value], TOKEN);
var entry = entries[passage.ProvenanceId.Value];
Assert.Multiple(() =>
{
Assert.That(allStored, Is.True);
Assert.That(passage.Text, Is.Empty, "The chat keeps only the reference.");
Assert.That(passage.SurroundingContent, Is.Empty);
Assert.That(passage.PageNumber, Is.EqualTo(3), "Where the passage was found stays with the chat, so the list of sources draws without the store.");
Assert.That(entry.Content, Is.EqualTo("The mixing console."));
Assert.That(entry.SurroundingContent, Is.EqualTo(new[] { "Before.", "After." }));
Assert.That(entry.Kind, Is.EqualTo(ProvenanceKind.DATA_SOURCE_PASSAGE));
});
}
[Test]
public async Task EveryPassageGetsAnEntryOfItsOwn()
{
await using var store = await TemporaryProvenanceStore.CreateAsync();
var answer = new ContentText { SourcePassages = [Passage($"{HANDBOOK_URL}#page=3", "The mixing console."), Passage($"{HANDBOOK_URL}#page=4", "The amplifier.")] };
await answer.StoreReceivedContentAsync(store.Client, Guid.NewGuid());
var ids = answer.SourcePassages.Select(passage => passage.ProvenanceId).ToList();
Assert.Multiple(() =>
{
Assert.That(ids, Has.None.Null);
Assert.That(ids, Is.Unique);
});
}
[Test]
public async Task APassageWithoutTextsIsLeftAsItIs()
{
await using var store = await TemporaryProvenanceStore.CreateAsync();
var image = Passage($"{HANDBOOK_URL}#page=7", string.Empty) with { Category = RetrievalContentCategory.IMAGE, Type = RetrievalContentType.IMAGE_PHOTO };
var answer = new ContentText { SourcePassages = [image] };
var allStored = await answer.StoreReceivedContentAsync(store.Client, Guid.NewGuid());
var chatsWithEntries = await store.Client.GetChatIdsAsync(TOKEN);
Assert.Multiple(() =>
{
Assert.That(allStored, Is.True);
Assert.That(answer.SourcePassages.Single().ProvenanceId, Is.Null, "The AI was only told about the image, so there is no text to keep.");
Assert.That(chatsWithEntries, Is.Empty);
});
}
[Test]
public async Task AStoredPassageIsNotStoredTwice()
{
await using var store = await TemporaryProvenanceStore.CreateAsync();
var chatId = Guid.NewGuid();
var answer = new ContentText { SourcePassages = [Passage(HANDBOOK_URL, "The mixing console.")] };
await answer.StoreReceivedContentAsync(store.Client, chatId);
var firstId = answer.SourcePassages.Single().ProvenanceId;
await answer.StoreReceivedContentAsync(store.Client, chatId);
await using var context = store.CreateContext();
Assert.Multiple(() =>
{
Assert.That(answer.SourcePassages.Single().ProvenanceId, Is.EqualTo(firstId));
Assert.That(context.Entries.Count(), Is.EqualTo(1));
});
}
[Test]
public async Task WithoutAStoreThePassagesKeepTheirTexts()
{
var answer = new ContentText { SourcePassages = [Passage(HANDBOOK_URL, "The mixing console.") with { SurroundingContent = ["Before."] }] };
var allStored = await answer.StoreReceivedContentAsync(new NoProvenanceStoreClient("SQLite", "not available"), Guid.NewGuid());
var passage = answer.SourcePassages.Single();
Assert.Multiple(() =>
{
Assert.That(allStored, Is.False, "The caller tells the user, who should know why the chat file stays large.");
Assert.That(passage.ProvenanceId, Is.Null);
Assert.That(passage.Text, Is.EqualTo("The mixing console."), "The chat stays readable without the store.");
Assert.That(passage.SurroundingContent, Is.EqualTo(new[] { "Before." }));
});
}
[Test]
public void StoringNeverStopsTheAnswer()
{
// A passage without its text, as an ERI server sent one before its fields were checked:
var answer = new ContentText { SourcePassages = [Passage(HANDBOOK_URL, null!)] };
Assert.DoesNotThrowAsync(() => answer.StoreReceivedContentAsync(Guid.NewGuid()), "Completing the answer waits for this, and a throw left its chat running until the app restarted.");
}
[Test]
public async Task AStoredPassageIsSavedWithItsReferenceOnly()
{
await using var store = await TemporaryProvenanceStore.CreateAsync();
var answer = new ContentText { Text = "The console has 24 channels.", SourcePassages = [Passage($"{HANDBOOK_URL}#page=3", "The mixing console.", pageNumber: 3)] };
await answer.StoreReceivedContentAsync(store.Client, Guid.NewGuid());
var json = JsonSerializer.Serialize(answer, WorkspaceBehaviour.JSON_OPTIONS);
var loaded = JsonSerializer.Deserialize(json, WorkspaceBehaviour.JSON_OPTIONS)!;
using var document = JsonDocument.Parse(json);
var storedPassage = document.RootElement.GetProperty("source_passages")[0];
Assert.Multiple(() =>
{
Assert.That(storedPassage.GetProperty("provenance_id").GetGuid(), Is.EqualTo(answer.SourcePassages.Single().ProvenanceId));
Assert.That(storedPassage.GetProperty("text").GetString(), Is.Empty, "The text lies in the store.");
Assert.That(loaded.SourcePassages.Single().ProvenanceId, Is.EqualTo(answer.SourcePassages.Single().ProvenanceId));
});
}
[Test]
public void StoredPassagesOfOneSourceAreToldApartByTheirEntries()
{
var first = Passage(HANDBOOK_URL, string.Empty) with { ProvenanceId = Guid.NewGuid() };
var second = Passage(HANDBOOK_URL, string.Empty) with { ProvenanceId = Guid.NewGuid() };
List passages = [first];
passages.MergePassages([second, first]);
Assert.That(passages, Is.EqualTo(new[] { first, second }), "Without their texts, two stored passages of one page would otherwise count as one.");
}
[Test]
public void TheKindFollowsWhatThePassageIs()
{
Assert.Multiple(() =>
{
Assert.That(Passage(HANDBOOK_URL, "text").GetProvenanceKind(), Is.EqualTo(ProvenanceKind.DATA_SOURCE_PASSAGE));
Assert.That((Passage(HANDBOOK_URL, "text") with { Type = RetrievalContentType.TEXT_TRANSCRIPT }).GetProvenanceKind(), Is.EqualTo(ProvenanceKind.DATA_SOURCE_PASSAGE));
Assert.That((Passage(HANDBOOK_URL, "text") with { Type = RetrievalContentType.TEXT_EMAIL }).GetProvenanceKind(), Is.EqualTo(ProvenanceKind.MAIL), "A mail is a mail, whether a mailbox data source or a tool handed it over.");
Assert.That((Passage(HANDBOOK_URL, "text") with { Type = RetrievalContentType.TEXT_WEBSITE }).GetProvenanceKind(), Is.EqualTo(ProvenanceKind.WEB_PAGE));
});
}
private static SourcePassage Passage(string sourceUrl, string text, int? pageNumber = null) => new()
{
SourceUrl = sourceUrl,
DataSourceName = "Handbooks",
Path = sourceUrl,
Category = RetrievalContentCategory.TEXT,
Type = RetrievalContentType.TEXT_DOCUMENT,
Text = text,
PageNumber = pageNumber,
};
}