Повторный экзамен: Исследовательский бриф (план → шаги → сводка): solution.py

solution.py

@@ -0,0 +1,106 @@
+import os
+from typing import TypedDict, List, Optional, Dict, Any
+
+from langchain_openai import ChatOpenAI
+from langchain_tavily import TavilySearchResults
+from langgraph.graph import StateGraph, START, END
+from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import create_react_agent
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# ---------- 1. Состояние ----------
+class BriefState(TypedDict):
+    topic: str
+    outline: List[str] | None
+    step_index: int
+    notes: List[str]
+    final_brief: str | None
+
+# ---------- 2. Инструменты ----------
+search_tool = TavilySearchResults()
+
+# ---------- 3. Узлы ----------
+def outline_node(state: BriefState) -> Dict[str, Any]:
+    """Создаёт план из 4–5 пунктов по теме."""
+    llm = ChatOpenAI(model="gpt-4o-mini")
+    prompt = (
+        f"Составь план исследования по теме: {state['topic']}. "
+        "План должен содержать 4–5 пунктов, каждый пункт – заголовок без описания."
+    )
+    plan = llm.invoke(prompt).content.strip()
+    outline = [line.strip("- ").strip() for line in plan.splitlines() if line.strip()]
+    return {"outline": outline, "step_index": 0, "notes": []}
+
+def research_step_node(state: BriefState) -> Dict[str, Any]:
+    """Для текущего пункта плана делает поиск и формирует заметку."""
+    llm = ChatOpenAI(model="gpt-4o-mini")
+    current_index = state["step_index"]
+    if state["outline"] is None or current_index >= len(state["outline"]):
+        return {"step_index": current_index, "notes": state["notes"]}
+
+    topic = state["outline"][current_index]
+    # Поиск
+    search_query = f"{state['topic']} - {topic}"
+    results = search_tool.invoke({"query": search_query})
+    # Сводка
+    summary_prompt = (
+        f"Сделай краткую заметку (5–8 предложений) по теме '{topic}'. "
+        f"Используй найденные источники: {results}"
+    )
+    note = llm.invoke(summary_prompt).content.strip()
+    notes = state["notes"] + [f"**{topic}**\n{note}"]
+    return {"step_index": current_index + 1, "notes": notes}
+
+def synthesize_node(state: BriefState) -> Dict[str, Any]:
+    """Объединяет все заметки в связный бриф."""
+    llm = ChatOpenAI(model="gpt-4o-mini")
+    combined = "\n\n".join(state["notes"])
+    prompt = (
+        f"Собери из следующих заметок связный исследовательский бриф (½–1 страница):\n{combined}"
+    )
+    brief = llm.invoke(prompt).content.strip()
+    return {"final_brief": brief}
+
+# ---------- 4. Граф ----------
+def build_graph() -> StateGraph[BriefState]:
+    graph = StateGraph(BriefState)
+
+    # Добавляем узлы
+    graph.add_node("outline", outline_node)
+    graph.add_node("research_step", research_step_node)
+    graph.add_node("synthesize", synthesize_node)
+
+    # Переходы
+    graph.set_entry_point("outline")
+    graph.add_edge("outline", "research_step")
+    graph.add_conditional_edges(
+        "research_step",
+        lambda state: "synthesize" if state["step_index"] >= len(state["outline"] or []) else "research_step",
+        {
+            "synthesize": "synthesize",
+            "research_step": "research_step",
+        },
+    )
+    graph.add_edge("synthesize", END)
+
+    return graph
+
+# ---------- 5. Запуск ----------
+if __name__ == "__main__":
+    topic = input("Введите тему исследования: ").strip()
+    initial_state: BriefState = {
+        "topic": topic,
+        "outline": None,
+        "step_index": 0,
+        "notes": [],
+        "final_brief": None,
+    }
+
+    graph = build_graph()
+    chain = graph.compile()
+
+    result = chain.invoke(initial_state)
+    print("\n=== Итоговый исследовательский бриф ===\n")
+    print(result["final_brief"])