fix: main.py — Повторный экзамен: Исследовательский бриф (план → шаги → сводка)
This commit is contained in:
@@ -1,22 +1,19 @@
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
from dotenv import load_dotenv
|
||||
from typing import TypedDict, List, Optional
|
||||
from typing import TypedDict, List, Annotated
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_tavily import TavilySearchResults
|
||||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||||
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
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# LLM configuration
|
||||
# ----------------------------------------------------------------------
|
||||
# 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"),
|
||||
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
|
||||
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
|
||||
# ----------------------------------------------------------------------
|
||||
# Nodes
|
||||
# ----------------------------------------------------------------------
|
||||
def outline_node(state: BriefState) -> BriefState:
|
||||
if state["outline"] is None:
|
||||
"""Generate a 4-5 item outline for the given topic."""
|
||||
prompt = (
|
||||
f"Generate 4-5 bullet points outlining a research plan for the topic: "
|
||||
f"{state['topic']}. Return as a JSON array of strings."
|
||||
f"Create a concise outline (4-5 bullet points) for a short research brief on the topic:\n"
|
||||
f"\"{state['topic']}\"\n"
|
||||
"Each bullet should be a short phrase suitable as a section heading."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
try:
|
||||
outline = json.loads(response.content)
|
||||
if isinstance(outline, list):
|
||||
state["outline"] = outline
|
||||
except Exception:
|
||||
state["outline"] = []
|
||||
return state
|
||||
response = llm.invoke([HumanMessage(content=prompt)])
|
||||
outline_text = response.content.strip()
|
||||
# Split on newlines and strip bullet characters
|
||||
items = [line.lstrip("-• ").strip() for line in outline_text.splitlines() if line.strip()]
|
||||
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 []):
|
||||
current_step = state["outline"][state["step_index"]]
|
||||
# Perform web search
|
||||
search_results = tavily_search(current_step)
|
||||
# Summarize results
|
||||
summary_prompt = (
|
||||
f"Summarize the following search results into 5-8 sentences:\n{search_results}"
|
||||
)
|
||||
summary = llm.invoke(summary_prompt).content.strip()
|
||||
state["notes"].append(summary)
|
||||
state["step_index"] += 1
|
||||
return state
|
||||
"""Research one outline item, produce a short note, and store it."""
|
||||
outline = state["outline"]
|
||||
idx = state["step_index"]
|
||||
if outline is None or idx >= len(outline):
|
||||
return state # safety
|
||||
|
||||
current_topic = outline[idx]
|
||||
# Use web search tool
|
||||
search_result = web_search(current_topic)
|
||||
|
||||
# Prompt LLM to write a 5-8 sentence note based on search result
|
||||
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 []):
|
||||
notes_text = "\n\n".join(state["notes"])
|
||||
synth_prompt = (
|
||||
f"Combine the following notes into a concise research brief (about one page). "
|
||||
f"Use headings for each point.\n\n{notes_text}"
|
||||
"""Combine all notes into a coherent brief."""
|
||||
if not state["notes"]:
|
||||
final = "No notes were collected."
|
||||
else:
|
||||
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()
|
||||
state["final_brief"] = brief
|
||||
return state
|
||||
synthesis_response = llm.invoke([HumanMessage(content=synthesis_prompt)])
|
||||
final = synthesis_response.content.strip()
|
||||
return {
|
||||
**state,
|
||||
"final_brief": final,
|
||||
}
|
||||
|
||||
# Build the LangGraph
|
||||
graph = StateGraph(BriefState)
|
||||
graph.add_node("outline", outline_node)
|
||||
graph.add_node("research_step", research_step_node)
|
||||
graph.add_node("synthesize", synthesize_node)
|
||||
# ----------------------------------------------------------------------
|
||||
# Graph construction
|
||||
# ----------------------------------------------------------------------
|
||||
workflow = StateGraph(BriefState)
|
||||
|
||||
graph.set_entry_point("outline")
|
||||
graph.add_conditional_edges(
|
||||
workflow.add_node("outline", outline_node)
|
||||
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",
|
||||
lambda state: (
|
||||
"research_step"
|
||||
if state["step_index"] < len(state["outline"] or [])
|
||||
else "synthesize"
|
||||
),
|
||||
|
||||
def continue_condition(state: BriefState) -> str:
|
||||
if state["outline"] is None:
|
||||
return END
|
||||
if state["step_index"] < len(state["outline"]):
|
||||
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()
|
||||
|
||||
# Tool: Generate brief using the graph
|
||||
@tool
|
||||
def generate_brief(topic: str) -> str:
|
||||
"""Generate a research brief for the given topic."""
|
||||
initial_state: BriefState = {
|
||||
# ----------------------------------------------------------------------
|
||||
# Demo execution
|
||||
# ----------------------------------------------------------------------
|
||||
async def main():
|
||||
topic = "Как студенту безопасно подключать MCP к LangChain"
|
||||
# Initialize state
|
||||
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
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[generate_brief, tavily_search],
|
||||
backend=backend,
|
||||
system_prompt="You are a helpful research assistant. Use the provided tools to generate a research brief.",
|
||||
# Run the graph
|
||||
result = await graph.ainvoke(
|
||||
init_state,
|
||||
config={"configurable": {"thread_id": "demo-1"}},
|
||||
)
|
||||
|
||||
# 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())
|
||||
Reference in New Issue
Block a user