73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
from langgraph.graph import StateGraph, START, END
|
|
from state import ReflectState
|
|
from nodes import draft_answer, reflect, rewrite
|
|
|
|
def should_continue(state: ReflectState) -> str:
|
|
"""Функция-условие: продолжать или завершить"""
|
|
# Проверяем вердикт
|
|
if state["verdict"] == "ok":
|
|
print(f"✅ Ответ принят после {state['round']} раундов")
|
|
return "end"
|
|
|
|
# Проверяем лимит раундов
|
|
if state["round"] >= state["max_rounds"]:
|
|
print(f"⚠️ Достигнут лимит раундов ({state['max_rounds']}). Возвращаем последнюю версию.")
|
|
return "end"
|
|
|
|
# Требуется доработка
|
|
print(f"🔄 Запускаем раунд доработки {state['round'] + 1}/{state['max_rounds']}")
|
|
return "rewrite"
|
|
|
|
def build_reflection_graph(max_rounds: int = 2):
|
|
"""Строит граф с рефлексией"""
|
|
# Инициализация графа с состоянием
|
|
graph = StateGraph(ReflectState)
|
|
|
|
# Добавляем узлы
|
|
graph.add_node("draft_answer", draft_answer)
|
|
graph.add_node("reflect", reflect)
|
|
graph.add_node("rewrite", rewrite)
|
|
|
|
# Добавляем рёбра
|
|
graph.add_edge(START, "draft_answer")
|
|
graph.add_edge("draft_answer", "reflect")
|
|
|
|
# Условное ребро после рефлексии
|
|
graph.add_conditional_edges(
|
|
"reflect",
|
|
should_continue,
|
|
{
|
|
"rewrite": "rewrite",
|
|
"end": END
|
|
}
|
|
)
|
|
|
|
# После rewrite возвращаемся на reflect для повторной оценки
|
|
graph.add_edge("rewrite", "reflect")
|
|
|
|
# Компилируем граф
|
|
return graph.compile()
|
|
|
|
# Пример использования
|
|
if __name__ == "__main__":
|
|
# Создаём граф
|
|
app = build_reflection_graph(max_rounds=2)
|
|
|
|
# Входное состояние
|
|
initial_state = {
|
|
"question": "Объясни студенту разницу между tool и resource в MCP",
|
|
"draft": "",
|
|
"critique": "",
|
|
"verdict": "needs_revision",
|
|
"round": 1,
|
|
"max_rounds": 2
|
|
}
|
|
|
|
# Запуск
|
|
final_state = app.invoke(initial_state)
|
|
|
|
# Вывод результата
|
|
print("\n" + "="*80)
|
|
print("ФИНАЛЬНЫЙ ОТВЕТ:")
|
|
print("="*80)
|
|
print(final_state["draft"]) |