spb/polyllm

1"use client";2import * as React from "react";3import ReactMarkdown, { type Components } from "react-markdown";4import remarkGfm from "remark-gfm";5import remarkMath from "remark-math";6import rehypeKatex from "rehype-katex";7import "katex/dist/katex.min.css";8import { CodeBlock } from "./code-block";9import { normalizeMath, splitStreamingMarkdown } from "@/lib/chat/markdown-blocks";10import { cn } from "@/lib/utils";1112/**13 * Chat Markdown renderer.14 *15 * - GFM (tables, task lists, strikethrough) + LaTeX (`$…$`, `$$…$$`, and the `\(…\)` / `\[…\]` forms16 *   models emit, normalised in `lib/chat/markdown-blocks`). Raw HTML is skipped (never injected).17 * - Streaming: the text is split into completed top-level blocks that are rendered once and memoised;18 *   only the trailing block re-parses on each token, so long answers stay cheap and tables/code do not flicker.19 */20const components = (wrap: boolean): Components => ({21  code({ className, children, ...props }) {22    const match = /language-(\w[\w+#-]*)/.exec(className ?? "");23    const text = String(children ?? "").replace(/\n$/, "");24    const isBlock = Boolean(match) || text.includes("\n");25    if (isBlock) return <CodeBlock code={text} lang={match?.[1]} wrap={wrap} />;26    return (27      <code className={className} {...props}>28        {children}29      </code>30    );31  },32  pre({ children }) {33    return <>{children}</>;34  },35  a({ href, children }) {36    return (37      <a href={href} target="_blank" rel="noopener noreferrer nofollow">38        {children}39      </a>40    );41  },42  table({ children }) {43    return (44      <div className="my-2 overflow-x-auto scrollbar-thin">45        <table>{children}</table>46      </div>47    );48  },49});5051const REMARK_PLUGINS = [remarkGfm, remarkMath];52const REHYPE_PLUGINS = [[rehypeKatex, { output: "html", throwOnError: false, strict: false }]] as never[];5354const Block = React.memo(function Block({ content, comps }: { content: string; comps: Components }) {55  return (56    <ReactMarkdown remarkPlugins={REMARK_PLUGINS} rehypePlugins={REHYPE_PLUGINS} components={comps} skipHtml>57      {content}58    </ReactMarkdown>59  );60});6162export const Markdown = React.memo(function Markdown({ content, className, wrap = false, streaming = false }: { content: string; className?: string; wrap?: boolean; streaming?: boolean }) {63  const comps = React.useMemo(() => components(wrap), [wrap]);64  const normalized = React.useMemo(() => normalizeMath(content), [content]);65  const split = React.useMemo(() => (streaming ? splitStreamingMarkdown(normalized, { minTailChars: 24 }) : null), [normalized, streaming]);66  return (67    <div className={cn("prose-chat", wrap && "wrap-code", streaming && "caret", className)}>68      {split ? (69        <>70          {split.blocks.map((b, i) => (71            <Block key={i} content={b} comps={comps} />72          ))}73          {split.tail ? <Block content={split.tail} comps={comps} /> : null}74        </>75      ) : (76        <Block content={normalized} comps={comps} />77      )}78    </div>79  );80});81