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

main.py

@@ -1,22 +1,19 @@
 import os
-import json
 import asyncio
-from dotenv import load_dotenv
+from typing import TypedDict, List, Annotated
-from typing import TypedDict, List, Optional
 
-from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langchain_tavily import TavilySearchResults
+from langchain_core.messages import HumanMessage
 from langchain.tools import tool
-
+from langchain_tavily import TavilySearchResults
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-from langgraph.graph import StateGraph, START, END
+# ----------------------------------------------------------------------
-
+# Configuration
-# Load environment variables
+# ----------------------------------------------------------------------
-load_dotenv()
-
-# LLM configuration
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -24,13 +21,32 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# Backend for deepagents
+backend = CompositeBackend(
-backend = CompositeBackend([
+    [
-    LocalShellBackend(workspace_dir="./workspace"),
+        LocalShellBackend(workspace_dir="./workspace"),
-    FilesystemBackend(),
+        FilesystemBackend(),
-])
+    ]
+)
 
-# TypedDict for graph state
+# ----------------------------------------------------------------------
+# Tools
+# ----------------------------------------------------------------------
+search_tool = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
+
+@tool
+def web_search(query: str) -> str:
+    """
+    Perform a web search using Tavily and return a concise summary of the top results.
+    The LLM will ask for a short note based on this summary.
+    """
+    results = search_tool.run(query)
+    # results is a list of dicts with 'url' and 'content'
+    summaries = [r.get("content", "") for r in results[:3]]
+    return "\n".join(summaries) if summaries else "No relevant results found."
+
+# ----------------------------------------------------------------------
+# State definition
+# ----------------------------------------------------------------------
 class BriefState(TypedDict):
     topic: str
     outline: List[str] | None
@@ -38,110 +54,145 @@ class BriefState(TypedDict):
     notes: List[str]
     final_brief: str | None
 
-# Tool: Web search using Tavily
+# ----------------------------------------------------------------------
-@tool
+# Nodes
-def tavily_search(query: str) -> str:
+# ----------------------------------------------------------------------
-    """Search the web for the query and return results."""
-    search = TavilySearchResults()
-    results = search.run(query)
-    return results
-
-# Outline node: generate research plan
 def outline_node(state: BriefState) -> BriefState:
-    if state["outline"] is None:
+    """Generate a 4-5 item outline for the given topic."""
-        prompt = (
+    prompt = (
-            f"Generate 4-5 bullet points outlining a research plan for the topic: "
+        f"Create a concise outline (4-5 bullet points) for a short research brief on the topic:\n"
-            f"{state['topic']}. Return as a JSON array of strings."
+        f"\"{state['topic']}\"\n"
-        )
+        "Each bullet should be a short phrase suitable as a section heading."
-        response = llm.invoke(prompt)
+    )
-        try:
+    response = llm.invoke([HumanMessage(content=prompt)])
-            outline = json.loads(response.content)
+    outline_text = response.content.strip()
-            if isinstance(outline, list):
+    # Split on newlines and strip bullet characters
-                state["outline"] = outline
+    items = [line.lstrip("-• ").strip() for line in outline_text.splitlines() if line.strip()]
-        except Exception:
+    return {
-            state["outline"] = []
+        **state,
-    return state
+        "outline": items,
+        "step_index": 0,
+        "notes": [],
+    }
 
-# Research step node: one search per iteration
 def research_step_node(state: BriefState) -> BriefState:
-    if state["step_index"] < len(state["outline"] or []):
+    """Research one outline item, produce a short note, and store it."""
-        current_step = state["outline"][state["step_index"]]
+    outline = state["outline"]
-        # Perform web search
+    idx = state["step_index"]
-        search_results = tavily_search(current_step)
+    if outline is None or idx >= len(outline):
-        # Summarize results
+        return state  # safety
-        summary_prompt = (
+
-            f"Summarize the following search results into 5-8 sentences:\n{search_results}"
+    current_topic = outline[idx]
-        )
+    # Use web search tool
-        summary = llm.invoke(summary_prompt).content.strip()
+    search_result = web_search(current_topic)
-        state["notes"].append(summary)
+
-        state["step_index"] += 1
+    # Prompt LLM to write a 5-8 sentence note based on search result
-    return state
+    note_prompt = (
+        f"Based on the following web search summary, write a short note (5-8 sentences) "
+        f"that could serve as a paragraph for a research brief about \"{state['topic']}\". "
+        f"Focus on the aspect: \"{current_topic}\".\n\n"
+        f"Search summary:\n{search_result}"
+    )
+    note_response = llm.invoke([HumanMessage(content=note_prompt)])
+    note = note_response.content.strip()
+
+    new_notes = state["notes"] + [f"## {current_topic}\n{note}"]
+    return {
+        **state,
+        "notes": new_notes,
+        "step_index": idx + 1,
+    }
 
-# Synthesize node: combine notes into final brief
 def synthesize_node(state: BriefState) -> BriefState:
-    if state["step_index"] >= len(state["outline"] or []):
+    """Combine all notes into a coherent brief."""
-        notes_text = "\n\n".join(state["notes"])
+    if not state["notes"]:
-        synth_prompt = (
+        final = "No notes were collected."
-            f"Combine the following notes into a concise research brief (about one page). "
+    else:
-            f"Use headings for each point.\n\n{notes_text}"
+        combined = "\n\n".join(state["notes"])
+        synthesis_prompt = (
+            f"Combine the following sections into a single cohesive research brief (about half to one page). "
+            f"Keep the headings, ensure logical flow, and add a brief introduction and conclusion.\n\n"
+            f"{combined}"
         )
-        brief = llm.invoke(synth_prompt).content.strip()
+        synthesis_response = llm.invoke([HumanMessage(content=synthesis_prompt)])
-        state["final_brief"] = brief
+        final = synthesis_response.content.strip()
-    return state
+    return {
+        **state,
+        "final_brief": final,
+    }
 
-# Build the LangGraph
+# ----------------------------------------------------------------------
-graph = StateGraph(BriefState)
+# Graph construction
-graph.add_node("outline", outline_node)
+# ----------------------------------------------------------------------
-graph.add_node("research_step", research_step_node)
+workflow = StateGraph(BriefState)
-graph.add_node("synthesize", synthesize_node)
 
-graph.set_entry_point("outline")
+workflow.add_node("outline", outline_node)
-graph.add_conditional_edges(
+workflow.add_node("research_step", research_step_node)
+workflow.add_node("synthesize", synthesize_node)
+
+workflow.add_edge(START, "outline")
+workflow.add_conditional_edges(
     "outline",
-    lambda state: "research_step" if state["outline"] is not None else END,
+    lambda state: "research_step" if state["outline"] else END,
 )
-graph.add_conditional_edges(
+
-    "research_step",
+def continue_condition(state: BriefState) -> str:
-    lambda state: (
+    if state["outline"] is None:
-        "research_step"
+        return END
-        if state["step_index"] < len(state["outline"] or [])
+    if state["step_index"] < len(state["outline"]):
-        else "synthesize"
+        return "research_step"
-    ),
+    return "synthesize"
+
+workflow.add_edge("research_step", "research_step")
+workflow.add_conditional_edges("research_step", continue_condition)
+workflow.add_edge("synthesize", END)
+
+graph = workflow.compile()
+
+# ----------------------------------------------------------------------
+# DeepAgent wrapper (required by the course)
+# ----------------------------------------------------------------------
+agent = create_deep_agent(
+    model=llm,
+    tools=[web_search],
+    backend=backend,
+    system_prompt="You are an AI research assistant that helps build short research briefs.",
 )
-graph.add_edge("synthesize", END)
 
-compiled_graph = graph.compile()
+# ----------------------------------------------------------------------
-
+# Demo execution
-# Tool: Generate brief using the graph
+# ----------------------------------------------------------------------
-@tool
+async def main():
-def generate_brief(topic: str) -> str:
+    topic = "Как студенту безопасно подключать MCP к LangChain"
-    """Generate a research brief for the given topic."""
+    # Initialize state
-    initial_state: BriefState = {
+    init_state: BriefState = {
         "topic": topic,
         "outline": None,
         "step_index": 0,
         "notes": [],
         "final_brief": None,
     }
-    final_state = compiled_graph.invoke(initial_state)
-    return final_state["final_brief"] or ""
 
-# Create deepagents agent
+    # Run the graph
-agent = create_deep_agent(
+    result = await graph.ainvoke(
-    model=llm,
+        init_state,
-    tools=[generate_brief, tavily_search],
+        config={"configurable": {"thread_id": "demo-1"}},
-    backend=backend,
-    system_prompt="You are a helpful research assistant. Use the provided tools to generate a research brief.",
-)
-
-# Demo execution
-async def main():
-    topic = "Как студенту безопасно подключать MCP к LangChain"
-    result = await agent.ainvoke(
-        {"messages": [HumanMessage(content=f"Generate a brief on: {topic}")]},
-        {"configurable": {"thread_id": "session-1"}},
     )
-    print(result["messages"][-1].content)
+
+    # Print results
+    print("\n--- Outline ---")
+    if result["outline"]:
+        for i, item in enumerate(result["outline"], 1):
+            print(f"{i}. {item}")
+
+    print("\n--- Research Steps ---")
+    for i, note in enumerate(result["notes"], 1):
+        print(f"[Step {i}]")
+        print(note)
+        print()
+
+    print("\n--- Final Brief ---")
+    print(result["final_brief"] or "No brief generated.")
 
 if __name__ == "__main__":
     asyncio.run(main())