meeting-assistant/MeetingAssistant/Speakers/ResemblyzerSpeakerIdentityMergeService.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Options;

namespace MeetingAssistant.Speakers;

public sealed class ResemblyzerSpeakerIdentityMergeService : ISpeakerIdentityMergeService
{
    private readonly IDbContextFactory<SpeakerIdentityDbContext> dbContextFactory;
    private readonly ResemblyzerVoiceClusterMatcher matcher;
    private readonly ResemblyzerVoiceVectorOutlierPruner outlierPruner;
    private readonly SpeakerIdentificationOptions options;
    private readonly ResemblyzerSpeakerRecognitionOptions resemblyzerOptions;
    private readonly ILogger<ResemblyzerSpeakerIdentityMergeService> logger;

    public ResemblyzerSpeakerIdentityMergeService(
        IDbContextFactory<SpeakerIdentityDbContext> dbContextFactory,
        ResemblyzerVoiceClusterMatcher matcher,
        ResemblyzerVoiceVectorOutlierPruner outlierPruner,
        IOptions<MeetingAssistantOptions> options,
        ILogger<ResemblyzerSpeakerIdentityMergeService> logger)
    {
        this.dbContextFactory = dbContextFactory;
        this.matcher = matcher;
        this.outlierPruner = outlierPruner;
        this.options = options.Value.SpeakerIdentification;
        resemblyzerOptions = this.options.Resemblyzer;
        this.logger = logger;
    }

    public async Task<SpeakerIdentityMergeResult> MergeRecentIdentitiesAsync(
        TimeSpan? recentIdentityAge,
        CancellationToken cancellationToken)
    {
        await using var context = await dbContextFactory.CreateDbContextAsync(cancellationToken);
        await SpeakerIdentitySchema.EnsureCreatedOrUpdatedAsync(context, cancellationToken);
        var identities = await context.SpeakerIdentities
            .AsSplitQuery()
            .Include(identity => identity.Aliases)
            .Include(identity => identity.CandidateNames)
            .Include(identity => identity.Snippets)
            .Include(identity => identity.VoiceVectors)
            .Include(identity => identity.References)
            .OrderByDescending(identity => identity.References.Count)
            .ThenBy(identity => identity.Id)
            .ToListAsync(cancellationToken);
        var cutoff = DateTimeOffset.UtcNow - (recentIdentityAge ?? options.MergeRecentIdentityAge);
        var recentIds = identities
            .Where(identity => identity.CreatedAt >= cutoff)
            .Select(identity => identity.Id)
            .ToList();
        var required = resemblyzerOptions.RequiredVectorsPerSpeaker;
        var attempts = 0;
        var mergedPairs = 0;
        var pendingAudits = new List<PendingMergeAudit>();

        foreach (var sourceId in recentIds)
        {
            var source = identities.SingleOrDefault(identity => identity.Id == sourceId);
            if (source is null)
            {
                continue;
            }

            var sourceVectors = SpeakerVoiceVectors.DecodeCompatible(
                    source,
                    resemblyzerOptions.ModelId,
                    logger)
                .Take(required * 2)
                .ToList();
            if (sourceVectors.Count < required * 2)
            {
                logger.LogInformation(
                    "Skipping Resemblyzer merge source identity {SourceIdentityId}: {VectorCount}/{RequiredVectorCount} compatible vectors",
                    source.Id,
                    sourceVectors.Count,
                    required * 2);
                continue;
            }

            var candidates = identities
                .Where(identity => identity.Id != source.Id)
                .Where(identity => identity.VoiceVectors.Any(vector =>
                    string.Equals(vector.ModelId, resemblyzerOptions.ModelId, StringComparison.Ordinal)))
                .Take(Math.Max(1, options.MaxMatchCandidates))
                .Select(identity => new ResemblyzerVoiceVectorCandidate(
                    identity.Id,
                    SpeakerVoiceVectors.DecodeCompatible(
                        identity,
                        resemblyzerOptions.ModelId,
                        logger)))
                .Where(candidate => candidate.Vectors.Count > 0)
                .ToList();
            if (candidates.Count == 0)
            {
                continue;
            }

            attempts++;
            var first = matcher.Match(sourceVectors.Take(required).ToList(), candidates);
            if (!first.Accepted || first.IdentityId is not { } targetId)
            {
                continue;
            }

            attempts++;
            var second = matcher.Match(sourceVectors.Skip(required).Take(required).ToList(), candidates);
            if (!second.Accepted || second.IdentityId != targetId)
            {
                logger.LogInformation(
                    "Rejected Resemblyzer merge for source identity {SourceIdentityId}: disjoint clusters selected {FirstIdentityId} and {SecondIdentityId}",
                    source.Id,
                    first.IdentityId,
                    second.IdentityId);
                continue;
            }

            var target = identities.SingleOrDefault(identity => identity.Id == targetId);
            if (target is null)
            {
                continue;
            }

            var targetName = target.GetDisplayName() ?? $"identity-{target.Id}";
            var sourceName = source.GetDisplayName() ?? $"identity-{source.Id}";
            SpeakerIdentityMerger.MergeIntoAndPrune(
                target,
                source,
                options.MaxSnippetsPerSpeaker,
                resemblyzerOptions.MaxVectorsPerIdentity,
                outlierPruner);
            pendingAudits.Add(new PendingMergeAudit(
                target,
                targetName,
                sourceName));
            context.SpeakerIdentities.Remove(source);
            identities.Remove(source);
            mergedPairs++;
        }

        await context.SaveChangesAsync(cancellationToken);
        foreach (var audit in pendingAudits)
        {
            try
            {
                await SpeakerIdentityTranscriptAudit.AppendMergedAsync(
                    audit.Identity.References,
                    audit.TargetName,
                    audit.SourceName,
                    cancellationToken);
            }
            catch (Exception exception) when (exception is not OperationCanceledException)
            {
                logger.LogError(
                    exception,
                    "Resemblyzer identity merge was saved for target {IdentityId}, but its transcript audit could not be written",
                    audit.Identity.Id);
            }
        }

        return new SpeakerIdentityMergeResult(
            recentIds.Count,
            identities.Count,
            attempts,
            mergedPairs);
    }

    private sealed record PendingMergeAudit(
        SpeakerIdentity Identity,
        string TargetName,
        string SourceName);

}