meeting-assistant/MeetingAssistant/Transcription/StreamingSpeechRecognitionPipeline.cs

using System.Threading.Channels;
using MeetingAssistant.Recording;

namespace MeetingAssistant.Transcription;

public abstract class StreamingSpeechRecognitionPipeline : ISpeechRecognitionPipeline
{
    private readonly IStreamingTranscriptionProvider transcriptionProvider;
    private readonly Channel<AudioChunk> audio = Channel.CreateUnbounded<AudioChunk>();
    private readonly Channel<TranscriptionSegment> liveTranscript = Channel.CreateUnbounded<TranscriptionSegment>();
    private readonly List<TranscriptionSegment> liveSegments = [];
    private SpeechRecognitionPipelineOptions options = SpeechRecognitionPipelineOptions.Default;
    private Task? transcriptionTask;
    private bool initialized;

    protected StreamingSpeechRecognitionPipeline(IStreamingTranscriptionProvider transcriptionProvider)
    {
        this.transcriptionProvider = transcriptionProvider;
    }

    public virtual Task InitializeAsync(CancellationToken cancellationToken)
    {
        return InitializeAsync(SpeechRecognitionPipelineOptions.Default, cancellationToken);
    }

    public virtual Task InitializeAsync(
        SpeechRecognitionPipelineOptions options,
        CancellationToken cancellationToken)
    {
        if (initialized)
        {
            return Task.CompletedTask;
        }

        initialized = true;
        this.options = options;
        transcriptionTask = Task.Run(() => TranscribeAsync(cancellationToken), CancellationToken.None);
        return Task.CompletedTask;
    }

    public virtual Task WaitUntilReadyAsync(CancellationToken cancellationToken)
    {
        return Task.CompletedTask;
    }

    public ValueTask WriteAsync(AudioChunk chunk, CancellationToken cancellationToken)
    {
        EnsureInitialized();
        return audio.Writer.WriteAsync(chunk, cancellationToken);
    }

    public async Task CompleteAsync(CancellationToken cancellationToken)
    {
        EnsureInitialized();
        audio.Writer.TryComplete();
        if (transcriptionTask is not null)
        {
            await transcriptionTask.WaitAsync(cancellationToken);
        }
    }

    public async IAsyncEnumerable<TranscriptionSegment> ReadLiveTranscriptAsync(
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        EnsureInitialized();
        await foreach (var segment in liveTranscript.Reader.ReadAllAsync(cancellationToken))
        {
            yield return segment;
        }
    }

    public async Task<IReadOnlyList<TranscriptionSegment>> ReadFinishedTranscriptAsync(
        string audioPath,
        SpeechRecognitionPipelineOptions options,
        CancellationToken cancellationToken)
    {
        EnsureInitialized();
        await CompleteAsync(cancellationToken);
        var finalSegments = await BuildFinishedTranscriptAsync(audioPath, liveSegments, options, cancellationToken);
        return finalSegments.Count == 0 ? liveSegments : finalSegments;
    }

    public virtual ValueTask DisposeAsync()
    {
        audio.Writer.TryComplete();
        liveTranscript.Writer.TryComplete();
        return ValueTask.CompletedTask;
    }

    protected abstract Task<IReadOnlyList<TranscriptionSegment>> BuildFinishedTranscriptAsync(
        string audioPath,
        IReadOnlyList<TranscriptionSegment> liveSegments,
        SpeechRecognitionPipelineOptions options,
        CancellationToken cancellationToken);

    private async Task TranscribeAsync(CancellationToken cancellationToken)
    {
        try
        {
            await foreach (var segment in transcriptionProvider.TranscribeAsync(
                audio.Reader.ReadAllAsync(cancellationToken),
                options,
                cancellationToken))
            {
                liveSegments.Add(segment);
                await liveTranscript.Writer.WriteAsync(segment, cancellationToken);
            }

            liveTranscript.Writer.TryComplete();
        }
        catch (Exception exception)
        {
            liveTranscript.Writer.TryComplete(exception);
        }
    }

    private void EnsureInitialized()
    {
        if (!initialized)
        {
            throw new InvalidOperationException("Speech recognition pipeline must be initialized before use.");
        }
    }
}