update main.py
This commit is contained in:
@@ -13,106 +13,32 @@ Run:
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import os
|
import os
|
||||||
from typing import TypedDict, List, Dict, Any
|
from typing import Dict, Any
|
||||||
|
|
||||||
import questionary
|
import questionary
|
||||||
from langgraph.graph import StateGraph
|
|
||||||
from langgraph.constants import START
|
from langgraph.constants import START
|
||||||
from langgraph.types import interrupt, Command
|
from langgraph.types import Command
|
||||||
from langgraph.checkpoint.memory import InMemorySaver
|
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):
|
# The node is already defined in graph.py; we just use the compiled graph.
|
||||||
"""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
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# 2. Node that triggers an interrupt
|
# Main loop that runs the graph and handles interrupts
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
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
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
async def run_graph() -> None:
|
async def run_graph() -> None:
|
||||||
"""Execute the graph, handle the custom interrupt, and resume execution.
|
"""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.
|
|
||||||
"""
|
|
||||||
from asyncio import run # Imported lazily to keep top‑level imports minimal.
|
from asyncio import run # Imported lazily to keep top‑level imports minimal.
|
||||||
|
|
||||||
graph = build_graph()
|
graph = build_graph()
|
||||||
config = {"configurable": {"thread_id": "demo-thread-1"}}
|
config = {"configurable": {"thread_id": "demo-thread-1"}}
|
||||||
|
|
||||||
# Initial state – ``foo`` can be any value; it is not used in this demo.
|
# 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):
|
async for chunk in graph.stream(init_state, config):
|
||||||
if "__interrupt__" in chunk:
|
if "__interrupt__" in chunk:
|
||||||
interrupt_obj = chunk["__interrupt__"][0]
|
interrupt_obj = chunk["__interrupt__"][0]
|
||||||
@@ -147,7 +73,7 @@ async def run_graph() -> None:
|
|||||||
print(final_state)
|
print(final_state)
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# 6. Entry point
|
# Entry point
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
import asyncio
|
import asyncio
|
||||||
@@ -156,4 +82,3 @@ if __name__ == "__main__":
|
|||||||
asyncio.run(run_graph())
|
asyncio.run(run_graph())
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
print("\nInterrupted by user.")
|
print("\nInterrupted by user.")
|
||||||
"""
|
|
||||||
|
|||||||
Reference in New Issue
Block a user