101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
from langgraph.graph import StateGraph, START, END
|
|
from state import ReflectState
|
|
from nodes import draft_answer, reflect, rewrite
|
|
|
|
def should_continue_or_retry(state: ReflectState) -> str:
|
|
"""
|
|
Расширенная логика: сначала проверяем try/except ошибки,
|
|
затем вердикт критика
|
|
"""
|
|
# 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"
|
|
|
|
# 3. Проверяем лимит раундов доработки
|
|
if state["round"] >= state["max_rounds"]:
|
|
print(f"⚠️ Достигнут лимит раундов ({state['max_rounds']}). Возвращаем последнюю версию.")
|
|
return "end"
|
|
|
|
# 4. Требуется доработка
|
|
print(f"🔄 Запускаем раунд доработки {state['round'] + 1}/{state['max_rounds']}")
|
|
return "rewrite"
|
|
|
|
def build_reflection_graph(max_rounds: int = 2, max_retries: int = 3):
|
|
"""
|
|
Строит граф с рефлексией и retry при ошибках
|
|
|
|
Структура:
|
|
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("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_or_retry,
|
|
{
|
|
"retry_draft": "draft_answer", # retry при ошибке генерации
|
|
"rewrite": "rewrite", # доработка по вердикту критика
|
|
"end": END
|
|
}
|
|
)
|
|
|
|
# После rewrite возвращаемся на reflect для повторной оценки
|
|
graph.add_edge("rewrite", "reflect")
|
|
|
|
# Компилируем граф
|
|
return graph.compile()
|
|
|
|
# Пример использования
|
|
if __name__ == "__main__":
|
|
# Создаём граф
|
|
app = build_reflection_graph(max_rounds=2, max_retries=3)
|
|
|
|
# Входное состояние
|
|
initial_state: ReflectState = {
|
|
"question": "Объясни студенту разницу между tool и resource в MCP",
|
|
"draft": "",
|
|
"critique": "",
|
|
"verdict": "needs_revision",
|
|
"round": 1,
|
|
"max_rounds": 2,
|
|
"error": None
|
|
}
|
|
|
|
# Запуск
|
|
final_state = app.invoke(initial_state)
|
|
|
|
# Вывод результата
|
|
print("\n" + "="*80)
|
|
print("ФИНАЛЬНЫЙ ОТВЕТ:")
|
|
print("="*80)
|
|
print(final_state["draft"]) |