command_flow.py

from typing_extensions import TypedDict, Literal
import random
from langgraph.graph import StateGraph, START
from langgraph.types import Command

# Define graph state
class State(TypedDict):
    foo: str

# Define the nodes
def node_a(state: State) -> Command[Literal["node_b", "node_c"]]:
    print("\nNode A:")
    print(f"  入力状態: {state}")
    value = random.choice(["a", "b"])
    print(f"  選択された値: {value}")
    
    if value == "a":
        goto = "node_b"
        print("  次のノード: B")
    else:
        goto = "node_c"
        print("  次のノード: C")

    new_state = {"foo": value}
    print(f"  更新後の状態: {new_state}")
    
    return Command(update=new_state, goto=goto)

def node_b(state: State):
    print("\nNode B:")
    print(f"  入力状態: {state}")
    new_state = {"foo": state["foo"] + "b"}
    print(f"  更新後の状態: {new_state}")
    return new_state

def node_c(state: State):
    print("\nNode C:")
    print(f"  入力状態: {state}")
    new_state = {"foo": state["foo"] + "c"}
    print(f"  更新後の状態: {new_state}")
    return new_state

def main():
    # Create the graph
    builder = StateGraph(State)
    builder.add_edge(START, "node_a")
    builder.add_node(node_a)
    builder.add_node(node_b)
    builder.add_node(node_c)
    
    # Compile the graph
    graph = builder.compile()
    
    # Run the graph multiple times to see different paths
    for i in range(3):
        print(f"\n===== 実行 {i+1} =====")
        print("初期状態: {'foo': ''}")
        result = graph.invoke({"foo": ""})
        print(f"\n最終状態: {result}")
        print("=" * 20)

if __name__ == "__main__":
    main()