add graph.py

2026-06-04 16:26:38 +00:00
parent a8dea48a5d
commit b6acca7a83
1 changed files with 73 additions and 0 deletions
@@ -0,0 +1,73 @@
+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"])