Wiki/test/Wiki.Application.Tests/AskUseCaseTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using Wiki.Application.Ask;
using Wiki.Application.Retrieval;
using Wiki.Application.Sync;
using Wiki.Core.Model;
using Wiki.Core.Parsing;
using Wiki.Core.Ports;

namespace Wiki.Application.Tests;

/// <summary>A judge that records requests and answers by simple rules, so the use case's wiring is testable.</summary>
public sealed class FakeJudge : IJudge
{
    public List<JudgeRequest> Requests { get; } = [];
    public string ModelId => "fake-judge";
    public int? MaxChoiceOptions { get; init; } = 255;
    public string Name => "fake-judge";

    public Task<JudgeResult> JudgeAsync(JudgeRequest request, CancellationToken ct)
    {
        Requests.Add(request);
        var answers = new Dictionary<string, Judgement>();
        foreach (var (id, q) in request.Questions)
        {
            answers[id] = q switch
            {
                TypedQuestion.Noul => new Judgement.Noul(id == AskUseCase.CoveredId ? 0.8 : 0.9),
                TypedQuestion.Choice c => Spread(c),
                TypedQuestion.Score => new Judgement.Score(0.5),
                _ => throw new InvalidOperationException(),
            };
        }
        return Task.FromResult(new JudgeResult(answers, new JudgeUsage(1, 100, request.Questions.Count, TimeSpan.FromMilliseconds(5))));
    }

    private static Judgement.Choice Spread(TypedQuestion.Choice c)
    {
        // First option gets 0.6, the rest share 0.4 equally (NONE included).
        var keys = c.Options.Keys.ToList();
        var probs = new Dictionary<string, double>();
        for (var i = 0; i < keys.Count; i++)
            probs[keys[i]] = keys.Count == 1 ? 1.0 : (i == 0 ? 0.6 : 0.4 / (keys.Count - 1));
        return new Judgement.Choice(keys[0], probs, 0.6);
    }

    public Task<AdapterHealth> CheckAsync(CancellationToken ct) => Task.FromResult(new AdapterHealth(Name, true, null, TimeSpan.Zero, DateTimeOffset.UnixEpoch));
}

public class AskUseCaseTests
{
    private readonly FakeStore _store = new();
    private readonly FakeCatalog _catalog = new();
    private readonly FakeIndex _index = new();
    private readonly FakeEmbeddings _model = new("m1");
    private readonly FakeJudge _judge = new();

    private AskUseCase UseCase => new(new RankUseCase(_index, _catalog), _catalog, new ParagraphChunker());

    private async Task SeedAsync()
    {
        _store.Files["wiki/udm-dns.md"] = "---\ntitle: UDM intercepts DNS\ntype: wiki\ntags: [dns]\n---\n\n# UDM intercepts DNS\n\nThe UDM answers every query.\n\n## Details\n\nEven dig @1.1.1.1 is answered locally.\n";
        _store.Files["wiki/ceph-osd.md"] = "---\ntitle: Ceph OSD\ntype: wiki\ntags: [ceph]\n---\n\n# Ceph OSD\n\nOSD fsid mismatch after rebuild.\n";
        _store.Files["wiki/long-note.md"] = "---\ntitle: Long\ntype: wiki\n---\n\n# Long\n\n" + string.Join("\n\n", Enumerable.Range(0, 40).Select(i => $"paragraph {i} " + new string('x', 400))) + "\n";
        var sync = new SyncUseCase(_store, _catalog, _index, [_model], TimeProvider.System, NullLogger<SyncUseCase>.Instance);
        await sync.RunAsync(false, CancellationToken.None);
    }

    [Fact]
    public async Task DefaultQuestionsPerQueryAndCoveredGate()
    {
        await SeedAsync();
        var udm = (await _catalog.GetAsync(DocumentKind.Wiki, "udm-dns", CancellationToken.None))!;
        var request = new AskRequest([EmbeddingText.For(udm)], null, new AskOptions(K: 2));

        var result = await UseCase.AskAsync(request, _model, _judge, CancellationToken.None);

        Assert.Equal("udm-dns", result.Context[0].Slug);
        Assert.Equal(0.8, result.Covered);
        Assert.Equal(["answered_1", "where_1", AskUseCase.CoveredId], result.Answers.Keys.OrderBy(k => k == AskUseCase.CoveredId ? 2 : k.StartsWith("answered") ? 0 : 1));
        var where = Assert.IsType<AskAnswer.Context>(result.Answers["where_1"]);
        Assert.Equal("udm-dns#p1", where.Top[0].Id);
        Assert.Contains("answers every query", where.Top[0].Text);
        Assert.Equal("UDM intercepts DNS", where.Top[0].Title);
        // every question rides in the one request that carries the state
        var sentToJudge = Assert.Single(_judge.Requests);
        Assert.Equal(3, sentToJudge.Questions.Count);
        Assert.Equal(1, result.Usage.Requests);
        Assert.Equal("fake-judge", result.JudgeModel);
    }

    [Fact]
    public async Task ContextChoiceOptionsAreTheParagraphIdsPlusNone()
    {
        await SeedAsync();
        var udm = (await _catalog.GetAsync(DocumentKind.Wiki, "udm-dns", CancellationToken.None))!;
        var request = new AskRequest([EmbeddingText.For(udm)], new Dictionary<string, AskQuestion>
        {
            ["src"] = new AskQuestion.ContextChoice("which note", ContextScope.Documents),
            ["fact"] = new AskQuestion.Noul("is it intercepted", "yes", "no"),
        }, new AskOptions(K: 2));

        var result = await UseCase.AskAsync(request, _model, _judge, CancellationToken.None);

        var sent = _judge.Requests.Single(r => r.Questions.ContainsKey("src"));
        var choice = Assert.IsType<TypedQuestion.Choice>(sent.Questions["src"]);
        Assert.Equal(result.Context.Select(d => d.Slug).Append(AskUseCase.None).OrderBy(x => x), choice.Options.Keys.OrderBy(x => x));
        Assert.Equal(string.Empty, choice.Options[result.Context[0].Slug]);
        var src = Assert.IsType<AskAnswer.Context>(result.Answers["src"]);
        Assert.Equal("udm-dns", src.Top[0].Slug);
        Assert.Null(src.Top[0].Text);
        Assert.IsType<AskAnswer.Noul>(result.Answers["fact"]);
        Assert.True(result.Answers.ContainsKey(AskUseCase.CoveredId));
    }

    [Fact]
    public async Task BudgetsCapLongDocumentsAndDropTheLowest()
    {
        await SeedAsync();
        var longNote = (await _catalog.GetAsync(DocumentKind.Wiki, "long-note", CancellationToken.None))!;
        var (context, meta) = await UseCase.BuildContextAsync([EmbeddingText.For(longNote)], _model, new AskOptions(K: 3, DocChars: 3000, StateChars: 2950), CancellationToken.None);
        Assert.Equal("long-note", context[0].Slug);
        Assert.True(context[0].Truncated);
        Assert.True(context[0].Paragraphs.Count < 40);
        Assert.Equal(2, meta.DroppedForBudget.Count);
        Assert.Single(context);
    }

    [Fact]
    public async Task RejectsAContextChoiceOverTheJudgeCap()
    {
        await SeedAsync();
        var udm = (await _catalog.GetAsync(DocumentKind.Wiki, "udm-dns", CancellationToken.None))!;
        var small = new FakeJudge { MaxChoiceOptions = 3 };
        var request = new AskRequest([EmbeddingText.For(udm)], null, new AskOptions(K: 3));
        await Assert.ThrowsAsync<ArgumentException>(() => UseCase.AskAsync(request, _model, small, CancellationToken.None));
    }
}