Added math rendering (#705)

app/MindWork AI Studio/Chat/ContentBlockComponent.razor

@@ -96,11 +96,25 @@
                             }
                             else
                             {
-                                <MudMarkdown Value="@NormalizeMarkdownForRendering(textContent.Text)" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
-                                @if (textContent.Sources.Count > 0)
-                                {
-                                    <MudMarkdown Value="@textContent.Sources.ToMarkdown()" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
-                                }
+                                var renderPlan = this.GetMarkdownRenderPlan(textContent.Text);
+                                <div @ref="this.mathContentContainer" class="chat-math-container">
+                                    @foreach (var segment in renderPlan.Segments)
+                                    {
+                                        var segmentContent = segment.GetContent(renderPlan.Source);
+                                        if (segment.Type is MarkdownRenderSegmentType.MARKDOWN)
+                                        {
+                                            <MudMarkdown @key="@segment.RenderKey" Value="@segmentContent" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
+                                        }
+                                        else
+                                        {
+                                            <MathJaxBlock @key="@segment.RenderKey" Value="@segmentContent" Class="mb-5" />
+                                        }
+                                    }
+                                    @if (textContent.Sources.Count > 0)
+                                    {
+                                        <MudMarkdown Value="@textContent.Sources.ToMarkdown()" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
+                                    }
+                                </div>
                             }
                         }
                     }
@@ -135,4 +149,4 @@
             }
         }
     </MudCardContent>
-</MudCard>
+</MudCard>

app/MindWork AI Studio/Chat/ContentBlockComponent.razor.cs

@@ -8,8 +8,20 @@ namespace AIStudio.Chat;
 /// <summary>
 /// The UI component for a chat content block, i.e., for any IContent.
 /// </summary>
-public partial class ContentBlockComponent : MSGComponentBase
+public partial class ContentBlockComponent : MSGComponentBase, IAsyncDisposable
 {
+    private const string CHAT_MATH_SYNC_FUNCTION = "chatMath.syncContainer";
+    private const string CHAT_MATH_DISPOSE_FUNCTION = "chatMath.disposeContainer";
+    private const string HTML_START_TAG = "<";
+    private const string HTML_END_TAG = "</";
+    private const string HTML_SELF_CLOSING_TAG = "/>";
+    private const string CODE_FENCE_MARKER_BACKTICK = "```";
+    private const string CODE_FENCE_MARKER_TILDE = "~~~";
+    private const string MATH_BLOCK_MARKER_DOLLAR = "$$";
+    private const string MATH_BLOCK_MARKER_BRACKET_OPEN = """\[""";
+    private const string MATH_BLOCK_MARKER_BRACKET_CLOSE = """\]""";
+    private const string HTML_CODE_FENCE_PREFIX = "```html";
+
     private static readonly string[] HTML_TAG_MARKERS =
     [
         "<!doctype",
@@ -79,9 +91,18 @@ public partial class ContentBlockComponent : MSGComponentBase
     [Inject]
     private RustService RustService { get; init; } = null!;
 
+    [Inject]
+    private IJSRuntime JsRuntime { get; init; } = null!;
+
     private bool HideContent { get; set; }
     private bool hasRenderHash;
     private int lastRenderHash;
+    private string cachedMarkdownRenderPlanInput = string.Empty;
+    private MarkdownRenderPlan cachedMarkdownRenderPlan = MarkdownRenderPlan.EMPTY;
+    private ElementReference mathContentContainer;
+    private string lastMathRenderSignature = string.Empty;
+    private bool hasActiveMathContainer;
+    private bool isDisposed;
 
     #region Overrides of ComponentBase
 
@@ -97,6 +118,12 @@ public partial class ContentBlockComponent : MSGComponentBase
         return base.OnParametersSetAsync();
     }
 
+    protected override async Task OnAfterRenderAsync(bool firstRender)
+    {
+        await this.SyncMathRenderIfNeededAsync();
+        await base.OnAfterRenderAsync(firstRender);
+    }
+
     /// <inheritdoc />
     protected override bool ShouldRender()
     {
@@ -194,32 +221,320 @@ public partial class ContentBlockComponent : MSGComponentBase
         CodeBlock = { Theme = this.CodeColorPalette },
     };
 
+    private MarkdownRenderPlan GetMarkdownRenderPlan(string text)
+    {
+        if (ReferenceEquals(this.cachedMarkdownRenderPlanInput, text) || string.Equals(this.cachedMarkdownRenderPlanInput, text, StringComparison.Ordinal))
+            return this.cachedMarkdownRenderPlan;
+
+        this.cachedMarkdownRenderPlanInput = text;
+        this.cachedMarkdownRenderPlan = BuildMarkdownRenderPlan(text);
+        return this.cachedMarkdownRenderPlan;
+    }
+
+    private async Task SyncMathRenderIfNeededAsync()
+    {
+        if (this.isDisposed)
+            return;
+
+        if (!this.TryGetCompletedMathRenderState(out var mathRenderSignature))
+        {
+            await this.DisposeMathContainerIfNeededAsync();
+            return;
+        }
+
+        if (string.Equals(this.lastMathRenderSignature, mathRenderSignature, StringComparison.Ordinal))
+            return;
+
+        await this.JsRuntime.InvokeVoidAsync(CHAT_MATH_SYNC_FUNCTION, this.mathContentContainer, mathRenderSignature);
+        this.lastMathRenderSignature = mathRenderSignature;
+        this.hasActiveMathContainer = true;
+    }
+
+    private async Task DisposeMathContainerIfNeededAsync()
+    {
+        if (!this.hasActiveMathContainer)
+        {
+            this.lastMathRenderSignature = string.Empty;
+            return;
+        }
+
+        try
+        {
+            await this.JsRuntime.InvokeVoidAsync(CHAT_MATH_DISPOSE_FUNCTION, this.mathContentContainer);
+        }
+        catch (JSDisconnectedException)
+        {
+        }
+        catch (ObjectDisposedException)
+        {
+        }
+
+        this.hasActiveMathContainer = false;
+        this.lastMathRenderSignature = string.Empty;
+    }
+
+    private bool TryGetCompletedMathRenderState(out string mathRenderSignature)
+    {
+        mathRenderSignature = string.Empty;
+
+        if (this.HideContent || this.Type is not ContentType.TEXT || this.Content.IsStreaming || this.Content is not ContentText textContent || textContent.InitialRemoteWait)
+            return false;
+
+        var renderPlan = this.GetMarkdownRenderPlan(textContent.Text);
+        mathRenderSignature = CreateMathRenderSignature(renderPlan);
+        return !string.IsNullOrEmpty(mathRenderSignature);
+    }
+
+    private static string CreateMathRenderSignature(MarkdownRenderPlan renderPlan)
+    {
+        var hash = new HashCode();
+        var mathSegmentCount = 0;
+
+        foreach (var segment in renderPlan.Segments)
+        {
+            if (segment.Type is not MarkdownRenderSegmentType.MATH_BLOCK)
+                continue;
+
+            mathSegmentCount++;
+            hash.Add(segment.Start);
+            hash.Add(segment.Length);
+            hash.Add(segment.GetContent(renderPlan.Source).GetHashCode(StringComparison.Ordinal));
+        }
+
+        return mathSegmentCount == 0
+            ? string.Empty
+            : $"{mathSegmentCount}:{hash.ToHashCode()}";
+    }
+
+    private static MarkdownRenderPlan BuildMarkdownRenderPlan(string text)
+    {
+        var normalized = NormalizeMarkdownForRendering(text);
+        if (string.IsNullOrWhiteSpace(normalized))
+            return MarkdownRenderPlan.EMPTY;
+
+        var normalizedSpan = normalized.AsSpan();
+        var segments = new List<MarkdownRenderSegment>();
+        var activeCodeFenceMarker = '\0';
+        var activeMathBlockFenceType = MathBlockFenceType.NONE;
+        var markdownSegmentStart = 0;
+        var mathContentStart = 0;
+
+        for (var lineStart = 0; lineStart < normalizedSpan.Length;)
+        {
+            var lineEnd = lineStart;
+            while (lineEnd < normalizedSpan.Length && normalizedSpan[lineEnd] is not '\r' and not '\n')
+                lineEnd++;
+
+            var nextLineStart = lineEnd;
+            if (nextLineStart < normalizedSpan.Length)
+            {
+                if (normalizedSpan[nextLineStart] == '\r')
+                    nextLineStart++;
+
+                if (nextLineStart < normalizedSpan.Length && normalizedSpan[nextLineStart] == '\n')
+                    nextLineStart++;
+            }
+
+            var trimmedLine = TrimWhitespace(normalizedSpan[lineStart..lineEnd]);
+            if (activeMathBlockFenceType is MathBlockFenceType.NONE && TryUpdateCodeFenceState(trimmedLine, ref activeCodeFenceMarker))
+            {
+                lineStart = nextLineStart;
+                continue;
+            }
+
+            if (activeCodeFenceMarker != '\0')
+            {
+                lineStart = nextLineStart;
+                continue;
+            }
+
+            if (activeMathBlockFenceType is MathBlockFenceType.NONE)
+            {
+                if (trimmedLine.SequenceEqual(MATH_BLOCK_MARKER_DOLLAR.AsSpan()))
+                {
+                    AddMarkdownSegment(markdownSegmentStart, lineStart);
+                    mathContentStart = nextLineStart;
+                    activeMathBlockFenceType = MathBlockFenceType.DOLLAR;
+                    lineStart = nextLineStart;
+                    continue;
+                }
+
+                if (trimmedLine.SequenceEqual(MATH_BLOCK_MARKER_BRACKET_OPEN.AsSpan()))
+                {
+                    AddMarkdownSegment(markdownSegmentStart, lineStart);
+                    mathContentStart = nextLineStart;
+                    activeMathBlockFenceType = MathBlockFenceType.BRACKET;
+                    lineStart = nextLineStart;
+                    continue;
+                }
+            }
+            else if (activeMathBlockFenceType is MathBlockFenceType.DOLLAR && trimmedLine.SequenceEqual(MATH_BLOCK_MARKER_DOLLAR.AsSpan()))
+            {
+                var (start, end) = TrimLineBreaks(normalizedSpan, mathContentStart, lineStart);
+                segments.Add(new(MarkdownRenderSegmentType.MATH_BLOCK, start, end - start));
+
+                markdownSegmentStart = nextLineStart;
+                activeMathBlockFenceType = MathBlockFenceType.NONE;
+                lineStart = nextLineStart;
+                continue;
+            }
+            else if (activeMathBlockFenceType is MathBlockFenceType.BRACKET && trimmedLine.SequenceEqual(MATH_BLOCK_MARKER_BRACKET_CLOSE.AsSpan()))
+            {
+                var (start, end) = TrimLineBreaks(normalizedSpan, mathContentStart, lineStart);
+                segments.Add(new(MarkdownRenderSegmentType.MATH_BLOCK, start, end - start));
+
+                markdownSegmentStart = nextLineStart;
+                activeMathBlockFenceType = MathBlockFenceType.NONE;
+                lineStart = nextLineStart;
+                continue;
+            }
+
+            lineStart = nextLineStart;
+        }
+
+        if (activeMathBlockFenceType is not MathBlockFenceType.NONE)
+            return new(normalized, [new(MarkdownRenderSegmentType.MARKDOWN, 0, normalized.Length)]);
+
+        AddMarkdownSegment(markdownSegmentStart, normalized.Length);
+        if (segments.Count == 0)
+            segments.Add(new(MarkdownRenderSegmentType.MARKDOWN, 0, normalized.Length));
+
+        return new(normalized, segments);
+
+        void AddMarkdownSegment(int start, int end)
+        {
+            if (end <= start)
+                return;
+
+            segments.Add(new(MarkdownRenderSegmentType.MARKDOWN, start, end - start));
+        }
+    }
+
     private static string NormalizeMarkdownForRendering(string text)
     {
-        var cleaned = text.RemoveThinkTags().Trim();
-        if (string.IsNullOrWhiteSpace(cleaned))
+        var textWithoutThinkTags = text.RemoveThinkTags();
+        var trimmed = TrimWhitespace(textWithoutThinkTags.AsSpan());
+        if (trimmed.IsEmpty)
             return string.Empty;
 
-        if (cleaned.Contains("```", StringComparison.Ordinal))
+        var cleaned = trimmed.Length == textWithoutThinkTags.Length
+            ? textWithoutThinkTags
+            : trimmed.ToString();
+
+        if (cleaned.Contains(CODE_FENCE_MARKER_BACKTICK, StringComparison.Ordinal))
             return cleaned;
 
         if (LooksLikeRawHtml(cleaned))
-            return $"```html{Environment.NewLine}{cleaned}{Environment.NewLine}```";
+            return $"{HTML_CODE_FENCE_PREFIX}{Environment.NewLine}{cleaned}{Environment.NewLine}{CODE_FENCE_MARKER_BACKTICK}";
 
         return cleaned;
     }
 
     private static bool LooksLikeRawHtml(string text)
     {
-        var content = text.TrimStart();
-        if (!content.StartsWith("<", StringComparison.Ordinal))
+        var content = text.AsSpan();
+        var start = 0;
+        while (start < content.Length && char.IsWhiteSpace(content[start]))
+            start++;
+
+        content = content[start..];
+        if (!content.StartsWith(HTML_START_TAG.AsSpan(), StringComparison.Ordinal))
             return false;
 
         foreach (var marker in HTML_TAG_MARKERS)
-            if (content.Contains(marker, StringComparison.OrdinalIgnoreCase))
+            if (content.IndexOf(marker.AsSpan(), StringComparison.OrdinalIgnoreCase) >= 0)
                 return true;
 
-        return content.Contains("</", StringComparison.Ordinal) || content.Contains("/>", StringComparison.Ordinal);
+        return content.IndexOf(HTML_END_TAG.AsSpan(), StringComparison.Ordinal) >= 0
+               || content.IndexOf(HTML_SELF_CLOSING_TAG.AsSpan(), StringComparison.Ordinal) >= 0;
+    }
+
+    private static bool TryUpdateCodeFenceState(ReadOnlySpan<char> trimmedLine, ref char activeCodeFenceMarker)
+    {
+        var fenceMarker = '\0';
+        if (trimmedLine.StartsWith(CODE_FENCE_MARKER_BACKTICK.AsSpan(), StringComparison.Ordinal))
+            fenceMarker = '`';
+        else if (trimmedLine.StartsWith(CODE_FENCE_MARKER_TILDE.AsSpan(), StringComparison.Ordinal))
+            fenceMarker = '~';
+
+        if (fenceMarker == '\0')
+            return false;
+
+        activeCodeFenceMarker = activeCodeFenceMarker == '\0'
+            ? fenceMarker
+            : activeCodeFenceMarker == fenceMarker
+                ? '\0'
+                : activeCodeFenceMarker;
+
+        return true;
+    }
+
+    private static ReadOnlySpan<char> TrimWhitespace(ReadOnlySpan<char> text)
+    {
+        var start = 0;
+        var end = text.Length - 1;
+
+        while (start < text.Length && char.IsWhiteSpace(text[start]))
+            start++;
+
+        while (end >= start && char.IsWhiteSpace(text[end]))
+            end--;
+
+        return start > end ? ReadOnlySpan<char>.Empty : text[start..(end + 1)];
+    }
+
+    private static (int Start, int End) TrimLineBreaks(ReadOnlySpan<char> text, int start, int end)
+    {
+        while (start < end && text[start] is '\r' or '\n')
+            start++;
+
+        while (end > start && text[end - 1] is '\r' or '\n')
+            end--;
+
+        return (start, end);
+    }
+
+    private enum MarkdownRenderSegmentType
+    {
+        MARKDOWN,
+        MATH_BLOCK,
+    }
+
+    private enum MathBlockFenceType
+    {
+        NONE,
+        DOLLAR,
+        BRACKET,
+    }
+
+    private sealed record MarkdownRenderPlan(string Source, IReadOnlyList<MarkdownRenderSegment> Segments)
+    {
+        public static readonly MarkdownRenderPlan EMPTY = new(string.Empty, []);
+    }
+
+    private sealed class MarkdownRenderSegment(MarkdownRenderSegmentType type, int start, int length)
+    {
+        private string? cachedContent;
+
+        public MarkdownRenderSegmentType Type { get; } = type;
+
+        public int Start { get; } = start;
+
+        public int Length { get; } = length;
+
+        public int RenderKey { get; } = HashCode.Combine(type, start, length);
+
+        public string GetContent(string source)
+        {
+            if (this.cachedContent is not null)
+                return this.cachedContent;
+
+            this.cachedContent = this.Start == 0 && this.Length == source.Length
+                ? source
+                : source.Substring(this.Start, this.Length);
+
+            return this.cachedContent;
+        }
     }
     
     private async Task RemoveBlock()
@@ -294,4 +609,14 @@ public partial class ContentBlockComponent : MSGComponentBase
         var result = await ReviewAttachmentsDialog.OpenDialogAsync(this.DialogService, this.Content.FileAttachments.ToHashSet());
         this.Content.FileAttachments = result.ToList();
     }
-}
+
+    public async ValueTask DisposeAsync()
+    {
+        if (this.isDisposed)
+            return;
+
+        this.isDisposed = true;
+        await this.DisposeMathContainerIfNeededAsync();
+        this.Dispose();
+    }
+}

app/MindWork AI Studio/Chat/MathJaxBlock.razor

@@ -0,0 +1,5 @@
+@namespace AIStudio.Chat
+
+<div class="@this.RootClass" data-chat-math-block="true" style="white-space: pre-wrap;">
+    @this.MathText
+</div>

app/MindWork AI Studio/Chat/MathJaxBlock.razor.cs

@@ -0,0 +1,18 @@
+using Microsoft.AspNetCore.Components;
+
+namespace AIStudio.Chat;
+
+public partial class MathJaxBlock
+{
+    [Parameter]
+    public string Value { get; init; } = string.Empty;
+
+    [Parameter]
+    public string Class { get; init; } = string.Empty;
+
+    private string RootClass => string.IsNullOrWhiteSpace(this.Class)
+        ? "chat-mathjax-block"
+        : $"chat-mathjax-block {this.Class}";
+
+    private string MathText => $"$${Environment.NewLine}{this.Value}{Environment.NewLine}$$";
+}