Chat widget: markdown styling, KaTeX math rendering, SSE streaming fix

- Add comprehensive CSS for markdown elements in chat bubbles (headers,
  lists, tables, blockquotes, links, code blocks, horizontal rules)
- Integrate KaTeX via marked-katex-extension for inline ($) and display
  ($$) math rendering with DOMPurify whitelist for MathML/SVG elements
- Add normalizeDisplayMath() pre-processor to handle LLM output where $$
  delimiters appear inline rather than on their own lines
- Add dark theme KaTeX styles matching the SpiceBook color palette
- Fix asyncio.to_thread usage in chat SSE endpoint for streaming
This commit is contained in:
Ryan Malloy 2026-02-23 14:25:42 -07:00
parent 99d1ca28d2
commit 6f239c185e
5 changed files with 260 additions and 25 deletions

View File

@ -23,12 +23,12 @@ def _sse_event(event: str, data: dict | list) -> str:
return f"event: {event}\ndata: {payload}\n\n" return f"event: {event}\ndata: {payload}\n\n"
def _build_notebook_context(req: ChatStreamRequest) -> str: async def _build_notebook_context(req: ChatStreamRequest) -> str:
"""Extract SPICE cells and markdown notes from the notebook for LLM context.""" """Extract SPICE cells and markdown notes from the notebook for LLM context."""
if not req.notebook or not req.notebook.notebook_id: if not req.notebook or not req.notebook.notebook_id:
return "" return ""
nb = load_notebook(settings.notebook_dir, req.notebook.notebook_id) nb = await asyncio.to_thread(load_notebook, settings.notebook_dir, req.notebook.notebook_id)
if nb is None: if nb is None:
return "" return ""
@ -64,7 +64,7 @@ async def chat_stream(req: ChatStreamRequest):
async def generate(): async def generate():
# Build context from the notebook if available # Build context from the notebook if available
context = _build_notebook_context(req) context = await _build_notebook_context(req)
question = req.question question = req.question
if req.notebook and req.notebook.title: if req.notebook and req.notebook.title:

View File

@ -28,7 +28,10 @@
"astro-seo-meta": "^5.2.0", "astro-seo-meta": "^5.2.0",
"clsx": "^2.1.0", "clsx": "^2.1.0",
"dompurify": "^3.3.1", "dompurify": "^3.3.1",
"katex": "^0.16.33",
"lucide-react": "^0.468.0", "lucide-react": "^0.468.0",
"marked": "^17.0.3",
"marked-katex-extension": "^5.1.7",
"react": "^19.0.0", "react": "^19.0.0",
"react-dom": "^19.0.0", "react-dom": "^19.0.0",
"satori": "^0.19.2", "satori": "^0.19.2",
@ -6123,6 +6126,31 @@
"integrity": "sha512-H8jvkz1O50L3dMZCsLqiuB2tA7muqbSg1AtGEkN0leAqGjsUzDJir3Zwr02BhqdcITPg3ei3mZ+HjMocAknhhg==", "integrity": "sha512-H8jvkz1O50L3dMZCsLqiuB2tA7muqbSg1AtGEkN0leAqGjsUzDJir3Zwr02BhqdcITPg3ei3mZ+HjMocAknhhg==",
"license": "MIT" "license": "MIT"
}, },
"node_modules/katex": {
"version": "0.16.33",
"resolved": "https://registry.npmjs.org/katex/-/katex-0.16.33.tgz",
"integrity": "sha512-q3N5u+1sY9Bu7T4nlXoiRBXWfwSefNGoKeOwekV+gw0cAXQlz2Ww6BLcmBxVDeXBMUDQv6fK5bcNaJLxob3ZQA==",
"funding": [
"https://opencollective.com/katex",
"https://github.com/sponsors/katex"
],
"license": "MIT",
"dependencies": {
"commander": "^8.3.0"
},
"bin": {
"katex": "cli.js"
}
},
"node_modules/katex/node_modules/commander": {
"version": "8.3.0",
"resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz",
"integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==",
"license": "MIT",
"engines": {
"node": ">= 12"
}
},
"node_modules/kleur": { "node_modules/kleur": {
"version": "3.0.3", "version": "3.0.3",
"resolved": "https://registry.npmjs.org/kleur/-/kleur-3.0.3.tgz", "resolved": "https://registry.npmjs.org/kleur/-/kleur-3.0.3.tgz",
@ -6479,6 +6507,28 @@
"url": "https://github.com/sponsors/wooorm" "url": "https://github.com/sponsors/wooorm"
} }
}, },
"node_modules/marked": {
"version": "17.0.3",
"resolved": "https://registry.npmjs.org/marked/-/marked-17.0.3.tgz",
"integrity": "sha512-jt1v2ObpyOKR8p4XaUJVk3YWRJ5n+i4+rjQopxvV32rSndTJXvIzuUdWWIy/1pFQMkQmvTXawzDNqOH/CUmx6A==",
"license": "MIT",
"bin": {
"marked": "bin/marked.js"
},
"engines": {
"node": ">= 20"
}
},
"node_modules/marked-katex-extension": {
"version": "5.1.7",
"resolved": "https://registry.npmjs.org/marked-katex-extension/-/marked-katex-extension-5.1.7.tgz",
"integrity": "sha512-CVFzrqwpXGVaHByqcVvO/JfzW/OMWrAF3pEfNYNIruzBzM64moANSHapCg1qbzEN+NGf5unHwkMfwJIXHzyDAw==",
"license": "MIT",
"peerDependencies": {
"katex": ">=0.16 <0.17",
"marked": ">=4 <18"
}
},
"node_modules/mdast-util-definitions": { "node_modules/mdast-util-definitions": {
"version": "6.0.0", "version": "6.0.0",
"resolved": "https://registry.npmjs.org/mdast-util-definitions/-/mdast-util-definitions-6.0.0.tgz", "resolved": "https://registry.npmjs.org/mdast-util-definitions/-/mdast-util-definitions-6.0.0.tgz",

View File

@ -29,7 +29,10 @@
"astro-seo-meta": "^5.2.0", "astro-seo-meta": "^5.2.0",
"clsx": "^2.1.0", "clsx": "^2.1.0",
"dompurify": "^3.3.1", "dompurify": "^3.3.1",
"katex": "^0.16.33",
"lucide-react": "^0.468.0", "lucide-react": "^0.468.0",
"marked": "^17.0.3",
"marked-katex-extension": "^5.1.7",
"react": "^19.0.0", "react": "^19.0.0",
"react-dom": "^19.0.0", "react-dom": "^19.0.0",
"satori": "^0.19.2", "satori": "^0.19.2",

View File

@ -12,33 +12,63 @@ import { useChatStore } from '../../lib/chat-store';
import type { ChatMessage } from '../../lib/chat-store'; import type { ChatMessage } from '../../lib/chat-store';
import { streamChat } from '../../lib/chat-api'; import { streamChat } from '../../lib/chat-api';
import { useNotebookStore } from '../../lib/notebook-store'; import { useNotebookStore } from '../../lib/notebook-store';
import { marked } from 'marked';
import markedKatex from 'marked-katex-extension';
import DOMPurify from 'dompurify';
import 'katex/dist/katex.min.css';
import '../../styles/chat-widget.css'; import '../../styles/chat-widget.css';
// ── Lightweight markdown ──────────────────────────────── // ── Markdown rendering (marked + KaTeX + DOMPurify) ──────
function escapeHtml(text: string): string { marked.setOptions({
return text breaks: true,
.replace(/&/g, '&amp;') gfm: true,
.replace(/</g, '&lt;') });
.replace(/>/g, '&gt;')
.replace(/"/g, '&quot;'); marked.use(markedKatex({
throwOnError: false,
nonStandard: true,
}));
// KaTeX generates SVG + spans with specific classes/attributes.
// DOMPurify must allow these through.
const KATEX_TAGS = [
'math', 'semantics', 'mrow', 'mi', 'mo', 'mn', 'msup', 'msub',
'mfrac', 'munderover', 'mover', 'munder', 'msqrt', 'mroot',
'mtable', 'mtr', 'mtd', 'mspace', 'mtext', 'menclose',
'annotation', 'annotation-xml',
];
const KATEX_ATTRS = [
'xmlns', 'encoding', 'mathvariant', 'displaystyle', 'scriptlevel',
'fence', 'stretchy', 'symmetric', 'lspace', 'rspace',
'linethickness', 'columnalign', 'rowalign', 'columnspacing',
'rowspacing', 'columnlines', 'rowlines', 'frame', 'framespacing',
'width', 'height', 'voffset', 'accent', 'accentunder',
'notation', 'minsize', 'maxsize', 'movablelimits',
'aria-hidden', 'focusable', 'role', 'tabindex',
'viewBox', 'preserveAspectRatio', 'd', 'fill', 'stroke',
'stroke-width', 'stroke-linecap', 'stroke-linejoin',
'transform', 'clip-path',
];
// LLMs often emit $$...$$ inline (e.g. "is:$$\frac{1}{s}$$\nNext")
// but marked-katex-extension requires $$ on its own line for display mode.
// This normalizer ensures $$ delimiters get their own lines.
function normalizeDisplayMath(text: string): string {
return text.replace(/\$\$([\s\S]*?)\$\$/g, (_match, inner: string) => {
const trimmed = inner.trim();
return `\n$$\n${trimmed}\n$$\n`;
});
} }
function renderMarkdown(text: string): string { function renderMarkdown(text: string): string {
let html = escapeHtml(text); const normalized = normalizeDisplayMath(text);
// Code blocks (``` ... ```) const raw = marked.parse(normalized, { async: false }) as string;
html = html.replace(/```(\w*)\n([\s\S]*?)```/g, (_m, _lang, code) => { return DOMPurify.sanitize(raw, {
return `<pre><code>${code.trim()}</code></pre>`; ADD_TAGS: KATEX_TAGS,
ADD_ATTR: [...KATEX_ATTRS, 'target', 'class', 'style'],
}); });
html = html.replace(/\*\*(.+?)\*\*/g, '<strong>$1</strong>');
html = html.replace(/(?<!\*)\*([^*]+?)\*(?!\*)/g, '<em>$1</em>');
html = html.replace(/`([^`]+?)`/g, '<code>$1</code>');
html = html.replace(
/\[([^\]]+?)\]\((https?:\/\/[^)]+?)\)/g,
'<a href="$2" target="_blank" rel="noopener noreferrer">$1</a>',
);
html = html.replace(/\n/g, '<br>');
return html;
} }
// ── Components ────────────────────────────────────────── // ── Components ──────────────────────────────────────────
@ -158,6 +188,12 @@ export default function ChatWidget() {
const inputRef = useRef<HTMLInputElement>(null); const inputRef = useRef<HTMLInputElement>(null);
const abortRef = useRef<AbortController | null>(null); const abortRef = useRef<AbortController | null>(null);
// Token batching: accumulate tokens in a ref, flush to state once per animation frame.
// Without this, React 19 batches all rapid set() calls from the SSE loop into one
// render at stream end — so the user sees nothing until the LLM finishes.
const pendingTokensRef = useRef('');
const flushRafRef = useRef(0);
const activeConv = getActiveConversation(); const activeConv = getActiveConversation();
const messages = activeConv?.messages ?? []; const messages = activeConv?.messages ?? [];
@ -225,8 +261,18 @@ export default function ChatWidget() {
setStatusText((evt.data as { text: string }).text); setStatusText((evt.data as { text: string }).text);
break; break;
case 'token': case 'token':
appendToLastAssistant((evt.data as { text: string }).text); // Accumulate in ref; flush to Zustand once per animation frame
setStatusText(''); pendingTokensRef.current += (evt.data as { text: string }).text;
if (!flushRafRef.current) {
flushRafRef.current = requestAnimationFrame(() => {
if (pendingTokensRef.current) {
appendToLastAssistant(pendingTokensRef.current);
pendingTokensRef.current = '';
}
flushRafRef.current = 0;
});
setStatusText('');
}
break; break;
case 'reasoning': case 'reasoning':
setReasoningText((prev) => prev + (evt.data as { text: string }).text); setReasoningText((prev) => prev + (evt.data as { text: string }).text);
@ -250,6 +296,15 @@ export default function ChatWidget() {
); );
} }
} finally { } finally {
// Flush any remaining buffered tokens
if (flushRafRef.current) {
cancelAnimationFrame(flushRafRef.current);
flushRafRef.current = 0;
}
if (pendingTokensRef.current) {
appendToLastAssistant(pendingTokensRef.current);
pendingTokensRef.current = '';
}
setStreaming(false); setStreaming(false);
setStatusText(''); setStatusText('');
abortRef.current = null; abortRef.current = null;

View File

@ -174,6 +174,133 @@
padding: 0; padding: 0;
} }
/* Markdown headings */
.chat-bubble.assistant h1,
.chat-bubble.assistant h2,
.chat-bubble.assistant h3,
.chat-bubble.assistant h4 {
color: var(--color-sb-text-bright);
font-weight: 600;
margin: 0.75rem 0 0.25rem;
line-height: 1.3;
}
.chat-bubble.assistant h1 { font-size: 1rem; }
.chat-bubble.assistant h2 { font-size: 0.9375rem; }
.chat-bubble.assistant h3 { font-size: 0.875rem; }
.chat-bubble.assistant h4 { font-size: 0.8125rem; }
.chat-bubble.assistant h1:first-child,
.chat-bubble.assistant h2:first-child,
.chat-bubble.assistant h3:first-child,
.chat-bubble.assistant h4:first-child {
margin-top: 0;
}
/* Paragraphs */
.chat-bubble.assistant p {
margin: 0.375rem 0;
}
.chat-bubble.assistant p:first-child {
margin-top: 0;
}
.chat-bubble.assistant p:last-child {
margin-bottom: 0;
}
/* Lists */
.chat-bubble.assistant ul,
.chat-bubble.assistant ol {
margin: 0.375rem 0;
padding-left: 1.25rem;
}
.chat-bubble.assistant li {
margin: 0.125rem 0;
}
.chat-bubble.assistant li > ul,
.chat-bubble.assistant li > ol {
margin: 0.125rem 0;
}
/* Horizontal rules */
.chat-bubble.assistant hr {
border: none;
border-top: 1px solid var(--color-sb-border);
margin: 0.625rem 0;
}
/* Blockquotes */
.chat-bubble.assistant blockquote {
border-left: 3px solid var(--color-sb-accent);
margin: 0.5rem 0;
padding: 0.375rem 0.625rem;
background: rgba(59, 130, 246, 0.08);
border-radius: 0 0.25rem 0.25rem 0;
color: var(--color-sb-text);
}
.chat-bubble.assistant blockquote p {
margin: 0.25rem 0;
}
/* Tables — wrapped in a scrollable container via overflow on the bubble */
.chat-bubble.assistant table {
width: 100%;
border-collapse: collapse;
margin: 0.5rem 0;
font-size: 0.75rem;
display: block;
overflow-x: auto;
}
.chat-bubble.assistant th,
.chat-bubble.assistant td {
padding: 0.3rem 0.5rem;
border: 1px solid var(--color-sb-border);
text-align: left;
}
.chat-bubble.assistant th {
background: var(--color-sb-surface);
color: var(--color-sb-text-bright);
font-weight: 600;
}
.chat-bubble.assistant tr:nth-child(even) {
background: rgba(255, 255, 255, 0.02);
}
/* Links */
.chat-bubble.assistant a {
color: var(--color-sb-accent);
text-decoration: none;
}
.chat-bubble.assistant a:hover {
text-decoration: underline;
}
/* KaTeX math — dark theme overrides */
.chat-bubble.assistant .katex {
font-size: 0.9em;
color: var(--color-sb-text-bright);
}
.chat-bubble.assistant .katex-display {
margin: 0.5rem 0;
overflow-x: auto;
overflow-y: hidden;
padding: 0.375rem 0;
}
.chat-bubble.assistant .katex-display > .katex {
white-space: normal;
}
/* Status messages */ /* Status messages */
.chat-status { .chat-status {
text-align: center; text-align: center;