meeting-assistant/MeetingAssistant/Speakers/AzureSpeechSpeakerIdentityMatcher.cs

using MeetingAssistant.Transcription;
using NAudio.Wave;

namespace MeetingAssistant.Speakers;

public interface ISpeakerIdentityDiarizationClient
{
    Task<IReadOnlyList<TranscriptionSegment>> DiarizeAsync(
        string wavPath,
        CancellationToken cancellationToken);
}

public sealed class AzureSpeechSpeakerIdentityMatcher : ISpeakerIdentityMatcher
{
    private readonly ISpeakerIdentityDiarizationClient diarizationClient;
    private readonly SpeakerIdentificationOptions options;
    private readonly ILogger<AzureSpeechSpeakerIdentityMatcher> logger;
    private readonly ISpeakerIdentityMatchValidator matchValidator;

    public AzureSpeechSpeakerIdentityMatcher(
        ISpeakerIdentityDiarizationClient diarizationClient,
        Microsoft.Extensions.Options.IOptions<MeetingAssistantOptions> options,
        ILogger<AzureSpeechSpeakerIdentityMatcher> logger,
        ISpeakerIdentityMatchValidator matchValidator)
    {
        this.diarizationClient = diarizationClient;
        this.options = options.Value.SpeakerIdentification;
        this.logger = logger;
        this.matchValidator = matchValidator;
    }

    public async Task<SpeakerIdentityMatch?> MatchAsync(
        SpeakerIdentityMatchRequest request,
        CancellationToken cancellationToken)
    {
        if (!options.Enabled || request.Candidates.Count == 0)
        {
            return null;
        }

        if (!await matchValidator.ValidateSampleAsync(request.UnknownSnippet, cancellationToken))
        {
            logger.LogInformation(
                "Speaker identity matching {DiarizedSpeaker} skipped because secondary validation rejected the unknown sample",
                request.DiarizedSpeaker);
            return null;
        }

        logger.LogInformation(
            "Speaker identity matching {DiarizedSpeaker} against {CandidateCount} candidate(s)",
            request.DiarizedSpeaker,
            request.Candidates.Count);
        var tempPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-speaker-match", $"{Guid.NewGuid():N}.wav");
        try
        {
            Directory.CreateDirectory(Path.GetDirectoryName(tempPath)!);
            var layout = WriteCompositeWav(tempPath, request);
            if (layout.UnknownSegment is null || layout.KnownSegments.Count == 0)
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} skipped because no readable snippets were available",
                    request.DiarizedSpeaker);
                return null;
            }

            logger.LogInformation(
                "Speaker identity matching {DiarizedSpeaker} composite layout: unknown {UnknownStart}-{UnknownEnd}, {KnownSegmentCount} known segment(s) across identity ids {IdentityIds}",
                request.DiarizedSpeaker,
                layout.UnknownSegment.Value.Start,
                layout.UnknownSegment.Value.End,
                layout.KnownSegments.Count,
                string.Join(", ", layout.KnownSegments.Select(segment => segment.IdentityId).Distinct().Order()));
            var segments = await DiarizeWithTimeoutAsync(tempPath, cancellationToken);
            if (segments.Count == 0)
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} skipped because diarization returned no segments",
                    request.DiarizedSpeaker);
                return null;
            }

            logger.LogInformation(
                "Speaker identity matching {DiarizedSpeaker} received {SegmentCount} diarized segment(s)",
                request.DiarizedSpeaker,
                segments.Count);
            foreach (var segment in segments.OrderBy(segment => segment.Start))
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} diarized turn {Start}-{End} as {Speaker}",
                    request.DiarizedSpeaker,
                    segment.Start,
                    segment.End,
                    segment.Speaker);
            }

            var unknownSpeaker = FindBestSpeaker(layout.UnknownSegment.Value, segments);
            if (string.IsNullOrWhiteSpace(unknownSpeaker))
            {
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} found no speaker for the unknown snippet",
                    request.DiarizedSpeaker);
                return null;
            }

            foreach (var known in layout.KnownSegments.GroupBy(segment => segment.IdentityId))
            {
                var matchingSnippetCount = known.Count(segment =>
                    string.Equals(
                        FindBestSpeaker(segment.Segment, segments),
                        unknownSpeaker,
                        StringComparison.Ordinal));
                var requiredMatches = known.Count() > 1 ? 2 : 1;
                logger.LogInformation(
                    "Speaker identity matching {DiarizedSpeaker} compared unknown speaker {UnknownSpeaker} with identity {IdentityId}: {MatchingSnippetCount}/{KnownSnippetCount} matching known snippet(s), required {RequiredMatches}",
                    request.DiarizedSpeaker,
                    unknownSpeaker,
                    known.Key,
                    matchingSnippetCount,
                    known.Count(),
                    requiredMatches);
                if (matchingSnippetCount >= requiredMatches)
                {
                    var validationRequest = new SpeakerIdentityMatchValidationRequest(
                        request.DiarizedSpeaker,
                        known.Key,
                        request.UnknownSnippet,
                        known.Select(segment => segment.Snippet).ToList());
                    if (!await matchValidator.ValidateMatchAsync(validationRequest, cancellationToken))
                    {
                        logger.LogInformation(
                            "Speaker identity matching {DiarizedSpeaker} rejected identity {IdentityId} after secondary validation",
                            request.DiarizedSpeaker,
                            known.Key);
                        continue;
                    }

                    logger.LogInformation(
                        "Speaker identity matching {DiarizedSpeaker} matched identity {IdentityId}",
                        request.DiarizedSpeaker,
                        known.Key);
                    return new SpeakerIdentityMatch(known.Key);
                }
            }

            logger.LogInformation(
                "Speaker identity matching {DiarizedSpeaker} found no matching identity",
                request.DiarizedSpeaker);
            return null;
        }
        finally
        {
            SpeakerCompositeWav.TryDelete(tempPath);
        }
    }

    private async Task<IReadOnlyList<TranscriptionSegment>> DiarizeWithTimeoutAsync(
        string wavPath,
        CancellationToken cancellationToken)
    {
        var matchTimeout = options.MatchTimeout;
        if (matchTimeout <= TimeSpan.Zero)
        {
            return await diarizationClient.DiarizeAsync(wavPath, cancellationToken);
        }

        using var timeout = new CancellationTokenSource(matchTimeout);
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeout.Token);
        try
        {
            return await diarizationClient.DiarizeAsync(wavPath, linked.Token);
        }
        catch (OperationCanceledException) when (timeout.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
        {
            logger.LogWarning(
                "Speaker identity diarization timed out after {Timeout}",
                matchTimeout);
            return [];
        }
    }

    private CompositeLayout WriteCompositeWav(string path, SpeakerIdentityMatchRequest request)
    {
        using var firstReader = SpeakerCompositeWav.OpenFirstReadableWave(
            request.Candidates.SelectMany(candidate => candidate.Snippets).Append(request.UnknownSnippet));
        if (firstReader is null)
        {
            return new CompositeLayout(null, []);
        }

        var silence = TimeSpan.FromSeconds(Math.Max(0, options.SilenceBetweenSnippetsSeconds));
        using var writer = new WaveFileWriter(path, firstReader.WaveFormat);
        var current = TimeSpan.Zero;
        var knownSegments = new List<KnownCompositeSegment>();

        foreach (var candidate in request.Candidates)
        {
            foreach (var snippet in candidate.Snippets.Where(storedSnippet => storedSnippet.Length > 0))
            {
                var segment = SpeakerCompositeWav.AppendSnippet(writer, firstReader.WaveFormat, snippet, current);
                if (segment is null)
                {
                    continue;
                }

                knownSegments.Add(new KnownCompositeSegment(candidate.IdentityId, segment.Value, snippet));
                current = segment.Value.End + silence;
                SpeakerCompositeWav.WriteSilence(writer, firstReader.WaveFormat, silence);
            }
        }

        var unknownSegment = SpeakerCompositeWav.AppendSnippet(
            writer,
            firstReader.WaveFormat,
            request.UnknownSnippet,
            current);
        return new CompositeLayout(unknownSegment, knownSegments);
    }

    private static string? FindBestSpeaker(TimeSegment segment, IReadOnlyList<TranscriptionSegment> segments)
    {
        return SpeakerCompositeWav.FindBestSpeaker(segment, segments);
    }

    private sealed record CompositeLayout(TimeSegment? UnknownSegment, IReadOnlyList<KnownCompositeSegment> KnownSegments);

    private sealed record KnownCompositeSegment(int IdentityId, TimeSegment Segment, byte[] Snippet);

}