from langgraph.graph import StateGraph, END
from langgraph.types import Command, interrupt
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict


# Define your state
class AgentState(TypedDict):
    messages: list


# Define a node that requires human approval
def human_approval_node(state: AgentState):
    # Prompt the human for approval
    print("==== flag ====")  # 中断恢复时会从该函数开头重新执行
    response = interrupt("Do you approve this action? (yes/no)")
    # Process the human's response
    if response.lower() == "yes":
        return {"messages": state["messages"] + ["Human approved."]}
    else:
        return {"messages": state["messages"] + ["Human denied."]}


# Build the graph
builder = StateGraph(AgentState)
builder.add_node("action_to_approve", lambda x: {"messages": x["messages"] + ["Performing action to be approved."]})
builder.add_node("human_approval", human_approval_node)
builder.add_node("final_step", lambda x: {"messages": x["messages"] + ["Action completed or denied."]})

builder.add_edge("action_to_approve", "human_approval")
builder.add_edge("human_approval", "final_step")
builder.set_entry_point("action_to_approve")
builder.set_finish_point("final_step")

# Configure the graph with a checkpointer
memory = InMemorySaver()
graph = builder.compile(checkpointer=memory)

# Initial state
thread_id = {"configurable": {"thread_id": "1"}}
initial_state = AgentState(messages=["Starting workflow."])

# Run the graph up to the interrupt
print("Running graph...")
for s in graph.stream(initial_state, thread_id):
    print(s)

# Output will show the interrupt message and pause.
# The human then provides input, for example, "yes".

# Resume the graph with the human's input using Command

user_input = input().strip()  # 模拟用户输入
print("\nResuming graph with human input...")
for s in graph.stream(Command(resume=f"{user_input}"), thread_id):
    print(s)