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

This commit is contained in:
2026-07-02 04:27:51 +00:00
parent 26e7a37f35
commit a84893b8ef
+148 -97
View File
@@ -1,22 +1,19 @@
import os import os
import json
import asyncio 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.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 import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langgraph.graph import StateGraph, START, END # ----------------------------------------------------------------------
# Configuration
# Load environment variables # ----------------------------------------------------------------------
load_dotenv()
# LLM configuration
llm = ChatOpenAI( llm = ChatOpenAI(
model="openai/gpt-oss-20b:free", model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1", base_url="https://openrouter.ai/api/v1",
@@ -24,13 +21,32 @@ llm = ChatOpenAI(
temperature=0.0, 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): class BriefState(TypedDict):
topic: str topic: str
outline: List[str] | None outline: List[str] | None
@@ -38,110 +54,145 @@ class BriefState(TypedDict):
notes: List[str] notes: List[str]
final_brief: str | None 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: 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) response = llm.invoke([HumanMessage(content=prompt)])
try: outline_text = response.content.strip()
outline = json.loads(response.content) # Split on newlines and strip bullet characters
if isinstance(outline, list): items = [line.lstrip("-• ").strip() for line in outline_text.splitlines() if line.strip()]
state["outline"] = outline return {
except Exception: **state,
state["outline"] = [] "outline": items,
return state "step_index": 0,
"notes": [],
}
# Research step node: one search per iteration
def research_step_node(state: BriefState) -> BriefState: 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: 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", "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, "topic": topic,
"outline": None, "outline": None,
"step_index": 0, "step_index": 0,
"notes": [], "notes": [],
"final_brief": None, "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 # Print results
async def main(): print("\n--- Outline ---")
topic = "Как студенту безопасно подключать MCP к LangChain" if result["outline"]:
result = await agent.ainvoke( for i, item in enumerate(result["outline"], 1):
{"messages": [HumanMessage(content=f"Generate a brief on: {topic}")]}, print(f"{i}. {item}")
{"configurable": {"thread_id": "session-1"}},
) print("\n--- Research Steps ---")
print(result["messages"][-1].content) 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__": if __name__ == "__main__":
asyncio.run(main()) asyncio.run(main())