langgraph 循环节点的静态和动态实现

32 阅读1分钟

ReAct 循环节点的两种实现方法,静态实现和动态实现,current小于5的话就不断累加

1.静态

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import MessagesState
from rich import print as rprint

class overAllState(MessagesState):
    current: int

def add_fn(state: overAllState):
    num = state["current"]
    print("add_fn", num)
    return {"current": num + 1}

def route(state: overAllState):
    if state["current"] < 5:
        return "add_fn"
    return END

builder = StateGraph(overAllState)
builder.add_node("add_fn", add_fn)

# 起点直接进 add_fn
builder.add_edge(START, "add_fn")

# add_fn 执行完后,用条件边决定:回自己 或 结束
builder.add_conditional_edges(
    "add_fn",
    route,
    path_map={
        "add_fn": "add_fn",
        END: END,
    },
)

graph = builder.compile()
display(graph)
res = graph.invoke({"current": 0})
rprint(res)

2,动态实现

from langgraph.graph import StateGraph, START, END
from langgraph.types import Command
from langgraph.graph.message import MessagesState
from rich import print as rprint
from typing import Literal

class overAllState(MessagesState):
    current: int

def add_fn(state: overAllState) -> Command[Literal["add_fn", "__end__"]]:
    num = state["current"]
    print("add_fn", num)

    if num < 5:
        # 更新状态并自循环回到当前节点
        return Command(
            update={"current": num + 1},
            goto="add_fn"   # 指向自身,形成循环
        )
    else:
        # 达到条件,不更新状态(或更新后)直接结束
        return Command(goto=END)

builder = StateGraph(overAllState)
builder.add_node("add_fn", add_fn)
builder.add_edge(START, "add_fn")  # 起点直接进入循环节点

graph = builder.compile()
display(graph)
res = graph.invoke({"current": 0})
rprint(res)