meeting-assistant/MeetingAssistant/Speakers/VenvResemblyzerVoiceEncoder.cs

using System.Globalization;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.RegularExpressions;
using MeetingAssistant.Transcription;
using Microsoft.Extensions.Options;

namespace MeetingAssistant.Speakers;

public interface IResemblyzerVoiceEncoder
{
    Task<IReadOnlyList<float[]>> EncodeAsync(
        IReadOnlyList<byte[]> wavSamples,
        CancellationToken cancellationToken);

    Task WarmUpAsync(CancellationToken cancellationToken);
}

public sealed partial class VenvResemblyzerVoiceEncoder : IResemblyzerVoiceEncoder
{
    private const string EnvironmentSchemaVersion = "venv-v2";
    private const string JsonStart = "__MEETING_ASSISTANT_RESEMBLYZER_JSON_START__";
    private const string JsonEnd = "__MEETING_ASSISTANT_RESEMBLYZER_JSON_END__";
    private const string NumpyBeforePython313Requirement = "numpy<2; python_version < '3.13'";
    private const string NumpyFromPython313Requirement = "numpy>=2,<3; python_version >= '3.13'";
    private const string LibrosaRequirement = "librosa>=0.9.1";
    private const string ScipyRequirement = "scipy>=1.2.1";
    private readonly SemaphoreSlim commandLock = new(1, 1);
    private readonly ICommandRunner commandRunner;
    private readonly ResemblyzerSpeakerRecognitionOptions options;
    private readonly ILogger<VenvResemblyzerVoiceEncoder> logger;
    private string? verifiedEnvironmentPythonPath;

    public VenvResemblyzerVoiceEncoder(
        ICommandRunner commandRunner,
        IOptions<MeetingAssistantOptions> options,
        ILogger<VenvResemblyzerVoiceEncoder> logger)
    {
        this.commandRunner = commandRunner;
        this.options = options.Value.SpeakerIdentification.Resemblyzer;
        this.logger = logger;
    }

    public async Task<IReadOnlyList<float[]>> EncodeAsync(
        IReadOnlyList<byte[]> wavSamples,
        CancellationToken cancellationToken)
    {
        if (wavSamples.Count == 0)
        {
            return [];
        }

        if (wavSamples.Any(sample => sample.Length == 0))
        {
            throw new InvalidDataException("Resemblyzer cannot encode an empty WAV sample.");
        }

        var runtimeFolder = VaultPath.Resolve(options.RuntimeFolder);
        var inputFolder = Path.Combine(runtimeFolder, "input", Guid.NewGuid().ToString("N"));
        Directory.CreateDirectory(inputFolder);
        try
        {
            for (var index = 0; index < wavSamples.Count; index++)
            {
                await File.WriteAllBytesAsync(
                    Path.Combine(inputFolder, $"{index:D4}.wav"),
                    wavSamples[index],
                    cancellationToken);
            }

            await commandLock.WaitAsync(cancellationToken);
            try
            {
                return await RunEncodingAsync(inputFolder, wavSamples.Count, cancellationToken);
            }
            finally
            {
                commandLock.Release();
            }
        }
        finally
        {
            try
            {
                Directory.Delete(inputFolder, recursive: true);
            }
            catch (DirectoryNotFoundException)
            {
            }
            catch (IOException exception)
            {
                logger.LogWarning(exception, "Could not remove Resemblyzer temporary input folder {InputFolder}", inputFolder);
            }
            catch (UnauthorizedAccessException exception)
            {
                logger.LogWarning(exception, "Could not remove Resemblyzer temporary input folder {InputFolder}", inputFolder);
            }
        }
    }

    private async Task<IReadOnlyList<float[]>> RunEncodingAsync(
        string inputFolder,
        int expectedCount,
        CancellationToken cancellationToken)
    {
        return await RunWithTimeoutAsync(async token =>
        {
            var pythonPath = await EnsureEnvironmentAsync(token);
            var scriptPath = Path.Combine(VaultPath.Resolve(options.RuntimeFolder), "encode.py");
            await File.WriteAllTextAsync(scriptPath, BuildEncodingScript(), token);
            var result = await RunRequiredAsync(
                pythonPath,
                [scriptPath, Path.GetFullPath(inputFolder)],
                "encoding",
                token);
            return ParseAndValidateVectors(result.StandardOutput, expectedCount);
        }, cancellationToken);
    }

    public async Task WarmUpAsync(CancellationToken cancellationToken)
    {
        await commandLock.WaitAsync(cancellationToken);
        try
        {
            await RunWithTimeoutAsync(async token =>
            {
                await EnsureEnvironmentAsync(token);
                return true;
            }, cancellationToken);
        }
        finally
        {
            commandLock.Release();
        }
    }

    private async Task<string> EnsureEnvironmentAsync(CancellationToken cancellationToken)
    {
        ValidateDependencySettings();
        var runtimeFolder = VaultPath.Resolve(options.RuntimeFolder);
        var environmentFolder = Path.Combine(runtimeFolder, "venv", BuildEnvironmentKey());
        var pythonPath = GetEnvironmentPythonPath(environmentFolder);
        var readyPath = Path.Combine(environmentFolder, ".ready");
        if (string.Equals(verifiedEnvironmentPythonPath, pythonPath, StringComparison.OrdinalIgnoreCase)
            && File.Exists(pythonPath)
            && File.Exists(readyPath))
        {
            return pythonPath;
        }

        if (File.Exists(pythonPath) && File.Exists(readyPath))
        {
            var verification = await commandRunner.RunAsync(
                pythonPath,
                WarmupArguments,
                cancellationToken);
            if (verification.ExitCode == 0)
            {
                verifiedEnvironmentPythonPath = pythonPath;
                return pythonPath;
            }

            logger.LogWarning(
                "Existing Resemblyzer environment verification failed with exit code {ExitCode}: {Error}. Reprovisioning it.",
                verification.ExitCode,
                verification.StandardError);
            File.Delete(readyPath);
        }

        Directory.CreateDirectory(environmentFolder);
        await RunRequiredAsync(
            options.PythonCommand,
            ["-m", "venv", "--clear", environmentFolder],
            "virtual environment creation",
            cancellationToken);
        await RunRequiredAsync(
            pythonPath,
            ["-m", "pip", "install", "--upgrade", "pip"],
            "pip upgrade",
            cancellationToken);
        await RunRequiredAsync(
            pythonPath,
            [
                "-m", "pip", "install",
                "--index-url", options.TorchIndexUrl,
                TorchRequirement
            ],
            "CPU PyTorch installation",
            cancellationToken);
        await RunRequiredAsync(
            pythonPath,
            [
                "-m", "pip", "install",
                .. RuntimeDependencyRequirements
            ],
            "Resemblyzer dependency installation",
            cancellationToken);
        await RunRequiredAsync(
            pythonPath,
            [
                "-m", "pip", "install", "--no-deps",
                ResemblyzerRequirement
            ],
            "Resemblyzer installation",
            cancellationToken);
        await RunRequiredAsync(
            pythonPath,
            WarmupArguments,
            "environment verification",
            cancellationToken);
        await File.WriteAllTextAsync(readyPath, BuildEnvironmentKey(), cancellationToken);
        verifiedEnvironmentPythonPath = pythonPath;
        return pythonPath;
    }

    private async Task<CommandResult> RunRequiredAsync(
        string fileName,
        IReadOnlyList<string> arguments,
        string operation,
        CancellationToken cancellationToken)
    {
        var result = await commandRunner.RunAsync(fileName, arguments, cancellationToken);
        ThrowIfFailed(result, operation);
        return result;
    }

    private static string BuildEncodingScript()
    {
        return
            "import json\n"
            + "import sys\n"
            + "import numpy as np\n"
            + "from pathlib import Path\n"
            + "from resemblyzer import VoiceEncoder, preprocess_wav\n"
            + "encoder = VoiceEncoder('cpu', verbose=False)\n"
            + "vectors = []\n"
            + "for path in sorted(Path(sys.argv[1]).glob('*.wav')):\n"
            + "    wav = preprocess_wav(path)\n"
            + "    vector = encoder.embed_utterance(wav)\n"
            + "    vectors.append(np.asarray(vector, dtype=np.float32).tolist())\n"
            + $"print('{JsonStart}')\n"
            + "print(json.dumps(vectors, allow_nan=False))\n"
            + $"print('{JsonEnd}')\n";
    }

    private static IReadOnlyList<float[]> ParseAndValidateVectors(string output, int expectedCount)
    {
        var json = ExtractJson(output);
        float[][]? vectors;
        try
        {
            vectors = JsonSerializer.Deserialize<float[][]>(json);
        }
        catch (JsonException exception)
        {
            throw new InvalidDataException("Resemblyzer returned malformed vector JSON.", exception);
        }

        if (vectors is null || vectors.Length != expectedCount)
        {
            throw new InvalidDataException(
                $"Resemblyzer returned {vectors?.Length ?? 0} vectors for {expectedCount} WAV samples.");
        }

        return vectors
            .Select((vector, index) => SpeakerVoiceVectors.Normalize(vector, $"Resemblyzer vector {index}"))
            .ToList();
    }

    private static string ExtractJson(string output)
    {
        var start = output.IndexOf(JsonStart, StringComparison.Ordinal);
        if (start < 0)
        {
            throw new InvalidDataException("Resemblyzer output did not contain the JSON start marker.");
        }

        start += JsonStart.Length;
        var end = output.IndexOf(JsonEnd, start, StringComparison.Ordinal);
        if (end < 0)
        {
            throw new InvalidDataException("Resemblyzer output did not contain the JSON end marker.");
        }

        return output[start..end].Trim();
    }

    private string BuildEnvironmentKey()
    {
        var settings = string.Join('\n', EnvironmentManifest);
        var hash = SHA256.HashData(Encoding.UTF8.GetBytes(settings));
        return Convert.ToHexString(hash)[..16].ToLowerInvariant();
    }

    private string TorchRequirement => $"torch=={options.TorchVersion}";

    private string ResemblyzerRequirement => $"Resemblyzer=={options.PackageVersion}";

    private static string[] WarmupArguments =>
    [
        "-c",
        "from resemblyzer import VoiceEncoder; VoiceEncoder('cpu', verbose=False); print('Resemblyzer warm-up complete')"
    ];

    private string[] RuntimeDependencyRequirements =>
    [
        NumpyBeforePython313Requirement,
        NumpyFromPython313Requirement,
        LibrosaRequirement,
        ScipyRequirement,
        $"webrtcvad-wheels=={options.WebRtcVadVersion}"
    ];

    private IEnumerable<string> EnvironmentManifest =>
    [
        EnvironmentSchemaVersion,
        options.PythonCommand,
        options.TorchIndexUrl,
        TorchRequirement,
        .. RuntimeDependencyRequirements,
        ResemblyzerRequirement
    ];

    private static string GetEnvironmentPythonPath(string environmentFolder)
    {
        return OperatingSystem.IsWindows()
            ? Path.Combine(environmentFolder, "Scripts", "python.exe")
            : Path.Combine(environmentFolder, "bin", "python");
    }

    private void ValidateDependencySettings()
    {
        if (string.IsNullOrWhiteSpace(options.PythonCommand))
        {
            throw new InvalidOperationException("The Resemblyzer Python command cannot be empty.");
        }

        ValidateVersion(options.PackageVersion, "Resemblyzer package");
        ValidateVersion(options.TorchVersion, "PyTorch");
        ValidateVersion(options.WebRtcVadVersion, "webrtcvad-wheels");
        if (!Uri.TryCreate(options.TorchIndexUrl, UriKind.Absolute, out var indexUri)
            || indexUri.Scheme != Uri.UriSchemeHttps)
        {
            throw new InvalidOperationException(
                $"Invalid Resemblyzer PyTorch index URL '{options.TorchIndexUrl}'.");
        }
    }

    private static void ValidateVersion(string version, string dependency)
    {
        if (!PackageVersionPattern().IsMatch(version))
        {
            throw new InvalidOperationException($"Invalid {dependency} version '{version}'.");
        }
    }

    private async Task<T> RunWithTimeoutAsync<T>(
        Func<CancellationToken, Task<T>> operation,
        CancellationToken cancellationToken)
    {
        using var timeoutSource = options.CommandTimeout > TimeSpan.Zero
            ? new CancellationTokenSource(options.CommandTimeout)
            : null;
        using var linkedSource = timeoutSource is null
            ? null
            : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token);
        try
        {
            return await operation(linkedSource?.Token ?? cancellationToken);
        }
        catch (OperationCanceledException) when (
            !cancellationToken.IsCancellationRequested && timeoutSource?.IsCancellationRequested == true)
        {
            throw new TimeoutException(
                $"Resemblyzer command timed out after {options.CommandTimeout.ToString(null, CultureInfo.InvariantCulture)}.");
        }
    }

    private static void ThrowIfFailed(CommandResult result, string operation)
    {
        if (result.ExitCode != 0)
        {
            throw new InvalidOperationException(
                $"Resemblyzer {operation} failed with exit code {result.ExitCode}: {result.StandardError}");
        }
    }

    [GeneratedRegex("^[0-9A-Za-z.+-]+$", RegexOptions.CultureInvariant)]
    private static partial Regex PackageVersionPattern();
}