LangGraph の新機能Command について
LangGraph(Python 版)のver.0.2.58 において、Command
機能がリリースされた。
旧機能と比較しながらその特徴をメモ。
Command
とは
従来のLangGraph では、
- ノードはState の差分を返す
-
add_conditional_edges
に分岐候補のノード群と、分岐する次ノードを決定する関数(router)を設定する
で、ノードの遷移先を動的に制御していた。
サンプルコード
単純な一本道のグラフを生成する。LLM も使わないので、これだけならLangGraph を使う意味は無いのだが、まず雰囲気だけ。
from langgraph.graph import StateGraph, MessagesState, START, END
class State(MessagesState):
text: str
def node1(state: State):
return {
"messages": ["こんにちは"]
}
def node2(state: State):
return {
"text": "hoge"
}
builder = StateGraph(State)
builder.add_node("node1", node1)
builder.add_node("node2", node2)
builder.add_conditional_edges(
"node1",
lambda x: "next",
{"next": "node2"}
)
builder.add_edge(START, "node1")
builder.add_edge("node2", END)
graph = builder.compile()
グラフを図示するとこうなる。
Command
機能を使うと、
- ノードは
Command
クラスのオブジェクトを返す - オブジェクトの
update
パラメータにState の差分を、goto
パラメータに分岐する次ノード(変数でOK。現ノードの処理終了までに決定すれば良い)を入力
で、ノードの遷移先を動的に制御できる。
サンプルコード
同様の、単純な一本道のグラフを生成する。
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.types import Command
class State(MessagesState):
text: str
def node1(state: State) -> Command[Literal["node2"]]:
goto="node2"
return Command(
update={
"messages": ["こんにちは"]
},
goto=goto,
)
def node2(state: State):
return Command(
update={
"text": "hoge"
}
)
builder = StateGraph(State)
builder.add_node("node1", node1)
builder.add_node("node2", node2)
builder.add_edge(START, "node1")
builder.add_edge("node2", END)
graph = builder.compile()
Command
の利点
従来のadd_conditional_edges
と比較して、Command
を使うメリットを述べる。
Command
クラスで指定できる
ノードが返す変数の型アノテーションを個人的にはこれが一番大きい。以前はノードが返す変数はState の差分であったため、型アノテーションが書きにくかった。
なおgraph.get_graph()
する際にこの型アノテーションを参照しているようで、(実際の分岐先である、ノードが返すパラメータgoto
に関わらず)図示されるグラフの分岐先となる。
参考
ノードの関数はサンプルコードと同様(=動作する際の制御フローは同じ)だが、返り値の型アノテーションを滅茶苦茶に書いた。
def node1(state: State) -> Command[Literal["node2", END]]:
return Command(
update={
"messages": ["こんにちは"]
},
goto="node2", # 次のノードは必ずnode2
)
def node2(state: State) -> Command[Literal["node1"]]:
return Command(
update={
"text": "hoge"
}
)
ところがdisplay(Image(graph.get_graph().draw_mermaid_png()))
で図示すると、型アノテーション通りのものが表示されてしまう。これはちょっと微妙な点かも(State の型アノテーションは動作に関係するので...)。
子グラフの遷移先に親グラフのノードを指定できる
よく言われてるやつ。マルチエージェントを実装する際に、親グラフから子グラフ(サブグラフ)を呼び出すことが想定されるが、その自由度が上がった。
まず子グラフgraph_child
を定義し、そのノードnode2_child
の遷移先を親グラフのノードnode2_parent
とする。
from typing import Literal, Annotated
class State(MessagesState):
text: Annotated[str, lambda old, new: new]
def node1_child(state: State) -> Command[Literal["node2_child"]]:
return Command(
update={
"messages": ["こんにちは"]
},
goto="node2_child",
)
def node2_child(state: State):
return Command(
update={
"text": "hoge"
},
goto="node2_parent",
graph=Command.PARENT,
)
builder_child = StateGraph(State)
builder_child.add_node("node1_child", node1_child)
builder_child.add_node("node2_child", node2_child)
builder_child.add_edge(START, "node1_child")
graph_child = builder_child.compile()
この子グラフを、node1_parent
から呼び出すような親グラフgraph_parent
を定義する。
def node1_parent(state: State) -> Command[Literal["graph_child"]]:
return Command(
update={
"messages": ["こんばんは"]
},
goto="graph_child",
)
def node2_parent(state: State):
return Command(
update={
"text": "fuga"
}
)
builder_parent = StateGraph(State)
builder_parent.add_node("node1_parent", node1_parent)
builder_parent.add_node("node2_parent", node2_parent)
builder_parent.add_node("graph_child", graph_child)
builder_parent.add_edge(START, "node1_parent")
builder_parent.add_edge("node2_parent", END)
graph_parent = builder_parent.compile()
このとき、node1_parent
→ (親グラフ側の指定) → graph_child
(node1_child
→ node2_child
)→ (子グラフ側の指定) → node2_parent
のような制御になる。
for s in graph_parent.stream(
{"text": "piyo"},
{"recursion_limit": 100},
):
print(s)
print("---")
{'node1_parent': {'messages': ['こんばんは']}}
---
{'graph_child': [{'messages': [HumanMessage(content='こんばんは', additional_kwargs={}, response_metadata={}, id='73dde28b-7cda-4f6d-877a-d0d9868882c9'), HumanMessage(content='こんにちは', additional_kwargs={}, response_metadata={}, id='173ad002-5644-4de0-bcb3-1a7554d7c875')]}, {'text': 'piyo'}, {'text': 'hoge'}]}
---
{'node2_parent': {'text': 'fuga'}}
---
Command
を使わない実装では、子グラフの処理終了 →(親グラフ側の指定)→ 親グラフの次ノード しか出来なかったので、単純に便利になった。
追記
子グラフに親グラフと別のState を入力していても、graph=
の機能を使って親グラフに戻すと、子グラフのState が親グラフに反映されるので注意。
親グラフ
def node_graph_child(state: State):
graph_child.invoke({"messages": "おはようございます", "text": "hogefuga"})
builder_parent = StateGraph(State)
builder_parent.add_node("node1_parent", node1_parent)
builder_parent.add_node("node2_parent", node2_parent)
builder_parent.add_node("node_graph_child", node_graph_child)
builder_parent.add_edge(START, "node1_parent")
builder_parent.add_edge("node2_parent", END)
graph_parent = builder_parent.compile()
実行結果
{'node1_parent': {'messages': ['こんばんは']}}
---
{'node_graph_child': [{'messages': [HumanMessage(content='おはようございます', additional_kwargs={}, response_metadata={}, id='48b8d0be-bed7-42ff-ac5d-7dc088f141e0'), HumanMessage(content='こんにちは', additional_kwargs={}, response_metadata={}, id='af63375c-6b66-4a47-9252-942701d26a71')]}, {'text': 'hoge'}, {'text': 'hogehoge'}]}
---
{'node2_parent': {'text': 'fuga'}}
---
当然だが、子グラフの葉ノード(node3_child
)でgraph
とgoto
を指定せず、親グラフ側で戻した場合(builder_parent.add_edge("node_graph_child", "node2_parent")
)、子グラフでのState 情報は親グラフに反映されない。
他にも見つかったら追記予定。
余談
add_edge
とgoto
が両方設定しているとき、add_edge
による次ノード指定と、パラメータgoto
による遷移先指定では、前者が優先される。ただしその場合でも、update
によるState 更新は実施される。
動作確認コード
add_edge
が優先されるなら無限ループし、goto
が優先されるなら正常終了するグラフを作成。
class State(MessagesState):
text: str
def node1(state: State) -> Command[Literal["node2"]]:
return Command(
update={
"messages": ["こんにちは"]
},
goto="node2",
)
def node2(state: State):
return Command(
update={
"text": "hoge"
},
goto=END, #goto ではEND を指定
)
builder = StateGraph(State)
builder.add_node("node1", node1)
builder.add_node("node2", node2)
builder.add_edge(START, "node1")
builder.add_edge("node2", "node1") #add_edge ではnode1 を指定
graph = builder.compile()
for s in graph.stream(
{"messages": [("user", "Hello, wolrd!")]},
{"recursion_limit": 100},
):
print(s)
print("---")
実行したところ無限ループしたので、add_edge
が優先される仕様であることが分かった。またnode2
のupdate
によるState 更新も行われていることを確認。
{'node1': {'messages': ['こんにちは']}}
---
{'node2': {'text': 'hoge'}}
---
{'node1': {'messages': ['こんにちは']}}
---
{'node2': {'text': 'hoge'}}
---
...
Discussion