novelly/tests/Novelly.Api.Tests/NovelAgentServiceTests.cs

using System.Text.Json;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using Novelly.Api.Agent;
using Novelly.Api.Novels;

namespace Novelly.Api.Tests;

[TestFixture]
public class NovelAgentServiceTests : ServiceTestFixture
{
    private NovelAgentToolset _toolset = null!;

    protected override void OnSetUp() =>
        _toolset = new NovelAgentToolset(Novels, Characters, Arcs, Chapters, ChapterLabels, Beats, Tags, Locations, Questions, NullLogger<NovelAgentToolset>.Instance);

    private NovelAgentService BuildAgent(ScriptedModelClient model) => new(
        Db.Context,
        model,
        _toolset,
        Options.Create(new AgentOptions { MaxIterations = 4 }),
        NullLogger<NovelAgentService>.Instance,
        new SendAgentMessageRequestValidator());

    [Test]
    public async Task A_plain_reply_is_persisted_as_a_conversation()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;
        var model = new ScriptedModelClient([[new AgentTextBlock("Tell me about the ending.")]]);
        var agent = BuildAgent(model);

        var turn = await agent.SendMessageAsync(novelId, new SendAgentMessageRequest("Where do I start?"));

        Assert.That(turn.Content, Is.EqualTo("Tell me about the ending."));

        var conversation = (await agent.GetConversationAsync(turn.ConversationId))!;

        Assert.Multiple(() =>
        {
            Assert.That(conversation.Messages, Has.Count.EqualTo(2));
            Assert.That(conversation.Messages[0].Content, Is.EqualTo("Where do I start?"));
            Assert.That(conversation.Title, Is.EqualTo("Where do I start?"));
        });
    }

    [Test]
    public async Task Tool_calls_are_executed_against_real_novel_data()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;

        var model = new ScriptedModelClient([
            [ToolUse("t1", "create_character", new { name = "Ines", role = "Protagonist" })],
            [new AgentTextBlock("Added Ines as the protagonist.")]
        ]);

        var turn = await BuildAgent(model).SendMessageAsync(
            novelId, new SendAgentMessageRequest("Add a protagonist called Ines."));

        var characters = await Characters.ListAsync(novelId);

        Assert.Multiple(() =>
        {
            Assert.That(characters, Has.Count.EqualTo(1));
            Assert.That(characters[0].Name, Is.EqualTo("Ines"));
            Assert.That(turn.Content, Is.EqualTo("Added Ines as the protagonist."));
            Assert.That(turn.ToResponse().ToolCalls, Has.Count.EqualTo(1));
            Assert.That(turn.ToResponse().ToolCalls[0].Name, Is.EqualTo("create_character"));
        });
    }

    [Test]
    public async Task Every_tool_result_comes_back_in_a_single_user_turn()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;

        var model = new ScriptedModelClient([
            [
                ToolUse("t1", "create_character", new { name = "Ines" }),
                ToolUse("t2", "create_character", new { name = "Mara" })
            ],
            [new AgentTextBlock("Both added.")]
        ]);

        await BuildAgent(model).SendMessageAsync(novelId, new SendAgentMessageRequest("Add two characters."));

        var resultTurn = model.Transcripts[1][^1];
        var listed = await Characters.ListAsync(novelId);

        Assert.Multiple(() =>
        {
            Assert.That(resultTurn.Role, Is.EqualTo("user"));
            Assert.That(resultTurn.Content.OfType<AgentToolResultBlock>().Count(), Is.EqualTo(2));
            Assert.That(listed, Has.Count.EqualTo(2));
        });
    }

    [Test]
    public async Task A_failing_tool_is_reported_back_rather_than_thrown()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;

        var model = new ScriptedModelClient([
            [ToolUse("t1", "update_character", new { character_id = Guid.NewGuid().ToString(), name = "Ines" })],
            [new AgentTextBlock("That character does not exist yet — shall I create her?")]
        ]);

        var turn = await BuildAgent(model).SendMessageAsync(
            novelId, new SendAgentMessageRequest("Rename her."));

        var errorResult = model.Transcripts[1][^1].Content.OfType<AgentToolResultBlock>().Single();

        Assert.Multiple(() =>
        {
            Assert.That(errorResult.IsError, Is.True);
            Assert.That(errorResult.Content, Does.Contain("was not found"));
            Assert.That(turn.Content, Does.Contain("does not exist yet"));
        });
    }

    [Test]
    public async Task Unknown_tools_are_reported_without_breaking_the_loop()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;

        var model = new ScriptedModelClient([
            [ToolUse("t1", "summon_muse", new { })],
            [new AgentTextBlock("Sorry — I do not have that tool.")]
        ]);

        await BuildAgent(model).SendMessageAsync(novelId, new SendAgentMessageRequest("Summon the muse."));

        var result = model.Transcripts[1][^1].Content.OfType<AgentToolResultBlock>().Single();

        Assert.Multiple(() =>
        {
            Assert.That(result.IsError, Is.True);
            Assert.That(result.Content, Does.Contain("No such tool"));
        });
    }

    [Test]
    public async Task The_loop_stops_at_the_iteration_ceiling()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;

        var model = new ScriptedModelClient(
            Enumerable.Repeat<IReadOnlyList<AgentContentBlock>>(
                [ToolUse("t", "list_characters", new { })], 20).ToList());

        var turn = await BuildAgent(model).SendMessageAsync(
            novelId, new SendAgentMessageRequest("Keep going forever."));

        Assert.Multiple(() =>
        {
            Assert.That(model.Transcripts, Has.Count.EqualTo(4));
            Assert.That(turn.Content, Does.Contain("tool-call limit"));
        });
    }

    [Test]
    public async Task Follow_up_messages_continue_the_same_conversation()
    {
        var novelId = (await Novels.CreateAsync(new CreateNovelRequest("The Salt Road"))).Id;
        var model = new ScriptedModelClient([
            [new AgentTextBlock("First answer.")],
            [new AgentTextBlock("Second answer.")]
        ]);
        var agent = BuildAgent(model);

        var first = await agent.SendMessageAsync(novelId, new SendAgentMessageRequest("Question one."));
        var second = await agent.SendMessageAsync(
            novelId, new SendAgentMessageRequest("Question two.", first.ConversationId));

        var conversation = (await agent.GetConversationAsync(first.ConversationId))!;

        Assert.Multiple(() =>
        {
            Assert.That(second.ConversationId, Is.EqualTo(first.ConversationId));

            Assert.That(model.Transcripts[1], Has.Count.EqualTo(3));
            Assert.That(
                model.Transcripts[1].Select(m => m.Role),
                Is.EqualTo(new[] { "user", "assistant", "user" }));
            Assert.That(conversation.Messages, Has.Count.EqualTo(4));
        });
    }

    [Test]
    public void Every_tool_declares_an_object_schema_and_a_description()
    {
        Assert.That(_toolset.Definitions, Is.Not.Empty);

        Assert.Multiple(() =>
        {
            foreach (var tool in _toolset.Definitions)
            {
                Assert.That(string.IsNullOrWhiteSpace(tool.Description), Is.False, tool.Name);
                Assert.That(tool.InputSchema.GetProperty("type").GetString(), Is.EqualTo("object"), tool.Name);
                Assert.That(tool.InputSchema.TryGetProperty("properties", out _), Is.True, tool.Name);
            }

            Assert.That(_toolset.Definitions.Select(t => t.Name), Is.Unique);
        });
    }

    private static AgentToolUseBlock ToolUse(string id, string name, object input) =>
        new(id, name, JsonSerializer.SerializeToElement(input));
}

internal class ScriptedModelClient(IReadOnlyList<IReadOnlyList<AgentContentBlock>> script)
    : IAgentModelClient
{
    private int _turn;

    public List<IReadOnlyList<AgentChatMessage>> Transcripts { get; } = [];

    public Task<AgentModelResponse> CompleteAsync(
        string systemPrompt,
        IReadOnlyList<AgentChatMessage> messages,
        IReadOnlyList<AgentToolDefinition> tools,
        CancellationToken ct = default)
    {
        Transcripts.Add([.. messages]);

        var content = _turn < script.Count ? script[_turn] : [new AgentTextBlock("(no more script)")];
        _turn++;

        var stopReason = content.OfType<AgentToolUseBlock>().Any() ? "tool_use" : "end_turn";
        return Task.FromResult(new AgentModelResponse(content, stopReason));
    }
}