meeting-assistant/MeetingAssistant/Transcription/FunAsrMessageParser.cs

using System.Text.Json;

namespace MeetingAssistant.Transcription;

public sealed class FunAsrMessageParser
{
    private readonly FunAsrOptions options;

    public FunAsrMessageParser(FunAsrOptions options)
    {
        this.options = options;
    }

    public FunAsrParsedMessage Parse(string message)
    {
        using var document = JsonDocument.Parse(message);
        var root = document.RootElement;
        var mode = GetString(root, "mode");
        var isFinal = GetBoolean(root, "is_final");
        var shouldEmit = options.EmitOnlinePartials || !IsOnlineMode(mode);
        var segments = shouldEmit ? ReadSegments(root) : [];

        return new FunAsrParsedMessage(IsTerminal: isFinal && !IsOnlineMode(mode), segments);
    }

    private static List<TranscriptionSegment> ReadSegments(JsonElement root)
    {
        var sentenceSegments = ReadSentenceSegments(root, "sentence_info");
        if (sentenceSegments.Count > 0)
        {
            return sentenceSegments;
        }

        var stampSegments = ReadSentenceSegments(root, "stamp_sents");
        if (stampSegments.Count > 0)
        {
            return stampSegments;
        }

        var text = GetString(root, "text").Trim();
        if (string.IsNullOrWhiteSpace(text))
        {
            return [];
        }

        var (start, end) = ReadTimestampRange(root);
        return [new TranscriptionSegment(start, end, ReadSpeaker(root), text)];
    }

    private static List<TranscriptionSegment> ReadSentenceSegments(JsonElement root, string propertyName)
    {
        if (!root.TryGetProperty(propertyName, out var sentences) || sentences.ValueKind != JsonValueKind.Array)
        {
            return [];
        }

        var segments = new List<TranscriptionSegment>();
        foreach (var sentence in sentences.EnumerateArray())
        {
            var text = (GetString(sentence, "text") + GetString(sentence, "text_seg") + GetString(sentence, "punc")).Trim();
            if (string.IsNullOrWhiteSpace(text))
            {
                continue;
            }

            segments.Add(new TranscriptionSegment(
                TimeSpan.FromMilliseconds(GetDouble(sentence, "start")),
                TimeSpan.FromMilliseconds(GetDouble(sentence, "end")),
                ReadSpeaker(sentence),
                text));
        }

        return segments;
    }

    private static (TimeSpan Start, TimeSpan End) ReadTimestampRange(JsonElement root)
    {
        if (!root.TryGetProperty("timestamp", out var timestamp))
        {
            return (TimeSpan.Zero, TimeSpan.Zero);
        }

        if (timestamp.ValueKind == JsonValueKind.String)
        {
            var value = timestamp.GetString();
            if (string.IsNullOrWhiteSpace(value))
            {
                return (TimeSpan.Zero, TimeSpan.Zero);
            }

            using var document = JsonDocument.Parse(value);
            return ReadTimestampArray(document.RootElement);
        }

        return ReadTimestampArray(timestamp);
    }

    private static (TimeSpan Start, TimeSpan End) ReadTimestampArray(JsonElement timestamp)
    {
        if (timestamp.ValueKind != JsonValueKind.Array)
        {
            return (TimeSpan.Zero, TimeSpan.Zero);
        }

        double? start = null;
        double? end = null;
        foreach (var item in timestamp.EnumerateArray())
        {
            if (item.ValueKind != JsonValueKind.Array)
            {
                continue;
            }

            var values = item.EnumerateArray().ToArray();
            if (values.Length < 2)
            {
                continue;
            }

            start ??= ReadNumber(values[0]);
            end = ReadNumber(values[1]);
        }

        return (
            TimeSpan.FromMilliseconds(start ?? 0),
            TimeSpan.FromMilliseconds(end ?? start ?? 0));
    }

    private static string ReadSpeaker(JsonElement element)
    {
        var speaker = GetString(element, "spk_name");
        if (!string.IsNullOrWhiteSpace(speaker))
        {
            return speaker;
        }

        speaker = GetString(element, "speaker");
        if (!string.IsNullOrWhiteSpace(speaker))
        {
            return speaker;
        }

        if (element.TryGetProperty("spk", out var spk))
        {
            return $"Speaker {spk}";
        }

        return "Unknown";
    }

    private static bool IsOnlineMode(string mode)
    {
        return mode.Equals("online", StringComparison.OrdinalIgnoreCase)
            || mode.Equals("2pass-online", StringComparison.OrdinalIgnoreCase);
    }

    private static string GetString(JsonElement element, string propertyName)
    {
        return element.TryGetProperty(propertyName, out var property) && property.ValueKind == JsonValueKind.String
            ? property.GetString() ?? ""
            : "";
    }

    private static bool GetBoolean(JsonElement element, string propertyName)
    {
        return element.TryGetProperty(propertyName, out var property)
            && property.ValueKind is JsonValueKind.True or JsonValueKind.False
            && property.GetBoolean();
    }

    private static double GetDouble(JsonElement element, string propertyName)
    {
        return element.TryGetProperty(propertyName, out var property) ? ReadNumber(property) : 0;
    }

    private static double ReadNumber(JsonElement element)
    {
        return element.ValueKind == JsonValueKind.Number && element.TryGetDouble(out var value) ? value : 0;
    }
}

public sealed record FunAsrParsedMessage(bool IsTerminal, IReadOnlyList<TranscriptionSegment> Segments);