如何在AI流式数据中渲染mermaid图表

如何在AI流式数据中渲染Mermaid图表

接上一篇文章《如何优雅的在AI应用中渲染Markdown数据》，有同学提出希望可以支持mermaid图表渲染，最近我周末把这个功能加上来，接下来，我将结合 ds-markdown 介绍在AI对话场景中，如何优雅地处理流式输出的Mermaid图表代码，并实现实时渲染？本文将深入探讨一个完整的技术解决方案。

首先我们来看效果

在线DEMO

主要功能点

• 图表渲染
• 图表交互（放大、缩小、自适应大小、全屏、下载等）

技术架构设计

我们采用插件化架构，把mermaid作为ds-markdown的一个内置插件，因为需要引入多个package，如果放到ds-markdown里，会导致不需要该功能的也安装这些包，这样可以保持ds-markdown的轻量化

插件地址

流式数据处理流程

graph TD A[AI输出流式文本] --> B[React-Markdown解析] B --> C[检测Code代码块] C --> CODE{是否为mermaid} CODE --> |是| M[检测mermaid渲染] CODE --> |否| N[普通Code渲染] M -->|完整| E[立即渲染图表] M -->|不完整| F{尝试渲染Mermaid} F --> |可渲染| H[渲染不完整的图表] F --> |报错| K[继续显示原来的图表] H --> G[继续接收流式数据] K --> G G --> A E --> I[用户交互层]

核心技术实现

创建一个 mermaid 插件

import { createBuildInPlugin, mermaidId } from 'ds-markdown/plugins';

const plugin = createBuildInPlugin({
  id: mermaidId,
  // 用于转换mermaid代码
  rehypePlugin: [rehypeMermaid],
  // 用于检测代码的完整性
  remarkPlugin: [remarkMermaid],
  components: {
    // Mermaid渲染抓紧
    MermaidBlock: MermaidBlock as any,
  },
});

判断是否为mermaid，并渲染

这是整个方案的核心技术，通过rehype插件实现markdown到React组件的智能转换， remarkMermaid.ts部分代码

/**
 * Rehype插件：将markdown中的mermaid代码块转换为React组件
 */
export const rehypeMermaid: Plugin<[RehypeMermaidOptions?]> = (options = {}) => {
  return (tree) => {
    visit(tree, 'element', (node, index: number, parent: any) => {
      // 1. 检测是否为代码块结构
      if (node.tagName === 'pre' && node.children && node.children.length > 0 && node.children[0].tagName === 'code') {
        const codeNode = node.children[0];
        const className = codeNode.properties?.className || [];

        // 2. 检测是否包含mermaid语言标识，只需判断 className中是否包含 language-mermaid
        if (className.some((cls: string) => cls.includes('language-mermaid'))) {
          const rawCode = codeNode.children?.[0]?.value || '';

          // 3. 获取完整性信息（来自remark插件），可查看 完整性检测机制
          const isComplete = (codeNode as any)?.properties?.['data-is-code-block-complete'] ?? false;

          // 4. 清理代码内容
          const code = cleanMermaidCode(rawCode);

          // 5. 转换为自定义组件
          node.type = 'element';
          node.tagName = 'MermaidBlock';
          node.properties = {
            code,
            isComplete, // 传递完整性标志
            mermaidConfig: options.mermaidConfig,
          };
          node.children = [];
        }
      }
    });
  };
};

完整性检测机制

因为我们在跟图表交互的时候，比如复制、下载，会下载到不完整的代码，这就导致无用信息，我们需要在流式完整输出（mermaid代码已经输出完成）

这个是不完整的mermaid代码，因为没有 ``` 结尾

这个是不完整的mermaid代码
\`\`\` mermaid
graph TD
    A[AI输出流式文本] --> B[Markdown解析器]
    B --> C[检测Mermaid代码块]
    C --> D{代码块完整性检测}

这是完整的mermaid代码，因为有 ``` 结尾

这个是不完整的mermaid代码
\`\`\` mermaid
graph TD
    A[AI输出流式文本] --> B[Markdown解析器]
    B --> C[检测Mermaid代码块]
    C --> D{代码块完整性检测}
    D -->|完整| E[立即渲染图表]
    D -->|不完整| F[显示代码预览]
    F --> G[继续接收流式数据]
    G --> D
    E --> H[用户交互层]
\`\`\`

项目通过remark插件实现了智能的完整性检测：

/**
 * Remark插件：检测代码块的完整性
 * 在markdown解析阶段检测代码块是否有结束标记
 */
export const remarkMermaidCompleteness = (options: RemarkMermaidCompletenessOptions = {}) => {
  const { enabled = true } = options;

  return (tree: any, file: any) => {
    if (!enabled) return;
    // 从file对象获取原始文本
    const source = file.value || file.contents;

    visit(tree, 'code', (node: any, index: number | undefined, parent: any) => {
      // 检查是否是mermaid代码块
      if (node.lang === 'mermaid') {
        // 检查代码块是否完整
        const isComplete = isCodeBlockComplete(node, parent, source);
        // 在节点上添加完整性信息
        if (!node.data) {
          node.data = {
            hProperties: {},
          };
        }
        if (!node.data.hProperties) {
          node.data.hProperties = {};
        }
        node.data.hProperties['data-is-code-block-complete'] = isComplete;
      }
    });
  };
};



function isCodeBlockComplete(node: any, parent: any, source?: string): boolean {
  // 使用位置信息检测代码块完整性
  if (source && node.position) {
    const { start, end } = node.position;
    const fullCodeBlock = source.slice(start.offset, end.offset);

    // 检查是否以```结尾
    return fullCodeBlock.trim().endsWith('```');
  }

  // 备用检测逻辑
  const code = node.value || '';
  return code.trim().endsWith('```');
}

这里有个坑，我想通过 node.data.isComplete = ture来传递是否完整，但是在rehype中无法获取到改制，后面看源码才发现，内部值传递只能通过 node.data.hProperties, 所有有这个代码

if (!node.data.hProperties) { 
    node.data.hProperties = {};
}
node.data.hProperties['data-is-code-block-complete'] = isComplete;

关键特性：

• 基于AST位置信息的精确检测
• 支持多种markdown格式
• 实时更新完整性状态

高性能渲染服务

采用单例模式的MermaidService确保性能和资源管理：

class MermaidService {
  private static instance: MermaidService;
  private isInitialized = false;
  private initializationPromise: Promise<void> | null = null;

  public async render(id: string, code: string): Promise<{ svg: string }> {
    await this.initialize();
    return mermaid.render(id, code);
  }

  public async parse(code: string) {
    return await mermaid.parse(code);
  }
}

优势：

• 避免重复初始化
• 统一的错误处理
• 支持异步并发渲染
• 内存使用优化

智能渲染策略

RenderGraph组件实现了智能的渲染更新机制：

useEffect(() => {
  if (!code.trim()) {
    setSvgElement(null);
    return;
  }

  setError(null);
  const renderChart = async () => {
    const viewID = getMermaidId();
    try {
      // 先验证语法
      await mermaidService.parse(code);
      // 再进行渲染
      const { svg } = await mermaidService.render(viewID, code);
      const svgElement = parseSvgToJsx(svg);
      setSvgElement(svgElement);
    } catch (err) {
      setError(err.message);
    }
  };

  renderChart();
}, [code]);

特点：

• 语法验证在前，渲染在后
• 错误状态自动恢复
• SVG到JSX的高效转换
• 支持动态ID生成避免冲突

4. 双模式显示策略

针对流式场景，提供了图表和代码的双模式切换：

const MermaidBlock: React.FC<MermaidProps> = (props) => {
  const { code, isComplete = false } = props;
  const [activeSegmented, setActiveSegmented] = useState('mermaid');

  return (
    <GraphProvider isComplete={isComplete}>
      <CodeBlockWrap title={
        <Segmented
          segments={['diagram', 'code']}
          value={activeSegmented}
          onChange={setActiveSegmented}
        />
      }>
        <div style={{ display: activeSegmented === 'mermaid' ? 'block' : 'none' }}>
          <RenderGraph code={code} isComplete={isComplete} />
        </div>
        <div style={{ display: activeSegmented === 'code' ? 'block' : 'none' }}>
          <RenderCode code={code} />
        </div>
      </CodeBlockWrap>
    </GraphProvider>
  );
};

高级功能实现

图表交互增强

通过svg-pan-zoom库实现了丰富的图表交互：

class PanZoomState {
  public updateElement(diagramView: SVGElement, config?: { pan: Point; zoom: number }) {
    this.pzoom = panzoom(diagramView, {
      center: true,
      fit: true,
      panEnabled: true,
      zoomEnabled: true,
      maxZoom: 12,
      minZoom: 0.2,
      onPan: (pan) => this.handlePanChange(pan),
      onZoom: (zoom) => this.handleZoomChange(zoom),
    });
  }
}

功能特性：

• 平滑的缩放和平移
• 自适应视图居中
• 状态持久化
• 响应式resize处理

导出功能

支持图表的PNG导出和剪贴板复制：

export interface RenderGraphRef {
  download: () => Promise<boolean>;
  copy: () => Promise<boolean>;
  getSvg: () => SVGElement | null;
}

const downloadPng = async ({ id, width, height }) => {
  const svgElement = document.querySelector(id);
  // SVG转Canvas转PNG的完整流程
  const canvas = document.createElement('canvas');
  const ctx = canvas.getContext('2d');
  // ... 转换逻辑
};

主题系统集成

与ds-markdown的主题系统无缝集成：

const mermaidConfig = {
  theme: 'default', // 支持 default, dark, forest, neutral 等主题
  flowchart: {
    useMaxWidth: true,
    htmlLabels: true,
  },
};

<ConfigProvider mermaidConfig={mermaidConfig}>
  <Markdown plugins={[plugin]}>{content}</Markdown>
</ConfigProvider>

实际应用场景

AI对话场景

// 模拟AI流式输出
const streamingExample = `
AI正在生成流程图...

\`\`\`mermaid
graph TD
    A[开始分析数据]
    A --> B{数据是否完整？}
    B -->|是| C[执行分析算法]
    B -->|否| D[数据清洗]
    D --> C
    C --> E[生成报告]
\`\`\`
`;

// 使用插件渲染
<Markdown
  plugins={[mermaidPlugin]}
  interval={16}  // 流式输出间隔
  disableTyping={false}  // 启用打字机效果
>
  {streamingExample}
</Markdown>

文档生成场景

适用于自动生成的技术文档，支持：

• API文档的序列图生成
• 系统架构图的实时更新
• 数据流程图的动态展示

错误边界

// 优雅的错误处理
try {
  await mermaidService.parse(code);
  const { svg } = await mermaidService.render(viewID, code);
} catch (err) {
  // 显示错误信息而不是崩溃
  setError(err.message);
}

如果有兴趣，大家可以查看源码 ds-markdown-mermaid-plugin