144 lines
4.5 KiB
Python
144 lines
4.5 KiB
Python
"""
|
||
LangGraph 循环 - 智能体的核心能力
|
||
构建一个简单的"尝试-评估-重试"智能体
|
||
"""
|
||
from langgraph.graph import StateGraph, START, END
|
||
from typing import TypedDict
|
||
|
||
# 1️⃣ 定义状态 - 累积信息
|
||
class AgentState(TypedDict):
|
||
task: str # 任务
|
||
attempt: int # 尝试次数
|
||
result: str # 当前结果
|
||
quality: str # 质量评估: "good" / "needs_improvement" / "bad"
|
||
feedback: str # 改进建议
|
||
|
||
# 2️⃣ 定义节点
|
||
def plan_node(state: AgentState):
|
||
"""规划节点 - 制定初始方案"""
|
||
print(f"\n[规划] 任务: {state['task']}")
|
||
print(f"[规划] 制定初始方案...")
|
||
|
||
# 模拟规划逻辑
|
||
if "诗" in state['task']:
|
||
state['result'] = "床前明月光"
|
||
state['quality'] = "needs_improvement"
|
||
elif "代码" in state['task'] or "code" in state['task'].lower():
|
||
state['result'] = "print(hello)" # 有 bug
|
||
state['quality'] = "bad"
|
||
else:
|
||
state['result'] = "这是一个初步方案"
|
||
state['quality'] = "needs_improvement"
|
||
|
||
state['attempt'] = 1
|
||
print(f"[规划] 第 {state['attempt']} 次尝试: {state['result']}")
|
||
return state
|
||
|
||
def execute_node(state: AgentState):
|
||
"""执行节点 - 根据反馈改进"""
|
||
state['attempt'] += 1
|
||
print(f"\n[执行] 第 {state['attempt']} 次尝试...")
|
||
print(f"[执行] 收到反馈: {state['feedback']}")
|
||
|
||
# 模拟根据反馈改进
|
||
if "诗" in state['task']:
|
||
if state['attempt'] == 2:
|
||
state['result'] = "床前明月光，疑是地上霜"
|
||
else:
|
||
state['result'] = "床前明月光，疑是地上霜。举头望明月，低头思故乡。"
|
||
elif "代码" in state['task'] or "code" in state['task'].lower():
|
||
if state['attempt'] == 2:
|
||
state['result'] = "print('hello')" # 修复了
|
||
else:
|
||
state['result'] = "print('hello, world!') # 完美!"
|
||
|
||
print(f"[执行] 改进后: {state['result']}")
|
||
return state
|
||
|
||
def evaluate_node(state: AgentState):
|
||
"""评估节点 - 判断质量"""
|
||
print(f"\n[评估] 检查第 {state['attempt']} 次结果...")
|
||
|
||
# 模拟评估逻辑
|
||
if state['attempt'] >= 3:
|
||
state['quality'] = "good"
|
||
state['feedback'] = "质量满意，通过！"
|
||
print(f"[评估] 结果: {state['quality']} - {state['feedback']}")
|
||
else:
|
||
if "诗" in state['task']:
|
||
state['feedback'] = "内容太短，需要更完整"
|
||
elif "代码" in state['task'] or "code" in state['task'].lower():
|
||
state['feedback'] = "代码有语法错误，需要修复"
|
||
else:
|
||
state['feedback'] = "需要更详细的内容"
|
||
state['quality'] = "needs_improvement"
|
||
print(f"[评估] 结果: {state['quality']} - {state['feedback']}")
|
||
|
||
return state
|
||
|
||
# 3️⃣ 路由函数 - 决定是继续循环还是结束
|
||
def should_retry(state: AgentState):
|
||
"""⭐ 核心：决定是重试还是结束"""
|
||
if state['quality'] == "good":
|
||
return "end" # 质量够了，结束
|
||
elif state['attempt'] >= 3:
|
||
return "end" # 达到最大重试次数，强制结束
|
||
else:
|
||
return "retry" # 继续改进
|
||
|
||
# 4️⃣ 构建图
|
||
graph = StateGraph(AgentState)
|
||
|
||
# 添加节点
|
||
graph.add_node("plan", plan_node)
|
||
graph.add_node("execute", execute_node)
|
||
graph.add_node("evaluate", evaluate_node)
|
||
|
||
# 添加边 - 注意这里的循环！
|
||
graph.add_edge(START, "plan") # 开始 -> 规划
|
||
graph.add_edge("plan", "evaluate") # 规划 -> 评估
|
||
graph.add_conditional_edges(
|
||
"evaluate",
|
||
should_retry,
|
||
{
|
||
"retry": "execute", # 需要改进 -> 执行(改进)
|
||
"end": END, # 满意 -> 结束
|
||
}
|
||
)
|
||
graph.add_edge("execute", "evaluate") # 执行完再评估 -> 形成循环！
|
||
|
||
# 编译
|
||
app = graph.compile()
|
||
|
||
# 5️⃣ 测试
|
||
print("=" * 55)
|
||
print("循环演示 - 智能体不断尝试直到满意")
|
||
print("=" * 55)
|
||
|
||
# 测试 1: 写诗
|
||
print("\n" + "=" * 55)
|
||
print("测试 1: 写一首诗")
|
||
print("=" * 55)
|
||
result = app.invoke({
|
||
"task": "写一首诗",
|
||
"attempt": 0,
|
||
"result": "",
|
||
"quality": "",
|
||
"feedback": ""
|
||
})
|
||
print(f"\n最终结果: {result['result']}")
|
||
print(f"尝试次数: {result['attempt']}")
|
||
|
||
# 测试 2: 写代码
|
||
print("\n" + "=" * 55)
|
||
print("测试 2: 写一段代码")
|
||
print("=" * 55)
|
||
result = app.invoke({
|
||
"task": "写一段代码",
|
||
"attempt": 0,
|
||
"result": "",
|
||
"quality": "",
|
||
"feedback": ""
|
||
})
|
||
print(f"\n最终结果: {result['result']}")
|
||
print(f"尝试次数: {result['attempt']}") |