Next AI Draw.io 核心实现深度分析

请关注公众号【碳硅化合物AI】

前言

大家好！上一篇我们聊了项目的整体架构，今天咱们深入代码，看看这 8 个核心模块是怎么实现的。我会从入口类开始，分析关键类的关系，然后用时序图展示流程，最后总结实现的关键点。

一、AI 工具调用系统的实现

1.1 入口类和关键类关系

AI 工具调用系统是整个项目的核心，它的入口在 app/api/chat/route.ts。让我先看看关键类的关系：

入口类 - app/api/chat/route.ts 的 POST 函数：

export async function POST(request: Request) {
    // 1. 验证文件上传
    const validation = validateFileParts(messages)
    
    // 2. 获取 AI 模型配置
    const modelConfig = getAIModel(clientOverrides)
    
    // 3. 获取系统提示词
    const systemPrompt = getSystemPrompt(modelConfig.modelId)
    
    // 4. 创建流式响应
    const result = await streamText({
        model: modelConfig.model,
        system: systemPrompt,
        tools: { display_diagram, edit_diagram, ... }
    })
    
    // 5. 返回流式响应
    return createUIMessageStreamResponse({ stream })
}

这个入口函数做了几件关键的事：验证输入、配置 AI 模型、定义工具、创建流式响应。设计挺清晰的。

1.2 关键流程时序图

让我用时序图展示一下 AI 工具调用的完整流程：

这个流程展示了从用户输入到图表更新的完整过程。关键点是工具调用是异步的，通过 addToolOutput 返回结果给 AI。

1.3 实现关键点

工具定义策略：
项目定义了 4 个工具，每个工具都有明确的职责：

• display_diagram - 创建新图表或重大结构更改
• edit_diagram - 增量编辑，支持 update/add/delete 操作
• append_diagram - 处理截断情况，继续生成
• get_shape_library - 获取图标库文档

这种设计让 AI 可以根据场景选择最合适的工具，既保证了灵活性，又避免了不必要的全量生成。

工具调用处理（hooks/use-diagram-tool-handlers.ts）：

export function useDiagramToolHandlers({...}) {
    const handleToolCall = async ({ toolCall }, addToolOutput) => {
        if (toolCall.toolName === "display_diagram") {
            const { xml } = toolCall.input as { xml: string }
            const isTruncated = !isMxCellXmlComplete(xml)
            
            if (isTruncated) {
                // 处理截断情况
                addToolOutput({ 
                    state: "output-error",
                    errorText: "XML was truncated..." 
                })
            } else {
                // 正常处理
                const result = onDisplayChart(xml)
                addToolOutput({ output: "Diagram displayed successfully" })
            }
        }
    }
}

这里有个巧妙的设计：通过 isMxCellXmlComplete 判断 XML 是否完整，如果不完整就返回错误，让 AI 使用 append_diagram 继续生成。

二、Draw.io XML 处理机制

2.1 入口类和关键类关系

XML 处理的核心在 lib/utils.ts，这个文件有 1681 行，包含了大量的 XML 处理逻辑：

@startuml
class XMLUtils {
    + validateAndFixXml(xml: string): string
    + extractDiagramXML(xml: string): string
    + isMxCellXmlComplete(xml: string): boolean
    + extractCompleteMxCells(xml: string): string
    + wrapWithMxFile(xml: string): string
}

class DiagramContext {
    + loadDiagram(xml: string): string | null
}

class ToolHandlers {
    + handleDisplayDiagram()
    + handleEditDiagram()
}

XMLUtils <-- DiagramContext
XMLUtils <-- ToolHandlers
@enduml

关键函数 - validateAndFixXml：

export function validateAndFixXml(xml: string): string {
    // 1. 检查 XML 大小（限制 1MB）
    if (xml.length > MAX_XML_SIZE) {
        throw new Error("XML too large")
    }
    
    // 2. 解析 XML
    const parser = new DOMParser()
    const doc = parser.parseFromString(xml, "text/xml")
    
    // 3. 验证结构
    // 4. 修复常见错误（缺失属性、无效引用等）
    // 5. 返回修复后的 XML
}

这个函数是 XML 处理的入口，它负责验证和修复 XML，确保生成的 XML 能被 draw.io 正确解析。

2.2 关键流程时序图

XML 处理的流程比较复杂，涉及到验证、修复、提取等多个步骤：

2.3 实现关键点

流式 XML 处理：
项目支持流式响应，但 AI 生成 XML 时可能被截断。项目通过 isMxCellXmlComplete 来判断：

export function isMxCellXmlComplete(xml: string): boolean {
    // 找到最后一个完整的 mxCell 结束位置
    const lastSelfClose = xml.lastIndexOf("/>")
    const lastMxCellClose = xml.lastIndexOf("</mxCell>")
    const lastValidEnd = Math.max(lastSelfClose, lastMxCellClose)
    
    // 检查后缀是否只是闭合标签
    const suffix = xml.slice(lastValidEnd + endOffset)
    return /^(\s*<\/[^>]+>)*\s*$/.test(suffix)
}

这个函数通过正则表达式判断 XML 是否完整，如果后缀只是闭合标签，就认为 XML 是完整的。

XML 修复策略：
validateAndFixXml 函数会修复多种常见错误：

• 缺失必需属性（id, parent 等）
• 无效的引用（source/target 指向不存在的 cell）
• 结构错误（嵌套错误、标签不匹配等）

这种自动修复机制提高了系统的健壮性。

三、多提供商 AI 集成模式

3.1 入口类和关键类关系

多提供商集成的核心在 lib/ai-providers.ts，它抽象了不同提供商的差异：

入口函数 - getAIModel：

export function getAIModel(overrides?: ClientOverrides): ModelConfig {
    const provider = overrides?.provider || process.env.AI_PROVIDER || "bedrock"
    const modelId = overrides?.modelId || process.env.AI_MODEL || "..."
    
    switch (provider) {
        case "openai":
            return {
                model: createOpenAI({ apiKey, baseURL }),
                modelId,
                headers: {}
            }
        case "anthropic":
            return {
                model: createAnthropic({ apiKey, baseURL }),
                modelId,
                headers: ANTHROPIC_BETA_HEADERS
            }
        // ... 其他提供商
    }
}

这个函数通过 switch 语句根据提供商类型创建对应的模型实例，统一了接口。

3.2 关键流程时序图

多提供商集成的流程：

3.3 实现关键点

配置优先级：

1. 客户端提供的配置（浏览器 localStorage）
2. 环境变量配置
3. 默认配置

这种设计让用户可以在浏览器中配置自己的 API 密钥，既保护了隐私，又提供了灵活性。

特殊功能支持：
不同提供商有不同的特殊功能，项目通过 providerOptions 和 headers 来支持：

• OpenAI 的推理模型（o1/o3）需要 reasoningSummary
• Anthropic 的思考预算需要 thinkingBudgetTokens
• Bedrock 的 beta 功能需要 anthropicBeta

这种设计既保持了统一接口，又支持了各提供商的特殊功能。

四、状态管理和数据流设计

4.1 入口类和关键类关系

状态管理的核心是 DiagramContext，它使用 React Context API：

入口组件 - DiagramProvider：

export function DiagramProvider({ children }) {
    const [chartXML, setChartXML] = useState<string>("")
    const [diagramHistory, setDiagramHistory] = useState<[...]>([])
    const [isDrawioReady, setIsDrawioReady] = useState(false)
    
    // 加载图表
    const loadDiagram = (xml: string, skipValidation?: boolean) => {
        const fixedXml = skipValidation ? xml : validateAndFixXml(xml)
        setChartXML(fixedXml)
        // 通过 ref 更新 Draw.io
    }
    
    // 导出图表
    const handleExport = () => {
        // 保存到历史记录
        // 获取 SVG 预览
    }
    
    return (
        <DiagramContext.Provider value={{...}}>
            {children}
        </DiagramContext.Provider>
    )
}

4.2 关键流程时序图

数据流的完整流程：

4.3 实现关键点

状态管理策略：
项目使用 React Context + Hooks，没有使用 Redux。这种选择是合理的，因为：

1. 状态管理需求相对简单
2. 避免了 Redux 的复杂性
3. 代码更直观，易于理解

数据持久化：
图表状态通过 localStorage 持久化：

• next-ai-draw-io-diagram-xml - 当前图表 XML
• next-ai-draw-io-xml-snapshots - 历史记录
• next-ai-draw-io-messages - 聊天消息

这种设计让用户刷新页面后可以恢复之前的状态。

五、流式响应处理技巧

5.1 入口类和关键类关系

流式响应的核心在 app/api/chat/route.ts 和 components/chat-panel.tsx：

流式响应创建：

const result = await streamText({
    model: modelConfig.model,
    system: systemPrompt,
    tools: { display_diagram, edit_diagram, ... }
})

result.onToolCall('display_diagram', async ({ xml }) => {
    // 处理工具调用
})

return createUIMessageStreamResponse({
    stream: createUIMessageStream(result)
})

5.2 关键流程时序图

流式响应的处理流程：

5.3 实现关键点

渐进式渲染：
流式响应让用户可以实时看到 AI 的回复，提升了用户体验。项目通过 useChat hook 来处理流式响应：

const { messages, append, isLoading } = useChat({
    api: '/api/chat',
    onToolCall: handleToolCall,
    // ...
})

工具调用的流式处理：
工具调用在流式响应中实时返回，通过 onToolCall 回调处理。这种设计让工具调用和文本生成可以交错进行，提高了响应速度。

六、文件处理和图表历史系统

6.1 文件处理

文件处理的核心在 lib/use-file-processor.tsx 和 lib/pdf-utils.ts：

PDF 处理：

export async function extractPdfText(file: File): Promise<string> {
    const buffer = await file.arrayBuffer()
    const pdf = await getDocumentProxy(new Uint8Array(buffer))
    const { text } = await extractText(pdf, { mergePages: true })
    return text
}

文件上传限制：

• 最大文件大小：2MB
• 最大文件数量：5 个
• 支持类型：PDF、图像、文本文件

6.2 图表历史系统

历史系统的核心在 contexts/diagram-context.tsx：

const [diagramHistory, setDiagramHistory] = useState<Array<{
    svg: string,  // SVG 预览
    xml: string   // 完整 XML
}>>([])

const handleExport = () => {
    // 获取当前 SVG
    // 保存到历史记录
    setDiagramHistory([...diagramHistory, { svg, xml: chartXML }])
}

自动快照机制：
每次 AI 编辑前自动保存快照，用户可以随时恢复到任意版本。这种设计让用户可以放心地让 AI 修改图表。

七、Electron 桌面应用实现

7.1 入口类和关键类关系

Electron 应用的主进程入口在 electron/main/index.ts：

入口代码：

app.whenReady().then(async () => {
    // 注册 IPC 处理器
    registerIpcHandlers()
    
    // 启动 Next.js 服务器（生产环境）
    if (!isDev) {
        serverUrl = await startNextServer()
    }
    
    // 创建主窗口
    createWindow(serverUrl)
})

7.2 实现关键点

Next.js 集成：
Electron 应用通过启动 Next.js standalone 服务器来运行 Web 应用。这种设计让 Web 和桌面版本共享同一套代码。

原生功能：

• 使用 OS keychain 安全存储 API 密钥
• 原生文件对话框打开/保存 .drawio 文件
• 通过菜单管理配置预设

这些原生功能提升了桌面应用的用户体验。

八、国际化和 MCP Server 集成

8.1 国际化实现

国际化的核心在 lib/i18n/ 目录：

export const i18n = {
    defaultLocale: "en",
    locales: ["en", "zh", "ja"],
} as const

路由结构：

• /[lang]/ - 主页
• /[lang]/about - 关于页

使用 Next.js 动态路由实现多语言支持。

8.2 MCP Server 集成

MCP Server 在 packages/mcp-server/src/index.ts：

const server = new McpServer({
    name: "next-ai-drawio",
    version: "0.1.2",
})

// 注册工具
server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [
        {
            name: "display_diagram",
            description: "Display a diagram...",
            inputSchema: { ... }
        }
    ]
}))

核心功能：

• 创建图表
• 编辑图表
• 获取图表状态
• 实时浏览器预览

MCP Server 让 AI 代理（Claude Desktop、Cursor 等）可以直接操作图表，扩展了应用的使用场景。

总结说明

这 8 个核心模块构成了 Next AI Draw.io 的完整功能。每个模块都有清晰的职责和设计：

1. AI 工具调用系统 - 通过结构化工具让 AI 操作图表，设计巧妙
2. XML 处理机制 - 完善的验证和修复逻辑，支持流式处理
3. 多提供商集成 - 统一的接口设计，支持 11 种提供商
4. 状态管理 - 简单有效的 Context + Hooks 方案
5. 流式响应 - 实时更新，提升用户体验
6. 文件处理 - 支持 PDF、图像、文本文件，功能完整
7. Electron 应用 - 原生功能集成，用户体验好
8. 国际化与 MCP - 扩展了应用的使用场景