87 lines
2.3 KiB
Python
87 lines
2.3 KiB
Python
"""Human-in-the-loop: кастомное прерывание в LangGraph."""
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import TypedDict
|
|
|
|
import questionary
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.constants import START
|
|
from langgraph.graph import StateGraph
|
|
from langgraph.types import Command, interrupt
|
|
|
|
|
|
class State(TypedDict, total=False):
|
|
foo: str
|
|
human_value: str
|
|
|
|
|
|
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)
|
|
|
|
print(f"> Received an input from the interrupt: {answer}")
|
|
return {
|
|
"foo": state.get("foo", ""),
|
|
"human_value": answer,
|
|
}
|
|
|
|
|
|
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)
|
|
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_hitl() -> None:
|
|
graph = build_graph()
|
|
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
|
|
stream_input: dict | Command = {"foo": "abc"}
|
|
|
|
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)
|
|
|
|
if not interrupted:
|
|
break
|
|
|
|
for chunk in graph.stream(stream_input, config=config):
|
|
print(chunk)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_hitl()
|