From 92197f12d15036bdd2b5f00dfc88e89819c11e2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 16:08:49 +0000 Subject: [PATCH] add main.py --- main.py | 130 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..2db9803 --- /dev/null +++ b/main.py @@ -0,0 +1,130 @@ +""" +LangGraph Research Brief Agent +============================= + +This repository implements a LangGraph agent that generates a short research brief. +The agent follows the specification from the BroJS assignment: + +* Build an outline of 4‑5 points for a given topic. +* For each point perform one web search via Tavily and collect a concise note. +* Synthesize all notes into a coherent brief (≈½–1 page). +* +The implementation uses the official `langgraph` library, `langchain-openai` +for LLM calls and `langchain-tavily` for web searching. The agent is +exposed through a simple CLI that accepts a topic as an argument. +""" + +from __future__ import annotations + +import os +import sys +from typing import TypedDict, List + +from langgraph.graph import StateGraph, START, END +from langgraph.checkpoint.memory import MemorySaver +from langchain_openai import ChatOpenAI +from langchain_tavily import TavilySearchResults +from langchain_core.messages import HumanMessage, SystemMessage + +# --------------------------------------------------------------------------- +# 1. State definition +# --------------------------------------------------------------------------- +class BriefState(TypedDict): + topic: str + outline: List[str] | None + step_index: int + notes: List[str] + final_brief: str | None + +# --------------------------------------------------------------------------- +# 2. 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.2, +) + +search_tool = TavilySearchResults(max_results=3, tavily_api_key=os.getenv("TAVILY_API_KEY")) + +# --------------------------------------------------------------------------- +# 3. Nodes +# --------------------------------------------------------------------------- +async def outline_node(state: BriefState) -> dict: + """Generate a short outline of 4‑5 research points.""" + system = SystemMessage( + content="You are an assistant that creates a concise outline for a research brief.") + user = HumanMessage(content=f"Create 4–5 bullet points outlining the main aspects to cover when researching: {state['topic']}") + response = await llm.ainvoke([system, user]) + # Parse bullets + bullets = [line.strip("- ").strip() for line in response.content.splitlines() if line.strip().startswith("-")] + return {"outline": bullets, "step_index": 0, "notes": []} + +async def research_step_node(state: BriefState) -> dict: + """For the current outline point perform a web search and collect a short note.""" + point = state["outline"][state["step_index"]] + # Search via Tavily + results = await search_tool.ainvoke({"query": point}) + notes_text = "\n".join([f"{i+1}. {r['title']}: {r['content']}" for i, r in enumerate(results)]) + new_notes = state["notes"] + [f"**{point}**:\n{notes_text}"] + next_index = state["step_index"] + 1 + return {"notes": new_notes, "step_index": next_index} + +async def synthesize_node(state: BriefState) -> dict: + """Combine all notes into a single brief.""" + system = SystemMessage(content="You are an assistant that writes a concise research brief.") + user = HumanMessage( + content=f"Using the following notes, write a ½–1 page brief on {state['topic']}:\n\n{chr(10).join(state['notes'])}") + response = await llm.ainvoke([system, user]) + return {"final_brief": response.content.strip()} + +# --------------------------------------------------------------------------- +# 4. Graph definition +# --------------------------------------------------------------------------- +builder = StateGraph(BriefState) +builder.add_node("outline", outline_node) +builder.add_node("research_step", research_step_node) +builder.add_node("synthesize", synthesize_node) + +builder.set_entry_point("outline") +builder.add_edge("outline", "research_step") +# Loop until all points processed +builder.add_conditional_edges( + "research_step", + lambda state: "synthesize" if state["step_index"] >= len(state["outline"]) else "research_step", +) +builder.add_edge("synthesize", END) + +graph = builder.compile(checkpointer=MemorySaver()) + +# --------------------------------------------------------------------------- +# 5. CLI helper +# --------------------------------------------------------------------------- +async def run_brief(topic: str) -> None: + state: BriefState = {"topic": topic, "outline": None, "step_index": 0, "notes": [], "final_brief": None} + result = await graph.ainvoke(state) + print("\n=== Outline ===") + for i, point in enumerate(result["outline"]): + print(f"{i+1}. {point}") + print("\n=== Notes ===") + for note in result["notes"]: + print(note) + print("\n=== Final Brief ===") + print(result["final_brief"]) + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python main.py ''") + sys.exit(1) + topic = sys.argv[1] + import asyncio + asyncio.run(run_brief(topic)) + +# --------------------------------------------------------------------------- +# 6. Example usage (for documentation only, not executed by the script) +# --------------------------------------------------------------------------- +# Example 1: "How to integrate LangGraph with Tavily" +# Example 2: "Best practices for building research briefs in AI" +# Example 3: "Using LangChain and LangGraph for educational projects" +"""