109 lines
4.4 KiB
Python
109 lines
4.4 KiB
Python
import questionary
|
||
from typing import TypedDict, List, Any, Optional
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from langgraph.types import interrupt, Command
|
||
|
||
# 1. Описание состояния графа
|
||
class GraphState(TypedDict):
|
||
foo: str
|
||
human_value: Optional[str]
|
||
status: str
|
||
|
||
# 2. Узел с прерыванием
|
||
def human_approval_node(state: GraphState) -> GraphState:
|
||
"""
|
||
Узел, который ставит выполнение на паузу и ждет ответа пользователя.
|
||
При вызове interrupt() граф останавливается.
|
||
Когда вызывается Command(resume=...), значение из resume попадает в результат interrupt().
|
||
"""
|
||
# Формируем payload для прерывания
|
||
interrupt_payload = {
|
||
"type": "confirm",
|
||
"question": "Уверены, что хотите продолжить?",
|
||
"options": ["approve", "reject"]
|
||
}
|
||
|
||
# Вызываем interrupt.
|
||
# Если граф возобновлен через Command(resume=...),
|
||
# переменная 'user_response' получит значение из resume.
|
||
user_response = interrupt(interrupt_payload)
|
||
|
||
# Возвращаем обновленное состояние
|
||
return {
|
||
"human_value": user_response,
|
||
"status": "completed"
|
||
}
|
||
|
||
# 3. Сборка графа
|
||
def create_graph():
|
||
workflow = StateGraph(GraphState)
|
||
|
||
workflow.add_node("approval_node", human_approval_node)
|
||
workflow.add_edge(START, "approval_node")
|
||
workflow.add_edge("approval_node", END)
|
||
|
||
# Используем InMemorySaver для поддержки чекпоинтов (необходим для прерываний)
|
||
memory = InMemorySaver()
|
||
return workflow.compile(checkpointer=memory)
|
||
|
||
def run_workflow():
|
||
graph = create_graph()
|
||
config = {"configurable": {"thread_id": "123"}}
|
||
initial_state = {"foo": "initial_data", "human_value": None, "status": "running"}
|
||
|
||
# Запуск графа
|
||
print("--- Starting Graph Execution ---")
|
||
|
||
# Первый запуск (до прерывания)
|
||
# Используем stream_mode="updates" или "values".
|
||
# При interrupt в потоке появится специальный элемент.
|
||
events = graph.stream(initial_command_or_state(initial_state), config, stream_mode="updates")
|
||
|
||
final_state = None
|
||
|
||
# Итерируемся по событиям
|
||
for event in events:
|
||
# Проверяем наличие прерывания в потоке через ключ __interrupt__
|
||
if "__interrupt__" in event:
|
||
# Извлекаем payload из прерывания
|
||
# В LangGraph interrupt создает элемент в списке interrupts задачи
|
||
interrupt_info = event["__interrupt__"][0].value
|
||
|
||
print("\n--- Interrupt Detected ---")
|
||
print(interrupt_info)
|
||
print(f"!!! {interrupt_info['type'].upper()} !!!")
|
||
|
||
# 4. Сбор ответа пользователя через questionary
|
||
user_choice = questionary.select(
|
||
interrupt_info["question"],
|
||
choices=interrupt_info["options"]
|
||
).ask()
|
||
|
||
print(f"> Received an input from the interrupt: {user_choice}")
|
||
|
||
# 5. Возобновление графа через Command(resume=...)
|
||
# Передаем выбранный ответ обратно в граф
|
||
resume_events = graph.stream(
|
||
Command(resume=user_choice),
|
||
config,
|
||
stream_mode="updates"
|
||
)
|
||
|
||
for resume_event in resume_events:
|
||
final_state = resume_event
|
||
else:
|
||
# Обычное обновление узлов
|
||
final_state = event
|
||
|
||
# Вывод итогового состояния (из последнего события или snapshot)
|
||
print("\n--- Final State ---")
|
||
snapshot = graph.get_state(config)
|
||
print(snapshot.values)
|
||
|
||
def initial_command_or_state(initial_state):
|
||
"""Вспомогательная функция для корректного старта."""
|
||
return initial_state
|
||
|
||
if __name__ == "__main__":
|
||
run_workflow() |