From 9636a1f98e56528e0363d693a675a5b2af63674a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=93=D0=BB=D0=B5=D0=B1=20=D0=9D=D0=B8=D0=BA=D0=B8=D1=88?= =?UTF-8?q?=D0=B8=D0=BD?= Date: Thu, 4 Jun 2026 16:28:11 +0000 Subject: [PATCH] add: main.py --- main.py | 163 +++++++++++++++++++++++++++----------------------------- 1 file changed, 79 insertions(+), 84 deletions(-) diff --git a/main.py b/main.py index 4c5e55b..669548a 100644 --- a/main.py +++ b/main.py @@ -1,14 +1,20 @@ 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 -# --- LLM configuration (OpenRouter) --- +# 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", @@ -16,26 +22,22 @@ llm = ChatOpenAI( temperature=0.0, ) -# --- Backend for deepagents --- +# Backend for deepagents – local shell + filesystem backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# --- Tavily search tool --- +# Tavily search tool – real web search @tool def tavily_search(query: str) -> str: - """Search the web using Tavily and return a short summary of the top result.""" - from tavily import TavilyClient - client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) - results = client.search(query, max_results=1) - if not results: - return "No relevant information found." - # Return the first result's content (title + snippet) - first = results[0] - return f"{first.title}\n{first.snippet}" + """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 for LangGraph (used inside a tool) --- +# State definition class BriefState(TypedDict): topic: str outline: List[str] | None @@ -43,110 +45,103 @@ class BriefState(TypedDict): notes: List[str] final_brief: str | None -# --- LangGraph graph implementation --- -from langgraph.graph import StateGraph, START, END -from langgraph.graph.message import add_messages - -# Helper to format outline as a numbered list -def format_outline(outline: List[str]) -> str: - return "\n".join(f"{i+1}. {p}" for i, p in enumerate(outline)) - # Node: generate outline async def outline_node(state: BriefState) -> BriefState: - prompt = f"Create a concise 4–5 point outline for a research brief on the topic: {state['topic']}" + 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)]) - outline_text = response.content.strip() - # Assume the LLM returns a numbered list; split into lines - outline = [line.strip() for line in outline_text.splitlines() if line.strip()] - state['outline'] = outline - state['step_index'] = 0 - state['notes'] = [] + # 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: - current_point = state['outline'][state['step_index']] - # Use Tavily to get info - search_query = f"{current_point}" + step = state["outline"][state["step_index"]] + # Use Tavily to gather info + search_query = f"{state['topic']} {step}" search_result = tavily_search(search_query) - # Summarize the result with LLM - prompt = f"Summarize the following information in 5–8 sentences, focusing on the key points relevant to the research brief: {search_result}" - summary = await llm.ainvoke([HumanMessage(content=prompt)]) - state['notes'].append(f"**{current_point}**\n{summary.content.strip()}") - state['step_index'] += 1 + # 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_text = "\n\n".join(state['notes']) - prompt = f"Using the following notes, write a cohesive ½–1 page research brief. Include headings for each section.\n\n{notes_text}" + 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() + state["final_brief"] = brief.content.strip() return state -# Build the graph +# 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) -# Define transitions +# Entry point graph.set_entry_point("outline") +# After outline, loop research_step until all steps processed graph.add_conditional_edges( "outline", lambda _: "research_step", ) - -def research_cond(state: BriefState): - return "research_step" if state['step_index'] < len(state['outline']) else "synthesize" - -graph.add_conditional_edges("research_step", research_cond) +# 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) -# Compile the graph into a tool -from langgraph.prebuilt import create_react_agent +app = graph.compile() -# The graph will be used as a tool inside deepagents -@tool -async def run_brief(topic: str) -> str: - """Generate a research brief for the given topic.""" - # Initialize state - state: BriefState = { - "topic": topic, - "outline": None, - "step_index": 0, - "notes": [], - "final_brief": None, - } - # Run the graph - async for partial_state in graph.astream(state): - pass # we just wait for completion - # After completion, return the brief - return state['final_brief'] - -# --- DeepAgent setup --- +# DeepAgent wrapper agent = create_deep_agent( model=llm, - tools=[tavily_search, run_brief], + tools=[tavily_search], backend=backend, - system_prompt="You are a research assistant that creates concise research briefs. Use the provided tools to gather information and synthesize a brief.", + system_prompt="You are a research assistant that builds a brief.", ) -# --- Demo execution --- async def main(): - default_topic = "Как студенту безопасно подключать MCP к LangChain" - result = await agent.ainvoke( - {"messages": [HumanMessage(content=f"Create a research brief on: {default_topic}")]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print("\n--- Outline ---") - # The outline is part of the first tool call; extract it - for msg in result["messages"]: - if msg.type == "tool": - if "outline" in msg.content.lower(): - print(msg.content) - print("\n--- Final Brief ---") - print(result["messages"][-1].content) + # 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())