gqt
85 lines
1.8 KiB
Python
85 lines
1.8 KiB
Python
from typing import TypedDict
|
|
from loguru import logger
|
|
import time
|
|
from langgraph.types import interrupt, Command
|
|
|
|
|
|
class OverallState(TypedDict):
|
|
initial_input: str
|
|
node_a1: str
|
|
node_a2: str
|
|
node_b: str
|
|
|
|
|
|
def node_a1(state: OverallState) -> OverallState:
|
|
time.sleep(1)
|
|
logger.info("node_a1")
|
|
return OverallState(node_a1="node_a1")
|
|
|
|
|
|
def node_a2(state: OverallState) -> OverallState:
|
|
time.sleep(3)
|
|
logger.info("node_a2")
|
|
|
|
return OverallState(node_a2="node_a2")
|
|
|
|
|
|
def node_b(state: OverallState) -> OverallState:
|
|
time.sleep(1)
|
|
res = interrupt("hello")
|
|
logger.info("node_b")
|
|
return OverallState(node_b=f"node_b,res={res}")
|
|
|
|
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
builder = StateGraph(state_schema=OverallState)
|
|
builder.add_node("node_a1", node_a1)
|
|
builder.add_node("node_a2", node_a2)
|
|
builder.add_node("node_b", node_b)
|
|
|
|
builder.add_edge(START, "node_a1")
|
|
builder.add_edge(START, "node_a2")
|
|
builder.add_edge(["node_a1", "node_a2"], "node_b")
|
|
builder.add_edge("node_b", END)
|
|
|
|
checkpoint_saver = InMemorySaver()
|
|
|
|
graph = builder.compile(checkpointer=checkpoint_saver)
|
|
|
|
from IPython.display import display
|
|
|
|
display(graph)
|
|
|
|
from rich import print as rp
|
|
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": "12345"
|
|
}
|
|
}
|
|
|
|
for chunk in graph.stream(
|
|
{"initial_input": "init"},
|
|
stream_mode=["checkpoints"],
|
|
config=config
|
|
):
|
|
print("=" * 50)
|
|
rp(chunk)
|
|
|
|
print("=" * 50)
|
|
rp(list(graph.get_state_history(config=config)))
|
|
|
|
print("**" * 50)
|
|
|
|
for chunk in graph.stream(
|
|
Command(resume="interrupt_node_b_res"),
|
|
stream_mode=["checkpoints"],
|
|
config=config
|
|
):
|
|
print("=" * 50)
|
|
rp(chunk)
|
|
print("=" * 50)
|
|
rp(list(graph.get_state_history(config=config)))
|