147 lines
4.4 KiB
Python
147 lines
4.4 KiB
Python
import os
|
|
import json
|
|
import asyncio
|
|
from dotenv import load_dotenv
|
|
from typing import TypedDict, List, Optional
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_tavily import TavilySearchResults
|
|
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
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# LLM configuration
|
|
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
|
|
backend = CompositeBackend([
|
|
LocalShellBackend(workspace_dir="./workspace"),
|
|
FilesystemBackend(),
|
|
])
|
|
|
|
# TypedDict for graph state
|
|
class BriefState(TypedDict):
|
|
topic: str
|
|
outline: List[str] | None
|
|
step_index: int
|
|
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
|
|
def outline_node(state: BriefState) -> BriefState:
|
|
if state["outline"] is None:
|
|
prompt = (
|
|
f"Generate 4-5 bullet points outlining a research plan for the topic: "
|
|
f"{state['topic']}. Return as a JSON array of strings."
|
|
)
|
|
response = llm.invoke(prompt)
|
|
try:
|
|
outline = json.loads(response.content)
|
|
if isinstance(outline, list):
|
|
state["outline"] = outline
|
|
except Exception:
|
|
state["outline"] = []
|
|
return state
|
|
|
|
# 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
|
|
|
|
# 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}"
|
|
)
|
|
brief = llm.invoke(synth_prompt).content.strip()
|
|
state["final_brief"] = brief
|
|
return state
|
|
|
|
# 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.set_entry_point("outline")
|
|
graph.add_conditional_edges(
|
|
"outline",
|
|
lambda state: "research_step" if state["outline"] is not None else END,
|
|
)
|
|
graph.add_conditional_edges(
|
|
"research_step",
|
|
lambda state: (
|
|
"research_step"
|
|
if state["step_index"] < len(state["outline"] or [])
|
|
else "synthesize"
|
|
),
|
|
)
|
|
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 = {
|
|
"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.",
|
|
)
|
|
|
|
# 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)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |