meeting-assistant/MeetingAssistant/Recording/SpeakerAudioSampleCollector.cs

using MeetingAssistant.Speakers;
using MeetingAssistant.Transcription;

namespace MeetingAssistant.Recording;

internal sealed class SpeakerAudioSampleCollector
{
    private readonly object gate = new();
    private readonly RollingAudioBuffer audioBuffer;
    private readonly Dictionary<string, List<SpeakerAudioSample>> samplesBySpeaker = new(StringComparer.OrdinalIgnoreCase);
    private readonly int maxSamplesPerSpeaker;
    private readonly TimeSpan minimumUninterruptedSpeechDuration;
    private readonly TimeSpan maximumSegmentGap;
    private PendingSpeakerSpan? pendingSpan;

    public SpeakerAudioSampleCollector(TimeSpan bufferDuration, int maxSamplesPerSpeaker)
        : this(
            bufferDuration,
            maxSamplesPerSpeaker,
            TimeSpan.FromSeconds(30),
            TimeSpan.FromSeconds(1))
    {
    }

    public SpeakerAudioSampleCollector(
        TimeSpan bufferDuration,
        int maxSamplesPerSpeaker,
        TimeSpan minimumUninterruptedSpeechDuration,
        TimeSpan maximumSegmentGap)
    {
        audioBuffer = new RollingAudioBuffer(bufferDuration);
        this.maxSamplesPerSpeaker = Math.Max(1, maxSamplesPerSpeaker);
        this.minimumUninterruptedSpeechDuration = minimumUninterruptedSpeechDuration > TimeSpan.Zero
            ? minimumUninterruptedSpeechDuration
            : TimeSpan.Zero;
        this.maximumSegmentGap = maximumSegmentGap >= TimeSpan.Zero
            ? maximumSegmentGap
            : TimeSpan.Zero;
    }

    public void AppendAudio(AudioChunk chunk)
    {
        audioBuffer.Append(chunk);
    }

    public void Reset()
    {
        lock (gate)
        {
            samplesBySpeaker.Clear();
            pendingSpan = null;
            audioBuffer.Reset();
        }
    }

    public void TryAdd(TranscriptionSegment segment)
    {
        if (!IsDiarizedSpeaker(segment.Speaker))
        {
            return;
        }

        TranscriptionSegment sampleSegment;
        lock (gate)
        {
            sampleSegment = ExtendPendingSpan(segment);
        }

        var score = Score(sampleSegment, minimumUninterruptedSpeechDuration);
        if (score <= 0)
        {
            return;
        }

        var wavBytes = audioBuffer.TryExtractWav(sampleSegment.Start, sampleSegment.End);
        if (wavBytes.Length == 0)
        {
            return;
        }

        var sample = new SpeakerAudioSample(sampleSegment.Speaker, sampleSegment, wavBytes, score);
        lock (gate)
        {
            if (!samplesBySpeaker.TryGetValue(sampleSegment.Speaker, out var samples))
            {
                samples = [];
                samplesBySpeaker[sampleSegment.Speaker] = samples;
            }

            samples.Add(sample);
            var bestSamples = samples
                .OrderByDescending(candidate => candidate.Score)
                .Take(maxSamplesPerSpeaker)
                .ToList();
            samples.Clear();
            samples.AddRange(bestSamples);
        }
    }

    public IReadOnlyList<SpeakerAudioSample> Snapshot()
    {
        lock (gate)
        {
            return samplesBySpeaker.Values
                .SelectMany(samples => samples)
                .OrderBy(sample => sample.Speaker, StringComparer.OrdinalIgnoreCase)
                .ThenByDescending(sample => sample.Score)
                .ToList();
        }
    }

    private static bool IsDiarizedSpeaker(string speaker)
    {
        return !string.IsNullOrWhiteSpace(speaker) &&
            !string.Equals(speaker, "Unknown", StringComparison.OrdinalIgnoreCase);
    }

    private TranscriptionSegment ExtendPendingSpan(TranscriptionSegment segment)
    {
        if (pendingSpan is null ||
            !SpeakerSampleSpanSelector.CanExtend(pendingSpan.Speaker, pendingSpan.End, segment, maximumSegmentGap))
        {
            pendingSpan = new PendingSpeakerSpan(
                segment.Speaker,
                segment.Start,
                segment.End,
                [segment.Text]);
            return pendingSpan.ToSegment();
        }

        pendingSpan = pendingSpan.Extend(segment);
        return pendingSpan.ToSegment();
    }

    private static double Score(
        TranscriptionSegment segment,
        TimeSpan minimumUninterruptedSpeechDuration)
    {
        var durationSeconds = (segment.End - segment.Start).TotalSeconds;
        if (durationSeconds < minimumUninterruptedSpeechDuration.TotalSeconds)
        {
            return 0;
        }

        var words = segment.Text
            .Split(' ', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries)
            .Length;
        if (words < 3)
        {
            return 0;
        }

        var durationScore = Math.Min(durationSeconds / Math.Max(1, minimumUninterruptedSpeechDuration.TotalSeconds), 2);
        var wordScore = Math.Min(words / 60.0, 1);
        var sentenceBonus = segment.Text.TrimEnd().EndsWith('.') ||
            segment.Text.TrimEnd().EndsWith('?') ||
            segment.Text.TrimEnd().EndsWith('!')
                ? 5
                : 0;
        return durationScore * 70 + wordScore * 30 + sentenceBonus;
    }

    private sealed record PendingSpeakerSpan(
        string Speaker,
        TimeSpan Start,
        TimeSpan End,
        IReadOnlyList<string> TextParts)
    {
        public PendingSpeakerSpan Extend(TranscriptionSegment segment)
        {
            return this with
            {
                End = segment.End > End ? segment.End : End,
                TextParts = TextParts.Append(segment.Text).ToList()
            };
        }

        public TranscriptionSegment ToSegment()
        {
            return new TranscriptionSegment(
                Start,
                End,
                Speaker,
                string.Join(' ', TextParts.Where(part => !string.IsNullOrWhiteSpace(part))));
        }
    }
}