Unlocking Rich UI Component Rendering in AI Responses

LLMs have enabled us to solve a new class of problems with more flexibility than ever. But language models are inherently text-based, which has led to AI interfaces being heavily text dominated.

As someone who has been creating web technology experiences my entire life, I’m not satisfied with so much UI being replaced with text. At Vetted, we have been building a shopping research assistant, and shopping is inherently visual and interface-heavy. Products need to display images, and the UI needs to present structured data that lets users navigate between products and compare them.

Over the years, we have been experimenting with new ways to incorporate rich UI components into our LLM responses. We’ve done this by supplementing our text payloads with structured product data that renders as product cards and research components, but we weren’t happy with these elements being separated from the main text response. So, I set out to find a way for LLM output to embed UI components directly in our markdown output.

MDX Enhanced Dynamic Markdown

I built react-markdown-with-mdx: an HOC (Higher-Order Component) wrapper around react-markdown that enhances its markdown processing to support JSX component tags. You can register approved React components with it, ensuring only a safe subset of JSX component tags can be rendered. The library includes an optional validation helper function called mdxComponent, which pairs React components with zod validators to validate JSX attributes.

This lets you prompt LLM calls to generate JSX tags that can be rendered safely with a clean and easy integration. Here’s what it looks like in action in our prototype UI at Vetted:

The code looks something like this:

import ReactMarkdown from "react-markdown"
import { withMdx, mdxComponent, type MdxComponents} from "react-markdown-with-mdx"

const MdxReactMarkdown = withMdx(ReactMarkdown)

interface MdxMarkdownRendererProps {
  markdown: string
}

const MdxMarkdownRenderer: React.FC<MdxMarkdownRendererProps> = ({
  markdown,
}) => {
  return (
    <MdxReactMarkdown components={components}>
      {markdown}
    </MdxReactMarkdown>
  )
}

const components: MdxComponents = {
  "card-carousel": mdxComponent(
    MdxCardCarousel,
    z.object({ children: z.any() })
  ),
  "editorial-card": mdxComponent(
    MdxEditorialCard,
    z.object({
      id: z.string(),
      award: z.string().optional(),
      rating: z.string().optional(),
      ranking: z.string().optional(),
      children: z.any()
    })
  ),
  "product-card": mdxComponent(
    MdxProductCard,
    z.object({ name: z.string() })
  )
}

Unlike projects like MCP-UI, these components aren’t loaded in externally via an iframe that needs window message passing for integration, and they aren’t relegated to a separate message outside the main text. Instead, they’re processed into framework-native React components that are embedded directly in the main LLM-generated text. This essentially enables HTML component-like behavior in React and other JSX frameworks, letting you extend markdown with any UI component you desire!

To power the response streaming shown in the video, I needed to balance the HTML tag tree and truncate incomplete tags. This ensures the MDX parser receives valid HTML and blocks partial tag tokens until they’re complete. To enable this, I also created html-balancer-stream. It allows you to auto-close and balance unclosed tags or strip them out until they’re ready. It provides both streaming and non-streaming APIs.

How it Works: Powering up AI Markdown responses with MDX

I originally experimented with this by creating a prototype using HTML components. This approach worked well because HTML components are registered directly with the browser DOM. This allowed custom HTML components to be injected as HTML tag strings directly via innerHTML with no additional render hooks. This was fine for a prototype, but it was definitely not ready for production. Doing that would introduce major security risks and create side effects within our rendering framework. I wanted to find a way to do this safely, and tightly integrate it with React, which we use at Vetted.

For Markdown rendering, we already used the library react-markdown to convert LLM-generated markdown text into React components. Markdown supports adding raw HTML but react-markdown rightly does not support this. It converts HTML tags into sanitized strings because allowing arbitrary HTML in dynamic markdown content would be incredibly dangerous. My goal was to render the tags as framework-native React components to ensure that the components could only run trusted encapsulated code.

To determine the best way to do that, I started looking under the hood. react-markdown is powered by the unified project; a framework and community ecosystem for parsing and transforming syntax trees that can combine and convert various languages. react-markdown is actually a thin wrapper around unified that parses the markdown AST, converts it to an HTML AST, and then into React components.

It does this by parsing markdown with remark: a unified markdown processor that converts markdown into an mdast markdown AST. It then converts that to a hast HTML AST from rehype using remark-rehype. Finally, it uses hast-util-to-jsx-runtime to safely convert the hast AST into JSX components that can be rendered by React and other JSX-based rendering frameworks.

After understanding how react-markdown works, I felt the best approach for adding React component support to markdown was extending remark to enable JSX parsing. Fortunately, the react-markdown project mentions this in its readme and refers developers to the MDX project, which does exactly that. It lets you put JSX syntax directly in Markdown. Problem solved!

Of course, it’s never that simple. MDX is a very impressive project. It’s so impressive that it doesn’t just let you put static JSX components in your Markdown, it supports the entire JSX syntax, including imports, exports, and dynamic Javascript expressions.

Unfortunately, that takes us back to square one. Using it with dynamic LLM text output would be incredibly insecure and open up major security vulnerabilities. The MDX author is fully aware of this. MDX is designed for compile-time rendering of trusted static content, not runtime rendering.

Thankfully, MDX is built on the unified project which has a modular parsing and transformation pipeline. While I could not use MDX out of the box, I could leverage its remark-mdx parser to build an MDX AST in a custom pipeline. This let me add the exact subset of functionality I needed for static JSX components in markdown. Now, our LLM output could include JSX tags with safe conversion to framework-native runtime components!

To complete the pipeline, I needed to create a few new libraries. The MDX parser detects JSX tags in the markdown AST and adds them to the syntax tree as MDX nodes. With pipeline control, I could strip the MDX nodes with dynamic syntax while preserving the static ones.

I created rehype-mdx-elements to do this. It filters for static JSX tags and converts them into hast HTML AST elements so they can be processed by the hast-util-to-jsx-runtime library. It converts the HTML AST nodes into runtime components that must be registered with the utility, ensuring only trusted code gets executed.

I also needed to make remark-unravel-mdx, which removes unnecessary paragraph wrappers around MDX nodes, producing a cleaner component tree.

These libraries power the react-markdown-with-mdx processing pipeline. They enable the safe MDX-to-React component transformation shown in the demo.

Generating LLM JSX Output

To use this with AI-generated output, you can prompt an LLM with your JSX component tag and attribute schemas, their acceptable children, and examples of their usage. LLMs excel at generating static JSX tags because they share the same syntax as XML, a language they’re heavily trained on.

As JSX tags with many attributes can confuse LLMs, it’s best to keep attributes minimal. When you ask an LLM to generate a tag with coupled attributes, they can fall out of sync and lead to hallucinations. We avoid this by having the LLM fill in simple ID strings that components can look up in shared state to retrieve detailed data. The only data we ask the LLM to generate directly is context-sensitive content like text snippets from research sources.

Since the outputs are JSX tags, they can conform 1:1 with your existing components, so existing documentation and schemas can be shared between your prompts and codebase. This lets you keep their definitions and versions in sync.

I’ve long felt the need for stronger connections between LLM outputs and rendering frameworks. I’m hopeful react-markdown-with-mdx will help fill that gap!