diff --git a/graph.py b/graph.py new file mode 100644 index 0000000..5eb4cfa --- /dev/null +++ b/graph.py @@ -0,0 +1,66 @@ +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"]) \ No newline at end of file