79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
from langgraph.graph import StateGraph, START, END
|
||
from state import ReflectState
|
||
from nodes import generate_answer
|
||
|
||
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['round'] + 1}/{state['max_rounds']})...")
|
||
return "retry"
|
||
|
||
# Превышен лимит - завершаем с ошибкой
|
||
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, "generate_answer")
|
||
|
||
# Условное ребро после генерации
|
||
graph.add_conditional_edges(
|
||
"generate_answer",
|
||
should_retry,
|
||
{
|
||
"retry": "generate_answer", # Петля для повторной попытки
|
||
"end": END
|
||
}
|
||
)
|
||
|
||
return graph.compile()
|
||
|
||
# Пример использования
|
||
if __name__ == "__main__":
|
||
# Создаём граф
|
||
app = build_retry_graph(max_rounds=3)
|
||
|
||
# Входное состояние
|
||
initial_state = {
|
||
"question": "Объясни студенту разницу между tool и resource в MCP",
|
||
"draft": "",
|
||
"error": None,
|
||
"round": 1,
|
||
"max_rounds": 3
|
||
}
|
||
|
||
# Запуск
|
||
final_state = app.invoke(initial_state)
|
||
|
||
# Вывод результата
|
||
print("\n" + "="*80)
|
||
if final_state["error"] is None:
|
||
print("✅ ФИНАЛЬНЫЙ ОТВЕТ:")
|
||
print("="*80)
|
||
print(final_state["draft"])
|
||
else:
|
||
print("❌ НЕ УДАЛОСЬ СГЕНЕРИРОВАТЬ КАЧЕСТВЕННЫЙ ОТВЕТ")
|
||
print(f"Ошибка: {final_state['error']}") |