animated markdown for LLM chat messages

Markdown.tsx

'use client';

import React from 'react';
import ReactMarkdown from 'react-markdown';

import rehypeKatex from 'rehype-katex';
import rehypeRaw from 'rehype-raw';
import remarkGfm from 'remark-gfm';
import remarkMath from 'remark-math';

import { cn } from '@/lib/utils';
import 'katex/dist/katex.min.css';

const normalizeLaTeX = (text: string) =>
	text.replace(/\\\[/g, '$$$').replace(/\\\]/g, '$$$').replace(/\\\(/g, '$$$').replace(/\\\)/g, '$$$');

import { Link } from 'react-router-dom';
import { DocLink } from '../InternalDocLink';

const components = {
	table(props) {
		return (
			<div className="table-overflow" style={{ overflowX: 'auto' }}>
				<table>{props.children}</table>
			</div>
		);
	},
	a(props) {
		const href = props.href || '';
		if (href.startsWith('/doc/')) {
			return <DocLink id={href.split('/').at(2)!} />;
		}
		return (
			<Link to={href} target="_blank" rel="noopener noreferrer">
				{props.children}
			</Link>
		);
	},
	code: ({ node, className, children, ...props }: any) => {
		const match = /language-(\w+)/.exec(className || '');
		return match ? (
			<div {...props}>
				<SyntaxHighlighter style={style.docco} language={match[1]}>
					{children}
				</SyntaxHighlighter>
			</div>
		) : (
			<code {...props}>{children}</code>
		);
	},
};

const remarkMathOptions = { singleDollarTextMath: false };
const remarkPlugins = [[remarkMath, remarkMathOptions], remarkGfm] as const;
const rehypePlugins = [rehypeRaw, rehypeKatex] as const;

export const Markdown = React.memo(function Markdown({
	content,
	className = '',
}: {
	content: string;
	className?: string;
}) {
	return (
		<ReactMarkdown
			className={cn(
				'markdown prose prose-neutral w-full break-words dark:prose-invert prose-table:font-mono min-h-6',
				className
			)}
			remarkPlugins={remarkPlugins}
			rehypePlugins={rehypePlugins}
			children={normalizeLaTeX(content)}
			components={components}
		/>
	);
});

export default Markdown;

import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import style from 'react-syntax-highlighter/dist/esm/styles/hljs/atom-one-light';

interface SmoothTextProps {
	content: string;
	className?: string;
	sep?: string;
	animation?: string | null;
	animationDuration?: string;
	animationTimingFunction?: string;
	codeStyle?: any;
}

interface AnimatedImageProps {
	src: string;
	alt: string;
}

interface CustomRendererProps {
	rows: any[];
	stylesheet: any;
	useInlineStyles: boolean;
}

const AnimatedImage: React.FC<AnimatedImageProps> = ({ src, alt }) => {
	const [isLoaded, setIsLoaded] = React.useState(false);

	const imageStyle = isLoaded
		? {
				animation: 'var(--md-animation)',
				whiteSpace: 'pre-wrap',
			}
		: {
				display: 'none',
			};

	return <img src={src} alt={alt} onLoad={() => setIsLoaded(true)} style={imageStyle} />;
};

const TokenizedText = ({ input, sep }: any) => {
	const tokens = React.useMemo(() => {
		if (typeof input !== 'string') return null;

		let splitRegex;
		if (sep === 'word') {
			splitRegex = /(\s+)/;
		} else if (sep === 'char') {
			splitRegex = /(.)/;
		} else {
			throw new Error('Invalid separator');
		}

		return input.split(splitRegex).filter(token => token.length > 0);
	}, [input, sep]);

	return (
		<>
			{tokens?.map((token, index) => (
				<span
					className="md-animate"
					key={index}
				>
					{token}
				</span>
			))}
		</>
	);
};

function toMilliseconds(duration: string) {
	if (typeof duration === 'string') {
		if (duration.endsWith('ms')) {
			return parseFloat(duration);
		}
		if (duration.endsWith('s')) {
			return parseFloat(duration) * 1000;
		}
	}
	return typeof duration === 'number' ? duration : 0;
}

export const MarkdownAnimated: React.FC<SmoothTextProps> = ({
	animation: animationInit = 'fadeIn',
	animationDuration = '1s',
	...props
}) => {
	const [animation, setAnimation] = React.useState(animationInit);

	React.useEffect(() => {
		const id = setTimeout(() => setAnimation(animation), toMilliseconds(animationDuration));
		return () => clearTimeout(id);
		// eslint-disable-next-line react-hooks/exhaustive-deps
	}, [animationInit]);

	return <MarkdownAnimateText {...props} animation={animation} animationDuration={animationDuration} />;
};

/**
 * Inlined from:
 * https://github.com/data-maki/flowtoken/blob/main/src/components/AnimatedMarkdown.tsx
 */
export const MarkdownAnimateText: React.FC<SmoothTextProps> = ({
	content,
	className = '',
	sep = 'word',
	animation = 'fadeIn',
	animationDuration = '1s',
	animationTimingFunction = 'ease-in-out',
	codeStyle = null,
}) => {
	codeStyle = codeStyle || style.docco;

	const mdAnimation = React.useMemo(
		() => ({
			'--md-animation': `${animation} ${animationDuration} ${animationTimingFunction} 1`,
		}),
		[animation, animationDuration, animationTimingFunction]
	);

	// Memoize animateText function to prevent recalculations if props do not change
	const animateText: (text: string | Array<any>) => React.ReactNode = React.useCallback(
		(text: string | Array<any>) => {
			const processText: (input: any) => React.ReactNode = (input: any) => {
				if (Array.isArray(input)) {
					// Process each element in the array
					return input.map(element => processText(element));
				} else if (typeof input === 'string') {
					return <TokenizedText input={input} sep={sep} />;
				} else if (React.isValidElement(input)) {
					// If the element is a React component or element, clone it and process its children
					return input;
				} else {
					// Return non-string, non-element inputs unchanged (null, undefined, etc.)
					return input;
				}
			};
			if (!animation) {
				return text;
			}
			return processText(text);
		},
		[animation, sep]
	);

	// Memoize components object to avoid redefining components unnecessarily
	const components: any = React.useMemo(
		() => ({
			text: ({ node, ...props }: any) => animateText(props.children),
			h1: ({ node, ...props }: any) => <h1 {...props}>{animateText(props.children)}</h1>,
			h2: ({ node, ...props }: any) => <h2 {...props}>{animateText(props.children)}</h2>,
			h3: ({ node, ...props }: any) => <h3 {...props}>{animateText(props.children)}</h3>,
			h4: ({ node, ...props }: any) => <h4 {...props}>{animateText(props.children)}</h4>,
			h5: ({ node, ...props }: any) => <h5 {...props}>{animateText(props.children)}</h5>,
			h6: ({ node, ...props }: any) => <h6 {...props}>{animateText(props.children)}</h6>,
			p: ({ node, ...props }: any) => <p {...props}>{animateText(props.children)}</p>,
			li: ({ node, ...props }: any) => (
				<li {...props} style={mdAnimation}>
					{animateText(props.children)}
				</li>
			),
			a: ({ node, ...props }: any) => (
				<a {...props} href={props.href} target="_blank" rel="noopener noreferrer">
					{animateText(props.children)}
				</a>
			),
			strong: ({ node, ...props }: any) => <strong {...props}>{animateText(props.children)}</strong>,
			em: ({ node, ...props }: any) => <em {...props}>{animateText(props.children)}</em>,
			code: ({ node, className, children, ...props }: any) => {
				const match = /language-(\w+)/.exec(className || '');
				return match ? (
					<div {...props} className="md-animate">
						<SyntaxHighlighter style={codeStyle} language={match[1]} renderer={customRenderer}>
							{children}
						</SyntaxHighlighter>
					</div>
				) : (
					<code {...props}>{animateText(children)}</code>
				);
			},

			hr: ({ node, ...props }: any) => (
				<hr
					className="md-animate"
					{...props}
					style={{
						whiteSpace: 'pre-wrap',
					}}
				/>
			),
			img: ({ node, ...props }: any) => <AnimatedImage src={props.src} alt={props.alt} />,
			table: ({ node, ...props }: any) => (
				<div {...props} className="table-overflow" style={{ overflowX: 'auto' }}>
					<table>{props.children}</table>
				</div>
			),
			tr: ({ node, ...props }: any) => <tr {...props}>{animateText(props.children)}</tr>,
			td: ({ node, ...props }: any) => <td {...props}>{animateText(props.children)}</td>,
		}),
		[animateText]
	);

	return (
		<div style={mdAnimation} className="contents [&_.md-animate]:animate-[var(--md-animation)] [&_span.md-animate]:inline-block [&_span.md-animate]:whitespace-pre-wrap">
			<ReactMarkdown
				className={cn(
					'markdown prose prose-neutral w-full break-words dark:prose-invert prose-table:font-mono min-h-6',
					className
				)}
				components={components}
				remarkPlugins={remarkPlugins}
				rehypePlugins={rehypePlugins}
				children={normalizeLaTeX(content)}
			/>
		</div>
	);
};

const customRenderer: React.FC<CustomRendererProps> = ({ rows, stylesheet, useInlineStyles }) => {
	return rows.map((node, i) => (
		<div key={i} style={node.properties?.style || {}}>
			{node.children.map((token: any, key: string) => {
				// Extract and apply styles from the stylesheet if available and inline styles are used
				const tokenStyles =
					useInlineStyles && stylesheet
						? { ...stylesheet[token?.properties?.className[1]], ...token.properties?.style }
						: token.properties?.style || {};
				return (
					<span key={key} style={tokenStyles}>
						{token.children &&
							token.children[0].value.split(' ').map((word: string, index: number) => (
								<span
									className="md-animate"
									key={index}
								>
									{word + (index < token.children[0].value.split(' ').length - 1 ? ' ' : '')}
								</span>
							))}
					</span>
				);
			})}
		</div>
	));
};