告别状态混乱：Pydantic-AI图可视化让Jupyter工作流一目了然
项目地址: https://gitcode.com/GitHub_Trending/py/pydantic-ai 你是否曾在Jupyter Notebook中调试复杂AI工作流时迷失在嵌套函数调用的迷宫中？是否因状态管理混乱导致实验结果难以复现？Pydantic-AI的图可视化方案正是为解决这些痛点而生。本文将带你掌握如何在Jupyter环境中利用pydantic-graph实现工作流的可视化与状态管理，读完你将获得：
- 用类型提示定义有向图的核心方法
- Jupyter中实时渲染流程图的实操技巧
- 状态持久化确保实验可复现的最佳实践
- 结合AI Agent实现自动反馈循环的完整案例
为什么选择图可视化方案？
传统的控制流（if-else/for循环）在处理多步骤AI工作流时往往导致"意大利面条代码"，而Pydantic-AI的Graph模块提供了更优雅的解决方案。通过将工作流抽象为节点和边的集合，不仅使代码结构清晰可维护，更能直观展示状态流转过程。
图1：基于pydantic-graph的AI工作流执行示意图（源码参考）
Pydantic-Graph的核心优势在于：
- 类型安全：通过Python类型提示定义节点关系，静态检查避免运行时错误
- 状态管理：内置状态持久化机制，自动记录每个节点执行后的状态快照
- 可视化支持：原生集成Mermaid图表，在Jupyter中一键渲染流程图
- 无AI依赖：纯Python图引擎，可独立于LLM使用（但完美支持AI Agent场景）
核心组件快速上手
安装与环境配置
pydantic-graph已作为Pydantic-AI的核心依赖包含在内，通过以下命令确保安装最新版本：
pip install -U pydantic-ai
# 或使用uv包管理器（推荐）
uv add pydantic-ai
在Jupyter Notebook中验证安装：
from pydantic_graph import Graph, BaseNode
print(f"pydantic-graph版本: {BaseNode.__version__}") # 输出当前版本号
定义你的第一个状态机
以下示例实现一个简单的计数器状态机，当数值达到5的倍数时终止：
from dataclasses import dataclass
from pydantic_graph import BaseNode, End, Graph, GraphRunContext
@dataclass
class DivisibleBy5(BaseNode[None, None, int]):
foo: int
async def run(self, ctx: GraphRunContext) -> Increment | End[int]:
if self.foo % 5 == 0:
return End(self.foo) # 终止节点
else:
return Increment(self.foo) # 转移到Increment节点
@dataclass
class Increment(BaseNode):
foo: int
async def run(self, ctx: GraphRunContext) -> DivisibleBy5:
return DivisibleBy5(self.foo + 1) # 数值加1后回到判断节点
# 创建图实例并运行
fives_graph = Graph(nodes=[DivisibleBy5, Increment])
result = fives_graph.run_sync(DivisibleBy5(4))
print(f"最终结果: {result.output}") # 输出: 5
代码解析：这个两节点图展示了核心概念，
DivisibleBy5和Increment节点通过返回类型声明了它们之间的边关系，Graph类自动解析这些关系构建执行路径。
Jupyter中的可视化实践
Mermaid图表实时渲染
Pydantic-Graph内置Mermaid图表生成功能，在Jupyter中执行以下代码即可渲染流程图：
from IPython.display import Image, display
from graph_example import DivisibleBy5, fives_graph
# 生成并显示流程图
display(Image(fives_graph.mermaid_image(start_node=DivisibleBy5)))
图2：自动生成的计数器状态机流程图，蓝色节点表示起始点，红色箭头表示终止路径
状态持久化与可视化调试
当工作流变得复杂时，启用完整状态持久化可追踪每一步的状态变化：
from pydantic_graph import FullStatePersistence
# 创建带状态记录的图执行上下文
persistence = FullStatePersistence()
async with fives_graph.iter(DivisibleBy5(4), persistence=persistence) as run:
async for node in run:
print(f"执行节点: {node}，当前状态: {run.state}")
# 打印完整执行历史
for i, step in enumerate(persistence.history):
print(f"步骤{i}: {step.node} → 状态: {step.state}")
执行结果将展示从初始值4到终止值5的完整路径，配合logfire可实现更详细的可视化监控。
AI工作流实战案例
邮件生成与反馈循环
以下案例实现一个AI邮件写作助手，包含"写作-反馈"双节点循环：
from dataclasses import dataclass, field
from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_graph import BaseNode, End, Graph, GraphRunContext
# 定义状态模型
@dataclass
class State:
user: dict # 用户信息
email_drafts: list[str] = field(default_factory=list) # 存储历史草稿
# 邮件生成节点
@dataclass
class WriteEmail(BaseNode[State]):
feedback: str | None = None
async def run(self, ctx: GraphRunContext[State]) -> Feedback:
writer_agent = Agent(
'openai:gpt-4o',
output_type=str,
system_prompt=f"根据用户信息写邮件: {ctx.state.user}"
)
prompt = f"用户信息: {ctx.state.user}" + (f"修改建议: {self.feedback}" if self.feedback else "")
draft = await writer_agent.run(prompt)
ctx.state.email_drafts.append(draft)
return Feedback(draft)
# 反馈节点
@dataclass
class Feedback(BaseNode[State, None, str]):
email: str
async def run(self, ctx: GraphRunContext[State]) -> WriteEmail | End[str]:
critic_agent = Agent(
'google-gla:gemini-1.5-pro',
output_type=str,
system_prompt="评估邮件专业性和相关性，提供具体修改建议"
)
feedback = await critic_agent.run(f"评估邮件: {self.email}")
if "通过" in feedback: # 简单判断是否接受
return End(self.email)
return WriteEmail(feedback=feedback)
# 创建并运行图
email_graph = Graph(nodes=[WriteEmail, Feedback])
initial_state = State(user={"name": "张三", "interests": ["AI", "Python"]})
result = await email_graph.run(WriteEmail(), state=initial_state)
print(f"最终邮件: {result.output}")
图3：邮件生成-反馈循环流程图（完整代码），虚线表示条件跳转
在Jupyter中实时监控
配合Gradio界面可实现更直观的交互，但在Jupyter中也能通过以下代码实时查看状态变化：
# 安装必要依赖
!pip install ipywidgets
import ipywidgets as widgets
from IPython.display import display
# 创建状态显示组件
state_display = widgets.Textarea(description="当前状态:")
display(state_display)
# 实时更新状态
async with email_graph.iter(WriteEmail(), state=initial_state) as run:
async for node in run:
state_display.value = f"执行节点: {node}\n用户信息: {run.state.user}\n最新草稿: {run.state.email_drafts[-1][:50]}..."
高级特性与最佳实践
持久化到文件系统
对于长时间运行的实验，使用文件持久化确保中断后可恢复：
from pydantic_graph.persistence.file import FileStatePersistence
# 创建文件持久化器
persistence = FileStatePersistence("email_workflow.json")
# 保存当前状态
async with email_graph.iter(WriteEmail(), state=initial_state, persistence=persistence) as run:
... # 执行工作流
# 后续可从文件恢复
async with email_graph.iter_from_persistence(persistence) as resumed_run:
async for node in resumed_run:
print(f"恢复执行: {node}")
结合Mermaid扩展自定义样式
通过修改Mermaid配置自定义流程图样式：
mermaid_config = {
"theme": "dark",
"nodeStyle": {"fill": "#2563eb", "stroke": "#0f172a"},
"edgeStyle": {"stroke": "#64748b", "strokeWidth": 2}
}
display(Image(email_graph.mermaid_image(
start_node=WriteEmail,
config=mermaid_config
)))
图4：使用自定义Mermaid配置生成的深色主题流程图
常见问题与解决方案
Q: 如何处理循环依赖的节点定义？
A: 使用字符串引用或from __future__ import annotations，示例：
from __future__ import annotations # 放在文件顶部
@dataclass
class NodeA(BaseNode):
async def run(self, ctx: GraphRunContext) -> NodeB: # 此时NodeB尚未定义
return NodeB()
@dataclass
class NodeB(BaseNode):
async def run(self, ctx: GraphRunContext) -> NodeA:
return NodeA()
Q: 如何在Jupyter中实现交互式节点调试？
A: 结合ipdb实现断点调试：
@dataclass
class DebugNode(BaseNode):
async def run(self, ctx: GraphRunContext):
import ipdb; ipdb.set_trace() # 执行到此时会暂停
return NextNode()
总结与后续学习路径
通过本文介绍的Pydantic-Graph方案，你已掌握在Jupyter中实现AI工作流可视化的核心技能。关键要点回顾：
- 图抽象：用
BaseNode子类定义状态节点，通过返回类型声明边关系 - 状态管理：利用
FullStatePersistence追踪执行历史，确保实验可复现 - 可视化：
mermaid_image()在Jupyter中一键渲染流程图 - AI集成：与Pydantic-AI的Agent无缝协作，构建智能工作流
下一步推荐学习：
收藏本文并关注项目更新，下次将带来"分布式图执行与Temporal集成"的深度教程。如有疑问，欢迎在项目仓库提交issue或参与讨论。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
