diff --git a/main.py b/main.py
index bda4c38..79afd09 100644
--- a/main.py
+++ b/main.py
@@ -9,10 +9,11 @@ from langgraph.graph.message import add_messages
 
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
-from langchain.tools import tool
-
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+from langchain.tools import tool
+
+import questionary
 
 # ---------- LLM ----------
 llm = ChatOpenAI(
@@ -23,18 +24,20 @@ llm = ChatOpenAI(
 )
 
 # ---------- Backend for deepagents ----------
-backend = CompositeBackend([
-    LocalShellBackend(workspace_dir="./workspace"),
-    FilesystemBackend(),
-])
+backend = CompositeBackend(
+    [
+        LocalShellBackend(workspace_dir="./workspace"),
+        FilesystemBackend(),
+    ]
+)
 
-# ---------- Example tool ----------
+# ---------- Example tool (not used in this task but required for agent creation) ----------
 @tool
 def echo_tool(text: str) -> str:
     """Return the same text back."""
     return text
 
-# ---------- DeepAgent ----------
+# ---------- DeepAgent (required by the course) ----------
 agent = create_deep_agent(
     model=llm,
     tools=[echo_tool],
@@ -47,89 +50,65 @@ class GraphState(TypedDict):
     messages: Annotated[List[Any], add_messages]
     human_value: str
 
-# ---------- Node that triggers interrupt ----------
-def ask_human_node(state: GraphState):
-    # Prepare payload for interrupt
+# ---------- Node that triggers a custom interrupt ----------
+def ask_node(state: GraphState):
+    # If we are resumed, the payload will contain the answer
+    if isinstance(state, dict) and "answer" in state:
+        # Store the answer and finish
+        return {
+            "messages": state.get("messages", []),
+            "human_value": state["answer"],
+        }
+    # Otherwise raise an interrupt with the question payload
     payload = {
         "type": "confirm",
-        "question": "Do you want to continue?",
+        "question": "Do you want to continue the workflow?",
         "options": ["approve", "reject"],
     }
-    # Raise interrupt; execution will pause here
     return interrupt(payload)
 
-# ---------- Node after resume ----------
-def after_human_node(state: GraphState):
-    # The resumed payload will contain the answer under key 'answer'
-    answer = state.get("human_value", "no answer")
-    # Use deepagent to produce a final message
-    result = asyncio.run(
-        agent.ainvoke(
-            {"messages": [HumanMessage(content=f"User answered: {answer}")]},
-            {"configurable": {"thread_id": "deepagent-session"}},
-        )
-    )
-    # Append the agent's response to the message list
-    state["messages"].append(result["messages"][-1])
-    return state
+# ---------- Build the graph ----------
+graph_builder = StateGraph(GraphState)
+graph_builder.add_node("ask", ask_node)
+graph_builder.add_edge(START, "ask")
+graph_builder.add_edge("ask", END)
 
-# ---------- Build graph ----------
-graph = StateGraph(GraphState)
+graph = graph_builder.compile(checkpointer=InMemorySaver())
 
-graph.add_node("ask_human", ask_human_node)
-graph.add_node("after_human", after_human_node)
-
-graph.add_edge(START, "ask_human")
-graph.add_edge("ask_human", "after_human")
-graph.add_edge("after_human", END)
-
-# Use in-memory checkpointing
-graph.set_checkpoint_saver(InMemorySaver())
-app = graph.compile()
-
-# ---------- Runtime loop handling interrupt ----------
-async def run():
-    thread_id = "example-thread"
+# ---------- Execution loop handling interrupts ----------
+async def run_graph():
+    thread_id = "demo-thread"
     config = {"configurable": {"thread_id": thread_id}}
 
-    # Initial stream
-    stream = app.stream(
-        {"messages": []},
-        config,
-    )
+    # Initial state
+    state: GraphState = {"messages": [], "human_value": ""}
 
-    async for chunk in stream:
-        # Check for interrupt signal
-        if "__interrupt__" in chunk:
-            interrupt_payload = chunk["__interrupt__"][0].value
-            print("\n--- Interrupt received ---")
-            print(f"Type: {interrupt_payload.get('type')}")
-            print(f"Question: {interrupt_payload.get('question')}")
-            # Simple console input (could use questionary)
-            while True:
-                answer = input(f"Choose {interrupt_payload.get('options')}: ").strip()
-                if answer in interrupt_payload.get("options"):
-                    break
-                print("Invalid option, try again.")
-            # Add answer to payload
-            interrupt_payload["answer"] = answer
+    # Helper to process a stream until it finishes or hits an interrupt
+    async def process_stream(initial):
+        async for chunk in graph.astream(initial, config):
+            # Detect interrupt
+            if "__interrupt__" in chunk:
+                return chunk  # return the interrupt chunk
+            # Detect final state (contains human_value)
+            if "human_value" in chunk:
+                print("Workflow finished. Final state:")
+                print(chunk)
+                return None
+        return None
 
-            # Resume graph with the updated payload
-            resume_cmd = Command(resume=interrupt_payload)
-            resume_stream = app.stream(resume_cmd, config)
+    # First run - will hit the interrupt
+    interrupt_chunk = await process_stream(state)
+    while interrupt_chunk:
+        payload = interrupt_chunk["__interrupt__"][0].value
+        print("\n--- Human in the loop ---")
+        answer = questionary.select(
+            payload["question"], choices=payload["options"]
+        ).ask()
+        # Add answer to payload for resumption
+        payload["answer"] = answer
+        # Resume the graph with the updated payload
+        interrupt_chunk = await process_stream(Command(resume=payload))
 
-            async for resume_chunk in resume_stream:
-                if "__interrupt__" in resume_chunk:
-                    # Should not happen in this simple example
-                    continue
-                if "messages" in resume_chunk:
-                    # Final state reached
-                    final_state = resume_chunk
-                    print("\n--- Final state ---")
-                    for msg in final_state["messages"]:
-                        print(msg.content)
-                    return
-
-# ---------- Entry point ----------
+# ---------- Main entry ----------
 if __name__ == "__main__":
-    asyncio.run(run())
\ No newline at end of file
+    asyncio.run(run_graph())
\ No newline at end of file