diff --git a/graph.py b/graph.py index 51c4839..02df499 100644 --- a/graph.py +++ b/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']}") \ No newline at end of file + print("ФИНАЛЬНЫЙ ОТВЕТ:") + print("="*80) + print(final_state["draft"]) \ No newline at end of file