add

graph.py

@@ -1,66 +1,68 @@
 from langgraph.graph import StateGraph, START, END
 from state import ReflectState
-from nodes import draft_answer, reflect, rewrite
+from nodes import generate_answer
 
-def should_continue(state: ReflectState) -> str:
-    """Функция-условие: продолжать или завершить"""
-    # Проверяем вердикт
-    if state["verdict"] == "ok":
-        print(f"✅ Ответ принят после {state['round']} раундов")
+def should_retry(state: ReflectState) -> str:
+    """
+    Условие: продолжать попытки или завершить
+    Проверяет наличие ошибки и количество попыток
+    """
+    # Если ошибки нет - успех
+    if state["error"] is None:
+        print(f"✅ Ответ успешно сгенерирован за {state['round']} попыт(ок/ки)")
         return "end"
     
-    # Проверяем лимит раундов
-    if state["round"] >= state["max_rounds"]:
-        print(f"⚠️ Достигнут лимит раундов ({state['max_rounds']}). Возвращаем последнюю версию.")
+    # Если есть ошибка и не превышен лимит - повторяем
+    if state["round"] < state["max_rounds"]:
+        print(f"🔄 Повторная попытка ({state['round'] + 1}/{state['max_rounds']})...")
+        return "retry"
+    
+    # Превышен лимит - завершаем с ошибкой
+    print(f"⚠️ Достигнут лимит попыток ({state['max_rounds']}). Последняя ошибка: {state['error']}")
     return "end"
 
-    # Требуется доработка
-    print(f"🔄 Запускаем раунд доработки {state['round'] + 1}/{state['max_rounds']}")
-    return "rewrite"
+def build_retry_graph(max_rounds: int = 3):
+    """
+    Строит граф с retry-логикой через try/except
     
-def build_reflection_graph(max_rounds: int = 2):
-    """Строит граф с рефлексией"""
-    # Инициализация графа с состоянием
+    Структура:
+    START → generate_answer → should_retry
+        → (если error None) → END
+        → (если error и round < max_rounds) → generate_answer (с увеличенным round)
+        → иначе → END
+    """
     graph = StateGraph(ReflectState)
     
-    # Добавляем узлы
-    graph.add_node("draft_answer", draft_answer)
-    graph.add_node("reflect", reflect)
-    graph.add_node("rewrite", rewrite)
+    # Добавляем единственный узел (он же будет вызываться повторно)
+    graph.add_node("generate_answer", generate_answer)
     
     # Добавляем рёбра
-    graph.add_edge(START, "draft_answer")
-    graph.add_edge("draft_answer", "reflect")
+    graph.add_edge(START, "generate_answer")
     
-    # Условное ребро после рефлексии
+    # Условное ребро после генерации
     graph.add_conditional_edges(
-        "reflect",
-        should_continue,
+        "generate_answer",
+        should_retry,
         {
-            "rewrite": "rewrite",
+            "retry": "generate_answer",  # Петля для повторной попытки
             "end": END
         }
     )
     
-    # После rewrite возвращаемся на reflect для повторной оценки
-    graph.add_edge("rewrite", "reflect")
-    
-    # Компилируем граф
     return graph.compile()
 
 # Пример использования
 if __name__ == "__main__":
     # Создаём граф
-    app = build_reflection_graph(max_rounds=2)
+    app = build_retry_graph(max_rounds=3)
     
     # Входное состояние
     initial_state = {
         "question": "Объясни студенту разницу между tool и resource в MCP",
         "draft": "",
-        "critique": "",
-        "verdict": "needs_revision",
+        "error": None,
         "round": 1,
-        "max_rounds": 2
+        "max_rounds": 3
     }
     
     # Запуск
@@ -68,6 +70,10 @@ if __name__ == "__main__":
     
     # Вывод результата
     print("\n" + "="*80)
-    print("ФИНАЛЬНЫЙ ОТВЕТ:")
+    if final_state["error"] is None:
+        print("✅ ФИНАЛЬНЫЙ ОТВЕТ:")
         print("="*80)
         print(final_state["draft"])
+    else:
+        print("❌ НЕ УДАЛОСЬ СГЕНЕРИРОВАТЬ КАЧЕСТВЕННЫЙ ОТВЕТ")
+        print(f"Ошибка: {final_state['error']}")