meeting-assistant/MeetingAssistant/Transcription/FunAsrStreamingTranscriptionProvider.cs

using System.Text.Json;
using System.Threading.Channels;
using MeetingAssistant.Recording;
using Microsoft.Extensions.Options;

namespace MeetingAssistant.Transcription;

public sealed class FunAsrStreamingTranscriptionProvider : IStreamingTranscriptionProvider, IAsyncDisposable
{
    private readonly IFunAsrWebSocketConnectionFactory connectionFactory;
    private readonly IFunAsrBackendLifecycle backendLifecycle;
    private readonly MeetingAssistantOptions options;
    private readonly ILogger<FunAsrStreamingTranscriptionProvider> logger;

    public FunAsrStreamingTranscriptionProvider(
        IFunAsrWebSocketConnectionFactory connectionFactory,
        IFunAsrBackendLifecycle backendLifecycle,
        IOptions<MeetingAssistantOptions> options,
        ILogger<FunAsrStreamingTranscriptionProvider> logger)
    {
        this.connectionFactory = connectionFactory;
        this.backendLifecycle = backendLifecycle;
        this.options = options.Value;
        this.logger = logger;
    }

    public async IAsyncEnumerable<TranscriptionSegment> TranscribeAsync(
        IAsyncEnumerable<AudioChunk> audio,
        SpeechRecognitionPipelineOptions options,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var segments = Channel.CreateUnbounded<TranscriptionSegment>();
        var session = Task.Run(() => RunSessionAsync(audio, options, segments.Writer, cancellationToken), CancellationToken.None);

        await foreach (var segment in segments.Reader.ReadAllAsync(cancellationToken))
        {
            yield return segment;
        }

        await session;
    }

    private async Task RunSessionAsync(
        IAsyncEnumerable<AudioChunk> audio,
        SpeechRecognitionPipelineOptions pipelineOptions,
        ChannelWriter<TranscriptionSegment> segments,
        CancellationToken cancellationToken)
    {
        try
        {
            await backendLifecycle.EnsureStartedAsync(cancellationToken);
            await using var enumerator = audio.GetAsyncEnumerator(cancellationToken);
            if (!await enumerator.MoveNextAsync())
            {
                segments.TryComplete();
                return;
            }

            var firstChunk = enumerator.Current;
            var endpoint = new Uri(options.FunAsr.Endpoint);
            await using var connection = await connectionFactory.ConnectAsync(endpoint, cancellationToken);
            var parser = new FunAsrMessageParser(options.FunAsr);
            var finalReceived = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
            var receiveTask = ReceiveAsync(connection, parser, segments, finalReceived, cancellationToken);

            await SendStartAsync(connection, firstChunk, pipelineOptions.DictationWords, cancellationToken);
            await connection.SendBinaryAsync(firstChunk.Pcm, cancellationToken);

            while (await enumerator.MoveNextAsync())
            {
                await connection.SendBinaryAsync(enumerator.Current.Pcm, cancellationToken);
            }

            await connection.SendTextAsync("""{"is_speaking":false}""", cancellationToken);

            try
            {
                await finalReceived.Task.WaitAsync(options.FunAsr.FinalResultTimeout, cancellationToken);
            }
            catch (TimeoutException)
            {
                logger.LogWarning(
                    "Timed out waiting {Timeout} for the final FunASR result from {Endpoint}",
                    options.FunAsr.FinalResultTimeout,
                    endpoint);
            }

            await connection.CloseOutputAsync(cancellationToken);
            await receiveTask;
            segments.TryComplete();
        }
        catch (Exception exception)
        {
            segments.TryComplete(exception);
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (backendLifecycle is IAsyncDisposable disposable)
        {
            await disposable.DisposeAsync();
        }
    }

    private async Task ReceiveAsync(
        IFunAsrWebSocketConnection connection,
        FunAsrMessageParser parser,
        ChannelWriter<TranscriptionSegment> segments,
        TaskCompletionSource finalReceived,
        CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            var message = await connection.ReceiveTextAsync(cancellationToken);
            if (message is null)
            {
                finalReceived.TrySetResult();
                return;
            }

            var parsed = parser.Parse(message);
            foreach (var segment in parsed.Segments)
            {
                await segments.WriteAsync(segment, cancellationToken);
                logger.LogInformation(
                    "FunASR emitted segment at {Start} for {Speaker}: {Text}",
                    segment.Start,
                    segment.Speaker,
                    segment.Text);
            }

            if (parsed.IsTerminal)
            {
                finalReceived.TrySetResult();
                return;
            }
        }
    }

    private async Task SendStartAsync(
        IFunAsrWebSocketConnection connection,
        AudioChunk firstChunk,
        IReadOnlyList<string>? dictationWords,
        CancellationToken cancellationToken)
    {
        var message = new Dictionary<string, object?>
        {
            ["mode"] = options.FunAsr.Mode,
            ["chunk_size"] = NormalizeChunkSize(options.FunAsr.ChunkSize),
            ["chunk_interval"] = options.FunAsr.ChunkInterval,
            ["encoder_chunk_look_back"] = options.FunAsr.EncoderChunkLookBack,
            ["decoder_chunk_look_back"] = options.FunAsr.DecoderChunkLookBack,
            ["audio_fs"] = firstChunk.SampleRate,
            ["wav_name"] = options.FunAsr.WavName,
            ["wav_format"] = "pcm",
            ["is_speaking"] = true,
            ["hotwords"] = BuildHotwords(dictationWords),
            ["itn"] = options.FunAsr.Itn
        };

        await connection.SendTextAsync(JsonSerializer.Serialize(message), cancellationToken);
    }

    internal static string BuildHotwords(IReadOnlyList<string>? dictationWords)
    {
        var hotwords = (dictationWords ?? [])
            .Select(word => word.Trim())
            .Where(line => !string.IsNullOrWhiteSpace(line))
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToDictionary(word => word, _ => 20);

        return hotwords.Count == 0 ? "" : JsonSerializer.Serialize(hotwords);
    }

    private static int[] NormalizeChunkSize(int[] configuredChunkSize)
    {
        return configuredChunkSize.Length == 3 ? configuredChunkSize : [5, 10, 5];
    }
}