67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
from typing import TypedDict, Dict, Any
|
|
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
|
|
class GraphState(TypedDict):
|
|
human_value: str | None
|
|
__interrupt__: dict[str, Any] | None # for interrupt payload
|
|
__resume__: str | None # for resume payload
|
|
|
|
|
|
def interrupt_node(state: GraphState) -> Dict[str, Any]:
|
|
"""Trigger an interrupt asking the user to choose a value."""
|
|
return {
|
|
"__interrupt__": {
|
|
"type": "question",
|
|
"question": "Choose a value:",
|
|
"options": ["Option A", "Option B", "Option C"],
|
|
}
|
|
}
|
|
|
|
|
|
def resume_node(state: GraphState) -> Dict[str, Any]:
|
|
"""Store the user's choice and finish."""
|
|
# The chosen option is passed via the __resume__ key
|
|
if state.get("__resume__") is not None:
|
|
state["human_value"] = state["__resume__"]
|
|
return {"__end__": True}
|
|
|
|
|
|
builder = StateGraph(GraphState)
|
|
builder.add_node("interrupt", interrupt_node)
|
|
builder.add_node("resume", resume_node)
|
|
|
|
# The graph starts with the interrupt node
|
|
builder.set_entry_point("interrupt")
|
|
|
|
# After a successful resume, we end the graph
|
|
builder.add_edge(START, "interrupt")
|
|
builder.add_conditional_edges(
|
|
"interrupt",
|
|
lambda x: "__interrupt__" in x,
|
|
{"__interrupt__": "resume"},
|
|
)
|
|
builder.add_edge("resume", END)
|
|
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
# Run the graph and handle interrupts
|
|
state: GraphState = {"human_value": None, "__interrupt__": None, "__resume__": None}
|
|
while True:
|
|
result = graph.invoke(state)
|
|
if "__interrupt__" in result and result["__interrupt__"] is not None:
|
|
interrupt_info = result["__interrupt__"]
|
|
print(interrupt_info["question"])
|
|
for idx, opt in enumerate(interrupt_info["options"], 1):
|
|
print(f"{idx}. {opt}")
|
|
choice_idx = int(input("Enter number: ")) - 1
|
|
chosen = interrupt_info["options"][choice_idx]
|
|
# Resume the graph with the chosen value
|
|
state = graph.invoke(state, {"__resume__": chosen})
|
|
else:
|
|
# Graph finished
|
|
break
|
|
|
|
print("Final state:", state) |