add graph.py

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"])