add graph.py

graph.py

@@ -1,59 +1,68 @@
 from langgraph.graph import StateGraph, START, END
-from state import BriefState
-from nodes import outline_node, research_step_node, synthesize_node
+from state import ReflectState
+from nodes import generate_answer
 
-def should_continue_research(state: BriefState) -> str:
-    """Условие: продолжать исследование или переходить к синтезу"""
-    if state["step_index"] < len(state["outline"]):
-        return "continue_research"
-    else:
-        return "synthesize"
+def should_retry(state: ReflectState) -> str:
+    """
+    Условие: продолжать попытки или завершить
+    Проверяет наличие ошибки и количество попыток
+    """
+    # Если ошибки нет - успех
+    if state["error"] is None:
+        print(f"✅ Ответ успешно сгенерирован за {state['round']} попыт(ок/ки)")
+        return "end"
     
-def build_research_graph():
-    """Строит граф исследовательского агента"""
-    # Инициализация графа с состоянием
-    graph = StateGraph(BriefState)
+    # Если есть ошибка и не превышен лимит - повторяем
+    if state["round"] < state["max_rounds"]:
+        print(f"🔄 Повторная попытка ({state['round'] + 1}/{state['max_rounds']})...")
+        return "retry"
     
-    # Добавляем узлы
-    graph.add_node("outline", outline_node)
-    graph.add_node("research_step", research_step_node)
-    graph.add_node("synthesize", synthesize_node)
+    # Превышен лимит - завершаем с ошибкой
+    print(f"⚠️ Достигнут лимит попыток ({state['max_rounds']}). Последняя ошибка: {state['error']}")
+    return "end"
+
+def build_retry_graph(max_rounds: int = 3):
+    """
+    Строит граф с retry-логикой через try/except
+    
+    Структура:
+    START → generate_answer → should_retry
+        → (если error None) → END
+        → (если error и round < max_rounds) → generate_answer (с увеличенным round)
+        → иначе → END
+    """
+    graph = StateGraph(ReflectState)
+    
+    # Добавляем единственный узел (он же будет вызываться повторно)
+    graph.add_node("generate_answer", generate_answer)
     
     # Добавляем рёбра
-    graph.add_edge(START, "outline")
-    graph.add_edge("outline", "research_step")
+    graph.add_edge(START, "generate_answer")
     
-    # Условное ребро после research_step
+    # Условное ребро после генерации
     graph.add_conditional_edges(
-        "research_step",
-        should_continue_research,
+        "generate_answer",
+        should_retry,
         {
-            "continue_research": "research_step",  # Цикл по шагам
-            "synthesize": "synthesize"
+            "retry": "generate_answer",  # Петля для повторной попытки
+            "end": END
         }
     )
     
-    # После синтеза завершаем
-    graph.add_edge("synthesize", END)
-    
-    # Компилируем граф
     return graph.compile()
 
-# Пример использования (для тестирования)
+# Пример использования
 if __name__ == "__main__":
-    from dotenv import load_dotenv
-    load_dotenv()
-    
     # Создаём граф
-    app = build_research_graph()
+    app = build_retry_graph(max_rounds=3)
     
     # Входное состояние
-    initial_state: BriefState = {
-        "topic": "Как студенту безопасно подключать MCP к LangChain",
-        "outline": None,
-        "step_index": 0,
-        "notes": [],
-        "final_brief": None
+    initial_state = {
+        "question": "Объясни студенту разницу между tool и resource в MCP",
+        "draft": "",
+        "error": None,
+        "round": 1,
+        "max_rounds": 3
     }
     
     # Запуск
@@ -61,6 +70,10 @@ if __name__ == "__main__":
     
     # Вывод результата
     print("\n" + "="*80)
-    print("ИТОГОВЫЙ ИССЛЕДОВАТЕЛЬСКИЙ БРИФ")
+    if final_state["error"] is None:
+        print("✅ ФИНАЛЬНЫЙ ОТВЕТ:")
         print("="*80)
-    print(final_state["final_brief"])
+        print(final_state["draft"])
+    else:
+        print("❌ НЕ УДАЛОСЬ СГЕНЕРИРОВАТЬ КАЧЕСТВЕННЫЙ ОТВЕТ")
+        print(f"Ошибка: {final_state['error']}")