update main.py

main.py

@@ -13,106 +13,32 @@ Run:
 from __future__ import annotations
 
 import os
-from typing import TypedDict, List, Dict, Any
+from typing import Dict, Any
 
 import questionary
-from langgraph.graph import StateGraph
 from langgraph.constants import START
-from langgraph.types import interrupt, Command
+from langgraph.types import Command
 from langgraph.checkpoint.memory import InMemorySaver
+from graph import build_graph
 
 # ---------------------------------------------------------------------------
-# 1. Define the graph state
+# 1. Node that triggers an interrupt (defined in graph.py)
 # ---------------------------------------------------------------------------
-class GraphState(TypedDict):
-    """Minimal state used by the example.
-
-    * ``human_value`` – value supplied by the user during the interrupt.
-    * ``foo`` – placeholder for any initial data that might be needed later.
-    """
-
-    human_value: str | None
-    foo: int
+# The node is already defined in graph.py; we just use the compiled graph.
 
 # ---------------------------------------------------------------------------
-# 2. Node that triggers an interrupt
-# ---------------------------------------------------------------------------
-async def ask_user(state: GraphState) -> GraphState:
-    """Node that pauses execution and asks the user for confirmation.
-
-    The node returns a ``Command`` with an ``interrupt`` payload.  When the graph
-    receives this command it stops, yields control to the caller, and waits for
-    a resume command.
-    """
-    # Build interrupt payload – a simple dictionary that will be shown to the user.
-    payload: Dict[str, Any] = {
-        "type": "confirm",
-        "question": "Do you want to continue?",
-        "options": ["yes", "no"],
-    }
-    # Trigger interrupt – execution stops here until a resume command is sent.
-    return Command(resume=interrupt(payload))
-
-# ---------------------------------------------------------------------------
-# 3. Node that receives the user's answer and updates state
-# ---------------------------------------------------------------------------
-async def process_answer(state: GraphState) -> GraphState:
-    """Node executed after the graph resumes.
-
-    The ``state`` will contain the resume payload under ``__interrupt__``.
-    We extract the user response, store it in ``human_value`` and return the
-    updated state.
-    """
-    # ``__interrupt__`` is a list of interrupt objects; we only use the first.
-    interrupt_obj = state.get("__interrupt__", [])[0]
-    payload: Dict[str, Any] = interrupt_obj.value  # type: ignore[assignment]
-
-    # The user response will be added to the payload under ``answer``.
-    answer = payload.get("answer")
-    if not isinstance(answer, str):
-        raise ValueError("Interrupt payload missing 'answer' field")
-
-    state["human_value"] = answer
-    return state
-
-# ---------------------------------------------------------------------------
-# 4. Build the graph
-# ---------------------------------------------------------------------------
-def build_graph() -> StateGraph:
-    """Create and compile a simple graph with an interrupt node.
-
-    The graph consists of two nodes: ``ask_user`` (which pauses) and
-    ``process_answer`` (which continues after the user responds).
-    """
-    graph = StateGraph(GraphState)
-    graph.add_node("ask", ask_user)
-    graph.add_node("answer", process_answer)
-    graph.set_entry_point("ask")
-    graph.add_edge(START, "ask")
-    graph.add_edge("ask", "answer")
-
-    # Use an in‑memory checkpoint so we can resume after the interrupt.
-    return graph.compile(checkpointer=InMemorySaver())
-
-# ---------------------------------------------------------------------------
-# 5. Main loop that runs the graph and handles interrupts
+# Main loop that runs the graph and handles interrupts
 # ---------------------------------------------------------------------------
 async def run_graph() -> None:
-    """Execute the graph, handle the custom interrupt, and resume execution.
-
-    The function streams the graph output, looks for ``__interrupt__`` chunks,
-    prompts the user with *questionary*, then resumes the graph with the
-    selected answer.
-    """
+    """Execute the graph, handle the custom interrupt, and resume execution."""
     from asyncio import run  # Imported lazily to keep top‑level imports minimal.
 
     graph = build_graph()
     config = {"configurable": {"thread_id": "demo-thread-1"}}
 
     # Initial state – ``foo`` can be any value; it is not used in this demo.
-    init_state: GraphState = {"human_value": None, "foo": 42}
+    init_state: Dict[str, Any] = {"human_value": None, "foo": 42}
 
-    # Stream the graph until completion.
     async for chunk in graph.stream(init_state, config):
         if "__interrupt__" in chunk:
             interrupt_obj = chunk["__interrupt__"][0]
@@ -147,7 +73,7 @@ async def run_graph() -> None:
     print(final_state)
 
 # ---------------------------------------------------------------------------
-# 6. Entry point
+# Entry point
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
     import asyncio
@@ -156,4 +82,3 @@ if __name__ == "__main__":
         asyncio.run(run_graph())
     except KeyboardInterrupt:
         print("\nInterrupted by user.")
-"""