66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
from langgraph.graph import StateGraph, START, END
|
|
from state import BriefState
|
|
from nodes import outline_node, research_step_node, synthesize_node
|
|
|
|
def should_continue_research(state: BriefState) -> str:
|
|
"""Условие: продолжать исследование или переходить к синтезу"""
|
|
if state["step_index"] < len(state["outline"]):
|
|
return "continue_research"
|
|
else:
|
|
return "synthesize"
|
|
|
|
def build_research_graph():
|
|
"""Строит граф исследовательского агента"""
|
|
# Инициализация графа с состоянием
|
|
graph = StateGraph(BriefState)
|
|
|
|
# Добавляем узлы
|
|
graph.add_node("outline", outline_node)
|
|
graph.add_node("research_step", research_step_node)
|
|
graph.add_node("synthesize", synthesize_node)
|
|
|
|
# Добавляем рёбра
|
|
graph.add_edge(START, "outline")
|
|
graph.add_edge("outline", "research_step")
|
|
|
|
# Условное ребро после research_step
|
|
graph.add_conditional_edges(
|
|
"research_step",
|
|
should_continue_research,
|
|
{
|
|
"continue_research": "research_step", # Цикл по шагам
|
|
"synthesize": "synthesize"
|
|
}
|
|
)
|
|
|
|
# После синтеза завершаем
|
|
graph.add_edge("synthesize", END)
|
|
|
|
# Компилируем граф
|
|
return graph.compile()
|
|
|
|
# Пример использования (для тестирования)
|
|
if __name__ == "__main__":
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
# Создаём граф
|
|
app = build_research_graph()
|
|
|
|
# Входное состояние
|
|
initial_state: BriefState = {
|
|
"topic": "Как студенту безопасно подключать MCP к LangChain",
|
|
"outline": None,
|
|
"step_index": 0,
|
|
"notes": [],
|
|
"final_brief": None
|
|
}
|
|
|
|
# Запуск
|
|
final_state = app.invoke(initial_state)
|
|
|
|
# Вывод результата
|
|
print("\n" + "="*80)
|
|
print("ИТОГОВЫЙ ИССЛЕДОВАТЕЛЬСКИЙ БРИФ")
|
|
print("="*80)
|
|
print(final_state["final_brief"]) |