add agent.py
This commit is contained in:
@@ -0,0 +1,120 @@
|
|||||||
|
"""
|
||||||
|
LangGraph agent that builds a research brief.
|
||||||
|
|
||||||
|
The graph follows the specification from the assignment:
|
||||||
|
* Outline node – generates 4‑5 bullet points for the topic.
|
||||||
|
* Research step node – for each outline item performs one web search via Tavily and creates a short note.
|
||||||
|
* Synthesize node – combines all notes into a coherent brief.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from typing import TypedDict, List, Optional
|
||||||
|
|
||||||
|
from langgraph.graph import StateGraph, START, END
|
||||||
|
from langchain_openai import ChatOpenAI
|
||||||
|
from langchain_tavily.tools import TavilySearchResults
|
||||||
|
from langchain_core.messages import HumanMessage
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
# ---------- State definition --------------------------------------------
|
||||||
|
class BriefState(TypedDict):
|
||||||
|
topic: str
|
||||||
|
outline: List[str] | None
|
||||||
|
step_index: int
|
||||||
|
notes: List[str]
|
||||||
|
final_brief: Optional[str]
|
||||||
|
|
||||||
|
# ---------- LLM and tools ---------------------------------------------
|
||||||
|
llm = ChatOpenAI(
|
||||||
|
model="openai/gpt-oss-20b:free",
|
||||||
|
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
||||||
|
api_key=os.getenv("JOURNAL_MCP_PAT"),
|
||||||
|
temperature=0.5,
|
||||||
|
)
|
||||||
|
|
||||||
|
search_tool = TavilySearchResults(max_results=3, tavily_api_key=os.getenv("TAVILY_API_KEY"))
|
||||||
|
|
||||||
|
# ---------- Node functions ----------------------------------------------
|
||||||
|
async def outline(state: BriefState) -> BriefState:
|
||||||
|
"""Generate an outline of 4‑5 research points for the topic."""
|
||||||
|
prompt = (
|
||||||
|
f"You are a research assistant.\n"
|
||||||
|
f"Topic: {state['topic']}\n"
|
||||||
|
f"Provide 4–5 concise bullet points that could serve as sections of a short research brief."
|
||||||
|
)
|
||||||
|
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||||||
|
text = response.content.strip()
|
||||||
|
# split by newlines or bullets
|
||||||
|
lines = [l.strip("- ") for l in text.splitlines() if l.strip()]
|
||||||
|
state["outline"] = lines[:5] # ensure max 5
|
||||||
|
state["step_index"] = 0
|
||||||
|
state["notes"] = []
|
||||||
|
return state
|
||||||
|
|
||||||
|
async def research_step(state: BriefState) -> BriefState:
|
||||||
|
"""For the current outline item perform a web search and create a short note."""
|
||||||
|
idx = state["step_index"]
|
||||||
|
if state["outline"] is None or idx >= len(state["outline"]):
|
||||||
|
return state
|
||||||
|
point = state["outline"][idx]
|
||||||
|
# Search via Tavily tool
|
||||||
|
search_query = f"{point} topic"
|
||||||
|
results = await search_tool.ainvoke(search_query)
|
||||||
|
# Build a short note (5‑8 sentences) summarizing the first result
|
||||||
|
if results:
|
||||||
|
snippet = results[0].snippet or ""
|
||||||
|
note = f"**{point}:** {snippet[:200]}..."
|
||||||
|
else:
|
||||||
|
note = f"**{point}:** No relevant information found."
|
||||||
|
state["notes"].append(note)
|
||||||
|
state["step_index"] += 1
|
||||||
|
return state
|
||||||
|
|
||||||
|
async def synthesize(state: BriefState) -> BriefState:
|
||||||
|
"""Combine all notes into a coherent brief with headings."""
|
||||||
|
if not state.get("notes"):
|
||||||
|
state["final_brief"] = "No research was conducted."
|
||||||
|
return state
|
||||||
|
sections = [f"### {note.split(':')[0][2:]}\n{note.split(':',1)[1].strip()}" for note in state["notes"]]
|
||||||
|
brief = "\n\n".join(sections)
|
||||||
|
state["final_brief"] = brief
|
||||||
|
return state
|
||||||
|
|
||||||
|
# ---------- Graph construction -------------------------------------------
|
||||||
|
builder = StateGraph(BriefState)
|
||||||
|
builder.add_node("outline", outline)
|
||||||
|
builder.add_node("research_step", research_step)
|
||||||
|
builder.add_node("synthesize", synthesize)
|
||||||
|
|
||||||
|
builder.set_entry_point("outline")
|
||||||
|
builder.add_edge("outline", "research_step")
|
||||||
|
# loop until all steps processed
|
||||||
|
builder.add_conditional_edges(
|
||||||
|
"research_step",
|
||||||
|
lambda state: "synthesize" if state["step_index"] >= len(state.get("outline", [])) else "research_step",
|
||||||
|
)
|
||||||
|
builder.set_finish_point("synthesize")
|
||||||
|
|
||||||
|
BriefGraph = builder.compile()
|
||||||
|
|
||||||
|
# ---------- Demo runner -----------------------------------------------
|
||||||
|
async def run_demo(topic: str) -> None:
|
||||||
|
from langgraph.checkpoint.memory import MemorySaver
|
||||||
|
memory = MemorySaver()
|
||||||
|
state = {"topic": topic, "outline": None, "step_index": 0, "notes": [], "final_brief": None}
|
||||||
|
result = await BriefGraph.ainvoke(state, config={"configurable": {"thread_id": "demo"}}, checkpointer=memory)
|
||||||
|
print("\n=== Outline ===")
|
||||||
|
for i, p in enumerate(result["outline"]):
|
||||||
|
print(f"{i+1}. {p}")
|
||||||
|
print("\n=== Notes ===")
|
||||||
|
for n in result["notes"]:
|
||||||
|
print(n)
|
||||||
|
print("\n=== Final Brief ===")
|
||||||
|
print(result["final_brief"])
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import asyncio
|
||||||
|
default_topic = "Как студенту безопасно подключать MCP к LangChain"
|
||||||
|
asyncio.run(run_demo(default_topic))
|
||||||
Reference in New Issue
Block a user