148 lines
4.8 KiB
Python
148 lines
4.8 KiB
Python
import os
|
||
import asyncio
|
||
from typing import TypedDict, Annotated, List
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage
|
||
from langchain.tools import tool
|
||
from deepagents import create_deep_agent
|
||
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.graph.message import add_messages
|
||
from langchain_tavily import TavilySearchResults
|
||
|
||
# Load environment variables
|
||
from dotenv import load_dotenv
|
||
load_dotenv()
|
||
|
||
# LLM configuration – OpenRouter
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://openrouter.ai/api/v1",
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# Backend for deepagents – local shell + filesystem
|
||
backend = CompositeBackend([
|
||
LocalShellBackend(workspace_dir="./workspace"),
|
||
FilesystemBackend(),
|
||
])
|
||
|
||
# Tavily search tool – real web search
|
||
@tool
|
||
def tavily_search(query: str) -> str:
|
||
"""Search the web using Tavily and return a short summary."""
|
||
tavily = TavilySearchResults(max_results=3)
|
||
results = tavily.run(query)
|
||
# Concatenate titles and snippets
|
||
return "\n".join(f"{r['title']}: {r['content']}" for r in results)
|
||
|
||
# State definition
|
||
class BriefState(TypedDict):
|
||
topic: str
|
||
outline: List[str] | None
|
||
step_index: int
|
||
notes: List[str]
|
||
final_brief: str | None
|
||
|
||
# Node: generate outline
|
||
async def outline_node(state: BriefState) -> BriefState:
|
||
prompt = (
|
||
f"Generate a concise outline of 4–5 research steps for the topic: {state['topic']}\n"
|
||
"Return a JSON array of strings, each a single step."
|
||
)
|
||
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||
# Extract JSON array
|
||
import json, re
|
||
try:
|
||
array_text = re.search(r"\[.*\]", response.content, re.S).group(0)
|
||
outline = json.loads(array_text)
|
||
except Exception:
|
||
outline = ["Step 1: ...", "Step 2: ...", "Step 3: ...", "Step 4: ..."]
|
||
state["outline"] = outline
|
||
state["step_index"] = 0
|
||
state["notes"] = []
|
||
return state
|
||
|
||
# Node: research one step
|
||
async def research_step_node(state: BriefState) -> BriefState:
|
||
step = state["outline"][state["step_index"]]
|
||
# Use Tavily to gather info
|
||
search_query = f"{state['topic']} {step}"
|
||
search_result = tavily_search(search_query)
|
||
# Summarize with LLM
|
||
prompt = (
|
||
f"Using the following search results, write a concise note (5–8 sentences) for the step: {step}\n"
|
||
f"Search results:\n{search_result}\n"
|
||
"Note: keep it factual and cite sources if possible."
|
||
)
|
||
note = await llm.ainvoke([HumanMessage(content=prompt)])
|
||
state["notes"].append(note.content.strip())
|
||
state["step_index"] += 1
|
||
return state
|
||
|
||
# Node: synthesize final brief
|
||
async def synthesize_node(state: BriefState) -> BriefState:
|
||
notes = state["notes"]
|
||
outline = state["outline"]
|
||
prompt = (
|
||
"You are an academic writer. Using the following outline and notes, produce a cohesive research brief of ½–1 page.\n"
|
||
f"Outline: {outline}\n"
|
||
f"Notes: {notes}\n"
|
||
"Structure the brief with headings matching the outline steps."
|
||
)
|
||
brief = await llm.ainvoke([HumanMessage(content=prompt)])
|
||
state["final_brief"] = brief.content.strip()
|
||
return state
|
||
|
||
# Build LangGraph
|
||
graph = StateGraph(BriefState)
|
||
graph.add_node("outline", outline_node)
|
||
graph.add_node("research_step", research_step_node)
|
||
graph.add_node("synthesize", synthesize_node)
|
||
|
||
# Entry point
|
||
graph.set_entry_point("outline")
|
||
# After outline, loop research_step until all steps processed
|
||
graph.add_conditional_edges(
|
||
"outline",
|
||
lambda _: "research_step",
|
||
)
|
||
# After each research_step, decide whether to continue or synthesize
|
||
graph.add_conditional_edges(
|
||
"research_step",
|
||
lambda state: "synthesize" if state["step_index"] >= len(state["outline"]) else "research_step",
|
||
)
|
||
# Final node
|
||
graph.add_edge("synthesize", END)
|
||
|
||
app = graph.compile()
|
||
|
||
# DeepAgent wrapper
|
||
agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[tavily_search],
|
||
backend=backend,
|
||
system_prompt="You are a research assistant that builds a brief.",
|
||
)
|
||
|
||
async def main():
|
||
# Default topic
|
||
topic = os.getenv("DEFAULT_TOPIC", "Как студенту безопасно подключать MCP к LangChain")
|
||
# Run LangGraph to get outline and notes
|
||
state = await app.ainvoke({"topic": topic, "outline": None, "step_index": 0, "notes": [], "final_brief": None})
|
||
# Print outline
|
||
print("\n=== Outline ===")
|
||
for i, step in enumerate(state["outline"], 1):
|
||
print(f"{i}. {step}")
|
||
# Print notes per step
|
||
print("\n=== Notes ===")
|
||
for i, note in enumerate(state["notes"], 1):
|
||
print(f"[Step {i}] {note}\n")
|
||
# Print final brief
|
||
print("\n=== Final Brief ===")
|
||
print(state["final_brief"])
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|