diff --git a/main.py b/main.py index 2e7aaf8..2405acc 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ -"""Human-in-the-loop: кастомное прерывание (interrupt / resume) в LangGraph.""" +"""Human-in-the-loop: кастомное прерывание в LangGraph.""" from __future__ import annotations +import uuid from typing import TypedDict import questionary @@ -10,83 +11,76 @@ from langgraph.graph import StateGraph from langgraph.types import Command, interrupt -class GraphState(TypedDict): - """Состояние графа: начальные данные и ответ пользователя.""" - +class State(TypedDict, total=False): foo: str human_value: str -def human_node(state: GraphState) -> GraphState: - """Узел с кастомным прерыванием — ждёт ответ пользователя.""" - payload = interrupt( - { - "type": "alert", - "question": "Уверены что хотите продолжить?", - "allow_responds": ["approve", "reject"], - } - ) +def human_node(state: State) -> dict: + payload = { + "type": "alert", + "question": "Уверены, что хотите продолжить?", + "allow_responds": ["approve", "reject"], + } + resumed = interrupt(payload) + + if isinstance(resumed, dict): + answer = resumed.get("answer", "") + else: + answer = str(resumed) - answer = payload.get("answer", "") - print(f"!!! {payload.get('type', 'alert')} !!!") print(f"> Received an input from the interrupt: {answer}") - return { "foo": state.get("foo", ""), "human_value": answer, } -def _ask_user(payload: dict) -> dict: - """Показать вопрос и записать ответ в payload.""" +def build_graph(): + builder = StateGraph(State) + builder.add_node("node", human_node) + builder.add_edge(START, "node") + return builder.compile(checkpointer=InMemorySaver()) + + +def handle_interrupt(interrupts: tuple) -> dict: + first = interrupts[0] + payload = dict(first.value if hasattr(first, "value") else first) + + print("Произошла остановка") print(payload) - options = payload.get("allow_responds") or ["approve", "reject"] - question = payload.get("question", "Выберите вариант:") - choice = questionary.select(question, choices=options).ask() - if choice is None: - choice = options[0] - updated = dict(payload) - updated["answer"] = choice - return updated + print(f"!!! {payload.get('type', 'alert')} !!!") + + choice = questionary.select( + payload["question"], + choices=payload["allow_responds"], + ).ask() + + payload["answer"] = choice or payload["allow_responds"][0] + return payload -def run_graph() -> GraphState: - """Запуск графа с обработкой прерывания и возобновлением.""" - builder = StateGraph(GraphState) - builder.add_node("human_node", human_node) - builder.add_edge(START, "human_node") +def run_hitl() -> None: + graph = build_graph() + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + stream_input: dict | Command = {"foo": "abc"} - graph = builder.compile(checkpointer=InMemorySaver()) - config = {"configurable": {"thread_id": "hitl-interrupt-demo"}} - initial: GraphState = {"foo": "initialized", "human_value": ""} - - final_state: GraphState = initial - - stream = graph.stream(initial, config) - for chunk in stream: - if "__interrupt__" in chunk: - print("Произошла остановка") - interrupt_tuple = chunk["__interrupt__"] - payload = dict(interrupt_tuple[0].value) - resumed_payload = _ask_user(payload) - - resume_stream = graph.stream(Command(resume=resumed_payload), config) - for resume_chunk in resume_stream: - if "human_node" in resume_chunk: - final_state = resume_chunk["human_node"] - print({"node": resume_chunk}) - elif "human_node" in chunk: - final_state = chunk["human_node"] + while True: + interrupted = False + for chunk in graph.stream(stream_input, config=config): + if "__interrupt__" in chunk: + payload = handle_interrupt(chunk["__interrupt__"]) + stream_input = Command(resume=payload) + interrupted = True + break print(chunk) - return final_state + if not interrupted: + break - -def main() -> None: - result = run_graph() - print("\n=== Итоговое состояние ===") - print(result) + for chunk in graph.stream(stream_input, config=config): + print(chunk) if __name__ == "__main__": - main() + run_hitl()