diff --git a/main.py b/main.py new file mode 100644 index 0000000..88f7f27 --- /dev/null +++ b/main.py @@ -0,0 +1,102 @@ +"""LangGraph research brief agent. + +This script implements a LangGraph agent that, given a topic, produces a short research brief. +The brief consists of: +1. An outline of 4–5 research points. +2. For each point, a single web search via Tavily and a short note. +3. A final synthesis of all notes into a coherent brief. + +The agent uses LangGraph's StateGraph and LangChain's Tavily and OpenAI LLM. +""" + +import os +from typing import TypedDict, List + +from langgraph.graph import StateGraph, END +from langchain_openai import ChatOpenAI +from langchain_tavily 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: str | None + +# --- LLM and Tavily ------------------------------------------------------- +llm = ChatOpenAI(temperature=0.2) +search = TavilySearchResults(max_results=1) + +# --- Node functions ------------------------------------------------------- +async def outline_node(state: BriefState) -> BriefState: + """Generate an outline of 4–5 research points for the topic.""" + prompt = ( + f"Generate a concise outline of 4–5 research points for the following topic. " + f"Return the points as a numbered list, one per line. Topic: {state['topic']}" + ) + response = await llm.agenerate([HumanMessage(content=prompt)]) + outline_text = response.generations[0][0].text.strip() + outline = [line.strip() for line in outline_text.splitlines() if line.strip()] + return {"outline": outline, "step_index": 0, "notes": []} + +async def research_step_node(state: BriefState) -> BriefState: + """Perform a single web search for the current outline point and store a short note.""" + idx = state["step_index"] + point = state["outline"][idx] + # Search via Tavily + search_results = await search.ainvoke(point) + snippet = search_results[0].snippet if search_results else "No snippet found." + # Summarize snippet with LLM + prompt = ( + f"You are a concise researcher. Based on the following snippet, write a 5–8 sentence note summarizing the key information. " + f"Snippet: {snippet}" + ) + note_resp = await llm.agenerate([HumanMessage(content=prompt)]) + note = note_resp.generations[0][0].text.strip() + notes = state["notes"] + [note] + return {"notes": notes, "step_index": idx + 1} + +async def synthesize_node(state: BriefState) -> BriefState: + """Combine all notes into a coherent brief with headings.""" + outline = state["outline"] + notes = state["notes"] + sections = [f"**{point}**\n{note}" for point, note in zip(outline, notes)] + brief = "\n\n".join(sections) + return {"final_brief": brief} + +# --- Graph construction --------------------------------------------------- +workflow = StateGraph(BriefState) +workflow.add_node("outline", outline_node) +workflow.add_node("research_step", research_step_node) +workflow.add_node("synthesize", synthesize_node) + +workflow.set_entry_point("outline") +workflow.add_conditional_edges( + "outline", + lambda state: "research_step" if state["outline"] else END, +) +workflow.add_conditional_edges( + "research_step", + lambda state: "research_step" if state["step_index"] < len(state["outline"]) else "synthesize", +) +workflow.add_edge("synthesize", END) + +graph = workflow.compile() + +# --- Runner ---------------------------------------------------------------- +async def run_brief(topic: str) -> str: + state: BriefState = {"topic": topic, "outline": None, "step_index": 0, "notes": [], "final_brief": None} + result = await graph.ainvoke(state) + return result["final_brief"] + +if __name__ == "__main__": + import asyncio + topic = os.getenv("TOPIC", "Как студенту безопасно подключать MCP к LangChain") + brief = asyncio.run(run_brief(topic)) + print("\n=== Brief ===") + print(brief)