AI-Studio/app/Tests/Chat/ConversationPartsTests.cs

using System.Text.Json;

using AIStudio.Chat;
using AIStudio.Tools.ToolCallingSystem;

namespace AIStudio.Tests.Chat;

/// <summary>
/// Checks what a conversation is counted as costing before it is sent.
/// </summary>
/// <remarks>
/// The number under the input field used to count the sentence being typed and nothing else, which
/// answers a question nobody asks: what decides whether the next message fits is everything that
/// travels with it. So what is collected here has to be what the message builder actually sends --
/// no more, because a number which counts something that stays behind is wrong in the direction
/// that makes a person stop writing.
///
/// Beyond the messages, a request carries the schema of every tool the model may call, and, while
/// it runs, everything those tools have returned so far. Both are invisible on the screen, and the
/// second one is where a window fills up fastest.
/// </remarks>
[TestFixture]
public sealed class ConversationPartsTests
{
    private string directory = string.Empty;

    [SetUp]
    public void CreateFiles()
    {
        this.directory = Path.Combine(Path.GetTempPath(), $"ai-studio-tests-{Guid.NewGuid():N}");
        Directory.CreateDirectory(this.directory);
    }

    [TearDown]
    public void RemoveFiles()
    {
        if (Directory.Exists(this.directory))
            Directory.Delete(this.directory, true);
    }

    [Test]
    public void TheWholeConversationCountsAndNotOnlyWhatIsBeingTyped()
    {
        var thread = new ChatThread
        {
            SystemPrompt = "You are helpful.",
            Blocks =
            [
                Block("What is the capital of France?"),
                Block("Paris."),
            ],
        };

        var parts = ConversationParts.Of(thread, "You are helpful.", "And of Italy?", null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.EqualTo(new[] { "You are helpful.", "What is the capital of France?", "Paris." }));
            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "And of Italy?" }));
        });
    }

    [Test]
    public void WhatIsStillBeingWrittenIsKeptApartFromWhatStands()
    {
        //
        // Both cost the same and both are counted. They are kept apart because of what happens
        // afterwards: a message which stands says the same thing forever and its count is worth
        // remembering, while the answer being streamed is a different text three seconds later.
        //
        var streaming = Block("The answer so far");
        ((ContentText)streaming.Content!).IsStreaming = true;
        var thread = new ChatThread { Blocks = [Block("A question."), streaming] };

        var parts = ConversationParts.Of(thread, string.Empty, "a draft", null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.EqualTo(new[] { "A question." }));
            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "The answer so far", "a draft" }));
        });
    }

    [Test]
    public void AnAnswerWhichIsFinishedStandsLikeAnyOtherMessage()
    {
        var finished = Block("The whole answer.");
        ((ContentText)finished.Content!).IsStreaming = false;

        var parts = ConversationParts.Of(new() { Blocks = [finished] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.EqualTo(new[] { "The whole answer." }));
            Assert.That(parts.GrowingTexts, Is.Empty);
        });
    }

    [Test]
    public void TheSystemPromptCountedIsTheOneWhichWouldBeSent()
    {
        //
        // Not the one standing in the thread. A chat template may replace it, retrieved data is
        // appended to it, a profile adds a paragraph and the tool policy adds another -- and
        // switching a profile while writing has to move the number, which it cannot do if the
        // thread's own field is what gets counted.
        //
        var thread = new ChatThread { SystemPrompt = "What the person typed." };

        var parts = ConversationParts.Of(thread, "What the request carries.", string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.That(parts.Texts, Is.EqualTo(new[] { "What the request carries." }));
    }

    [Test]
    public void ABlockHiddenFromTheUserStillCosts()
    {
        //
        // Hidden on the screen, not in the request: the message builder sends it like any other
        // block, so its tokens are gone whether or not anybody can see where they went.
        //
        var hidden = Block("An instruction the user does not see.");
        var thread = new ChatThread { Blocks = [new() { ContentType = hidden.ContentType, Role = hidden.Role, Content = hidden.Content, HideFromUser = true }] };

        var parts = ConversationParts.Of(thread, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.That(parts.Texts, Is.EqualTo(new[] { "An instruction the user does not see." }));
    }

    [Test]
    public void WithoutAConversationOnlyTheDraftCounts()
    {
        var parts = ConversationParts.Of(null, string.Empty, "Hello", null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.Empty);
            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "Hello" }));
        });
    }

    [TestCase("")]
    [TestCase("   ")]
    public void NothingWrittenIsNothingToCount(string draft)
    {
        var parts = ConversationParts.Of(null, string.Empty, draft, null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.Empty);
            Assert.That(parts.GrowingTexts, Is.Empty);
        });
    }

    [Test]
    public void ABlockWithoutTextIsSkippedBecauseItIsNeverSent()
    {
        //
        // The message builder drops a block whose text is empty, whatever else hangs off it. A
        // count which added that block's attachments would report tokens for a message which is
        // never built.
        //
        var document = this.WriteFile("notes.txt", "some content");
        var empty = Block(string.Empty);
        ((ContentText)empty.Content!).FileAttachments.Add(FileAttachment.FromPath(document));

        var parts = ConversationParts.Of(new() { Blocks = [empty] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.Empty);
            Assert.That(parts.Documents, Is.Empty);
        });
    }

    [Test]
    public void AttachmentsOfTheConversationAndOfTheComposerBothCount()
    {
        //
        // A document attached three messages ago is sent again with every further message, so it
        // costs its tokens again every time. That is exactly the thing a person cannot see and
        // which this number is for.
        //
        var older = this.WriteFile("older.txt", "older content");
        var draft = this.WriteFile("draft.txt", "draft content");
        var block = Block("Please read this.");
        ((ContentText)block.Content!).FileAttachments.Add(FileAttachment.FromPath(older));

        var parts = ConversationParts.Of(new() { Blocks = [block] }, string.Empty, "And this one.", [FileAttachment.FromPath(draft)], imagesAreSent: true, toolDefinitions: null);

        Assert.That(parts.Documents.Select(document => document.FileName), Is.EqualTo(new[] { "older.txt", "draft.txt" }));
    }

    [Test]
    public void AnAttachmentWhoseFileIsGoneCountsForNothing()
    {
        //
        // It is not sent either: the message builder reports it as unavailable and leaves it out.
        //
        var attachment = FileAttachment.FromPath(Path.Combine(this.directory, "never-existed.txt"));

        var parts = ConversationParts.Of(null, string.Empty, "Here", [attachment], imagesAreSent: true, toolDefinitions: null);

        Assert.That(parts.Documents, Is.Empty);
    }

    [Test]
    public void ImagesAreCountedSeparatelyFromDocuments()
    {
        var document = this.WriteFile("notes.txt", "content");
        var image = this.WriteFile("photo.png", "not really a png");

        var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(document), FileAttachment.FromPath(image)], imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Documents.Select(entry => entry.FileName), Is.EqualTo(new[] { "notes.txt" }));
            Assert.That(parts.Images, Is.EqualTo(1));
        });
    }

    [Test]
    public void AModelWhichTakesNoImagesIsSentNoneAndIsToldAboutNone()
    {
        //
        // The message builder leaves the pictures out entirely for such a model, so reporting them
        // as uncounted would tell a person about a cost which is not there.
        //
        var image = this.WriteFile("photo.png", "not really a png");

        var parts = ConversationParts.Of(null, string.Empty, "Look", [FileAttachment.FromPath(image)], imagesAreSent: false, toolDefinitions: null);

        Assert.That(parts.Images, Is.Zero);
    }

    [Test]
    public void ABlockWithoutTextCountsWhileItsToolsAreStillRunning()
    {
        //
        // While a model calls tools there is no text yet: the answer arrives in one piece at the
        // end, and everything in between travels with every further round of the same request. The
        // block which looks emptiest is therefore the one whose request is growing the fastest --
        // and the one which used to be skipped for having nothing to say.
        //
        var running = Block(string.Empty);
        ((ContentText)running.Content!).PendingToolConversation = ["What the web search found.", "What the page said."];

        var parts = ConversationParts.Of(new() { Blocks = [running] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.Empty);
            Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "What the web search found.", "What the page said." }));
        });
    }

    [Test]
    public void TwoToolResultsWhichReadTheSameCostTwice()
    {
        //
        // The request carries both, so both are paid for. Folding them into one would promise a
        // smaller request than the one which is sent -- and a model reading the same page twice is
        // not a rare accident but a thing that happens on any busy search.
        //
        var running = Block(string.Empty);
        ((ContentText)running.Content!).PendingToolConversation = ["The same page.", "The same page."];

        var parts = ConversationParts.Of(new() { Blocks = [running] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.That(parts.GrowingTexts, Is.EqualTo(new[] { "The same page.", "The same page." }));
    }

    [Test]
    public void OnceTheAnswerStandsTheToolConversationIsGone()
    {
        //
        // It travels with the rounds of one request and with nothing afterwards: the next request is
        // built from the messages alone. A number which kept counting it would report a window
        // fuller than it is, and would never fall back.
        //
        var answered = Block("Here is what I found.");
        var content = (ContentText)answered.Content!;
        content.PendingToolConversation = ["What the web search found."];
        content.EndToolRun();

        var parts = ConversationParts.Of(new() { Blocks = [answered] }, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions: null);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.EqualTo(new[] { "Here is what I found." }));
            Assert.That(parts.GrowingTexts, Is.Empty);
        });
    }

    [Test]
    public void TheToolSchemasCountAndTheyCountWithWhatStands()
    {
        //
        // Every request carries the schema of every offered tool, whether or not the model calls a
        // single one of them. They belong with the lasting texts: a schema is the same string all
        // session long, so its count is worth remembering.
        //
        var parts = ConversationParts.Of(null, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions:
        [
            Tool("web_search", "Searches the web.", """{"type":"object"}"""),
            Tool("read_web_page", "Reads one page.", """{"type":"string"}"""),
        ]);

        Assert.Multiple(() =>
        {
            Assert.That(parts.Texts, Is.EqualTo(new[]
            {
                """web_searchSearches the web.{"type":"object"}""",
                """read_web_pageReads one page.{"type":"string"}""",
            }));

            Assert.That(parts.GrowingTexts, Is.Empty);
        });
    }

    [Test]
    public void AToolWhichStatesNoArgumentsCountsLikeAnyOther()
    {
        //
        // A definition which never names a parameter schema leaves an empty JSON element behind,
        // and asking such an element for its text throws. A tool arriving from a plugin may well
        // say nothing about its arguments, and the number under the input field is not the place
        // to find that out.
        //
        var parts = ConversationParts.Of(null, string.Empty, string.Empty, null, imagesAreSent: true, toolDefinitions:
        [
            new() { Function = new() { Name = "ping", DescriptionForLLM = "Says hello." } },
        ]);

        Assert.That(parts.Texts, Is.EqualTo(new[] { "pingSays hello." }));
    }

    private static ToolDefinition Tool(string name, string description, string parameterSchema) => new()
    {
        Function = new()
        {
            Name = name,
            DescriptionForLLM = description,
            Parameters = JsonDocument.Parse(parameterSchema).RootElement.Clone(),
        },
    };

    private static ContentBlock Block(string text) => new()
    {
        ContentType = ContentType.TEXT,
        Role = ChatRole.USER,
        Content = new ContentText { Text = text },
    };

    private string WriteFile(string name, string content)
    {
        var path = Path.Combine(this.directory, name);
        File.WriteAllText(path, content);
        return path;
    }
}