meeting-assistant/MeetingAssistant.Tests/PyannoteTranscriptFinalizerTests.cs

using MeetingAssistant;
using MeetingAssistant.Transcription;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;

namespace MeetingAssistant.Tests;

public sealed class PyannoteTranscriptFinalizerTests
{
    [Fact]
    public async Task FinalizerAssignsPyannoteSpeakersToLiveWhisperSegmentsByOverlap()
    {
        var audioPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "meeting.wav");
        Directory.CreateDirectory(Path.GetDirectoryName(audioPath)!);
        await File.WriteAllBytesAsync(audioPath, [1, 2, 3, 4]);
        var commandRunner = new CapturingCommandRunner(
            """
            install noise
            __MEETING_ASSISTANT_PYANNOTE_JSON_START__
            [{"start":0.0,"end":1.5,"speaker":"SPEAKER_00"},{"start":1.5,"end":3.0,"speaker":"SPEAKER_01"}]
            __MEETING_ASSISTANT_PYANNOTE_JSON_END__
            """);
        var finalizer = CreateFinalizer(
            commandRunner,
            token: "hf_test",
            alignmentMode: PyannoteAlignmentMode.BestOverlap);
        var liveSegments = new[]
        {
            new TranscriptionSegment(TimeSpan.FromSeconds(0.2), TimeSpan.FromSeconds(1.2), "Unknown", "hello"),
            new TranscriptionSegment(TimeSpan.FromSeconds(1.7), TimeSpan.FromSeconds(2.4), "Unknown", "there")
        };

        var segments = await finalizer.FinalizeAsync(
            audioPath,
            liveSegments,
            SpeechRecognitionPipelineOptions.Default,
            CancellationToken.None);

        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("inspect"));
        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("--format"));
        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("build"));
        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("meeting-assistant-pyannote:local"));
        Assert.Equal("docker", commandRunner.Commands.Last().FileName);
        Assert.Contains("run", commandRunner.Commands.Last().Arguments);
        Assert.Contains("meeting-assistant-pyannote:local", commandRunner.Commands.Last().Arguments);
        Assert.Contains("HF_TOKEN", commandRunner.Commands.Last().Arguments);
        Assert.DoesNotContain("HF_TOKEN=hf_test", commandRunner.Commands.Last().Arguments);
        Assert.Equal("hf_test", commandRunner.Environment["HF_TOKEN"]);
        Assert.Contains("sh", commandRunner.Commands.Last().Arguments);
        Assert.DoesNotContain("bash", commandRunner.Commands.Last().Arguments);
        Assert.DoesNotContain(commandRunner.Commands.Last().Arguments, argument => argument.Contains("pip install", StringComparison.Ordinal));
        Assert.DoesNotContain(commandRunner.Commands.Last().Arguments, argument => argument.Contains("apt-get install", StringComparison.Ordinal));
        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("PIP_CACHE_DIR=/workspace/cache/pip", StringComparison.Ordinal));
        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("HF_HOME=/workspace/cache/huggingface", StringComparison.Ordinal));
        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("TORCH_HOME=/workspace/cache/torch", StringComparison.Ordinal));
        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("speaker_diarization", StringComparison.Ordinal));
        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("token=os.environ.get('HF_TOKEN')", StringComparison.Ordinal));
        Assert.DoesNotContain(commandRunner.Commands.Last().Arguments, argument => argument.Contains("use_auth_token", StringComparison.Ordinal));
        Assert.Collection(
            segments,
            first =>
            {
                Assert.Equal("SPEAKER_00", first.Speaker);
                Assert.Equal("hello", first.Text);
            },
            second =>
            {
                Assert.Equal("SPEAKER_01", second.Speaker);
                Assert.Equal("there", second.Text);
            });
    }

    [Fact]
    public async Task FinalizerCanBuildTranscriptSegmentsFromPyannoteTurns()
    {
        var audioPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "meeting.wav");
        Directory.CreateDirectory(Path.GetDirectoryName(audioPath)!);
        await File.WriteAllBytesAsync(audioPath, [1, 2, 3, 4]);
        var commandRunner = new CapturingCommandRunner(
            """
            __MEETING_ASSISTANT_PYANNOTE_JSON_START__
            [{"start":0.0,"end":1.0,"speaker":"SPEAKER_00"},{"start":1.0,"end":2.0,"speaker":"SPEAKER_01"}]
            __MEETING_ASSISTANT_PYANNOTE_JSON_END__
            """);
        var finalizer = CreateFinalizer(
            commandRunner,
            token: "hf_test",
            annotationSource: PyannoteAnnotationSource.ExclusiveSpeakerDiarization,
            alignmentMode: PyannoteAlignmentMode.PyannoteTurns);
        var liveSegments = new[]
        {
            new TranscriptionSegment(TimeSpan.Zero, TimeSpan.FromSeconds(2), "Unknown", "alpha beta gamma delta")
        };

        var segments = await finalizer.FinalizeAsync(
            audioPath,
            liveSegments,
            SpeechRecognitionPipelineOptions.Default,
            CancellationToken.None);

        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("exclusive_speaker_diarization", StringComparison.Ordinal));
        Assert.Collection(
            segments,
            first =>
            {
                Assert.Equal(TimeSpan.Zero, first.Start);
                Assert.Equal(TimeSpan.FromSeconds(1), first.End);
                Assert.Equal("SPEAKER_00", first.Speaker);
                Assert.Equal("alpha beta", first.Text);
            },
            second =>
            {
                Assert.Equal(TimeSpan.FromSeconds(1), second.Start);
                Assert.Equal(TimeSpan.FromSeconds(2), second.End);
                Assert.Equal("SPEAKER_01", second.Speaker);
                Assert.Equal("gamma delta", second.Text);
            });
    }

    [Fact]
    public async Task FinalizerPassesConfiguredNumSpeakersToPyannote()
    {
        var audioPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "meeting.wav");
        Directory.CreateDirectory(Path.GetDirectoryName(audioPath)!);
        await File.WriteAllBytesAsync(audioPath, [1, 2, 3, 4]);
        var commandRunner = new CapturingCommandRunner(
            """
            __MEETING_ASSISTANT_PYANNOTE_JSON_START__
            [{"start":0.0,"end":1.0,"speaker":"SPEAKER_00"}]
            __MEETING_ASSISTANT_PYANNOTE_JSON_END__
            """);
        var finalizer = CreateFinalizer(commandRunner, token: "hf_test");

        await finalizer.FinalizeAsync(
            audioPath,
            [new TranscriptionSegment(TimeSpan.Zero, TimeSpan.FromSeconds(1), "Unknown", "hello")],
            new SpeechRecognitionPipelineOptions(5),
            CancellationToken.None);

        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("num_speakers=5", StringComparison.Ordinal));
    }

    [Fact]
    public async Task FinalizerCanUseExplicitDiarizationOptions()
    {
        var audioPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "meeting.wav");
        Directory.CreateDirectory(Path.GetDirectoryName(audioPath)!);
        await File.WriteAllBytesAsync(audioPath, [1, 2, 3, 4]);
        var commandRunner = new CapturingCommandRunner(
            """
            __MEETING_ASSISTANT_PYANNOTE_JSON_START__
            [{"start":0.0,"end":1.0,"speaker":"SPEAKER_00"}]
            __MEETING_ASSISTANT_PYANNOTE_JSON_END__
            """);
        var finalizer = new PyannoteTranscriptFinalizer(
            commandRunner,
            Options.Create(new MeetingAssistantOptions()),
            NullLogger<PyannoteTranscriptFinalizer>.Instance);
        var explicitDiarization = new PyannoteDiarizationOptions
        {
            Enabled = true,
            DockerCommand = "docker",
            BaseImage = "python:3.11-slim",
            Image = "meeting-assistant-pyannote-azure:local",
            ModelsFolder = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "models"),
            Model = "custom/diarization",
            Token = "hf_azure",
            TokenEnv = "",
            CommandTimeout = TimeSpan.FromMinutes(1)
        };

        await finalizer.FinalizeAsync(
            audioPath,
            [new TranscriptionSegment(TimeSpan.Zero, TimeSpan.FromSeconds(1), "Unknown", "hello")],
            explicitDiarization,
            SpeechRecognitionPipelineOptions.Default,
            CancellationToken.None);

        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("meeting-assistant-pyannote-azure:local"));
        Assert.Contains(commandRunner.Commands.Last().Arguments, argument => argument.Contains("custom/diarization", StringComparison.Ordinal));
        Assert.Equal("hf_azure", commandRunner.Environment["HF_TOKEN"]);
    }

    [Fact]
    public async Task WarmUpBuildsImageAndDownloadsConfiguredModelWithoutAudioInput()
    {
        var commandRunner = new CapturingCommandRunner("");
        var finalizer = CreateFinalizer(commandRunner, token: "hf_test");
        var diarization = new PyannoteDiarizationOptions
        {
            Enabled = true,
            DockerCommand = "docker",
            BaseImage = "python:3.11-slim",
            Image = "meeting-assistant-pyannote-warmup:local",
            ModelsFolder = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "models"),
            Model = "pyannote/speaker-diarization-3.1",
            Token = "hf_test",
            TokenEnv = "",
            CommandTimeout = TimeSpan.FromMinutes(1)
        };

        await finalizer.WarmUpAsync(diarization, CancellationToken.None);

        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("inspect"));
        Assert.Contains(commandRunner.Commands, command => command.Arguments.Contains("build"));
        var runCommand = commandRunner.Commands.Last();
        Assert.Contains("run", runCommand.Arguments);
        Assert.Contains("meeting-assistant-pyannote-warmup:local", runCommand.Arguments);
        Assert.DoesNotContain(runCommand.Arguments, argument => argument.Contains("/workspace/input.wav", StringComparison.Ordinal));
        Assert.Contains(runCommand.Arguments, argument => argument.Contains("Pipeline.from_pretrained", StringComparison.Ordinal));
        Assert.Contains(runCommand.Arguments, argument => argument.Contains("pyannote/speaker-diarization-3.1", StringComparison.Ordinal));
        Assert.Equal("hf_test", commandRunner.Environment["HF_TOKEN"]);
    }

    [Fact]
    public async Task WarmUpSkipsWhenTokenIsMissing()
    {
        var commandRunner = new CapturingCommandRunner("");
        var finalizer = CreateFinalizer(commandRunner, token: null);

        await finalizer.WarmUpAsync(new PyannoteDiarizationOptions
        {
            Enabled = true,
            Token = "",
            TokenEnv = "",
            ModelsFolder = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "models")
        }, CancellationToken.None);

        Assert.Empty(commandRunner.Commands);
    }

    [Fact]
    public async Task FinalizerSkipsPyannoteWhenTokenIsMissing()
    {
        var commandRunner = new CapturingCommandRunner("");
        var finalizer = CreateFinalizer(commandRunner, token: null);
        var audioPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "meeting.wav");
        Directory.CreateDirectory(Path.GetDirectoryName(audioPath)!);
        await File.WriteAllBytesAsync(audioPath, [1, 2, 3, 4]);

        var segments = await finalizer.FinalizeAsync(
            audioPath,
            [new TranscriptionSegment(TimeSpan.Zero, TimeSpan.FromSeconds(1), "Unknown", "hello")],
            SpeechRecognitionPipelineOptions.Default,
            CancellationToken.None);

        Assert.Empty(segments);
        Assert.Empty(commandRunner.Commands);
    }

    private static PyannoteTranscriptFinalizer CreateFinalizer(
        CapturingCommandRunner commandRunner,
        string? token,
        PyannoteAnnotationSource annotationSource = PyannoteAnnotationSource.SpeakerDiarization,
        PyannoteAlignmentMode alignmentMode = PyannoteAlignmentMode.PyannoteTurns)
    {
        return new PyannoteTranscriptFinalizer(
            commandRunner,
            Options.Create(new MeetingAssistantOptions
            {
                WhisperLocal = new WhisperLocalOptions
                {
                    Diarization = new PyannoteDiarizationOptions
                    {
                        Enabled = true,
                        DockerCommand = "docker",
                        BaseImage = "python:3.11-slim",
                        Image = "meeting-assistant-pyannote:local",
                        ModelsFolder = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N"), "models"),
                        AnnotationSource = annotationSource,
                        AlignmentMode = alignmentMode,
                        Token = token,
                        TokenEnv = "",
                        CommandTimeout = TimeSpan.FromMinutes(1)
                    }
                }
            }),
            NullLogger<PyannoteTranscriptFinalizer>.Instance);
    }

    private sealed class CapturingCommandRunner : ICommandRunner
    {
        private readonly string output;

        public CapturingCommandRunner(string output)
        {
            this.output = output;
        }

        public IReadOnlyList<CapturedCommand> Commands { get; private set; } = [];

        public IReadOnlyDictionary<string, string> Environment { get; private set; } =
            new Dictionary<string, string>();

        public Task<CommandResult> RunAsync(
            string fileName,
            IReadOnlyList<string> arguments,
            CancellationToken cancellationToken,
            IReadOnlyDictionary<string, string>? environment = null)
        {
            Commands = Commands.Append(new CapturedCommand(fileName, arguments)).ToList();
            Environment = environment ?? new Dictionary<string, string>();
            if (arguments.Contains("inspect"))
            {
                return Task.FromResult(new CommandResult(1, "", "image missing"));
            }

            return Task.FromResult(new CommandResult(0, output, ""));
        }
    }

    private sealed record CapturedCommand(string FileName, IReadOnlyList<string> Arguments);
}