add

graph.py

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