50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
import sys
|
|
from langgraph.graph import StateGraph, START
|
|
from langgraph.types import interrupt, Command
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from typing import TypedDict, Annotated
|
|
from langgraph.graph.message import add_messages
|
|
import questionary
|
|
|
|
class GraphState(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
human_value: str | None
|
|
|
|
async def interrupt_node(state: GraphState) -> GraphState:
|
|
# Trigger interrupt asking user to confirm
|
|
payload = {
|
|
"type": "confirm",
|
|
"question": "Do you want to continue?",
|
|
"options": ["yes", "no"],
|
|
}
|
|
# interrupt returns None, graph pauses
|
|
interrupt(payload)
|
|
# After resume, payload will have 'answer'
|
|
answer = state.get("human_value")
|
|
return {"messages": state["messages"], "human_value": answer}
|
|
|
|
builder = StateGraph(GraphState)
|
|
builder.add_node("interrupt_node", interrupt_node)
|
|
builder.add_edge(START, "interrupt_node")
|
|
builder.add_edge("interrupt_node", START)
|
|
graph = builder.compile(checkpointer=InMemorySaver())
|
|
|
|
async def main():
|
|
config = {"configurable": {"thread_id": "session-1"}}
|
|
# Start stream
|
|
async for chunk in graph.stream({"messages": [], "human_value": None}, config):
|
|
if "__interrupt__" in chunk:
|
|
payload = chunk["__interrupt__"][0].value
|
|
answer = questionary.select(payload["question"], choices=payload["options"]).ask()
|
|
# resume with answer
|
|
payload["answer"] = answer
|
|
# store answer in state
|
|
chunk["human_value"] = answer
|
|
await graph.stream(Command(resume=payload), config)
|
|
else:
|
|
print(chunk)
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
asyncio.run(main())
|