diff --git a/main.py b/main.py index ebd6fd2..90b31d5 100644 --- a/main.py +++ b/main.py @@ -1,22 +1,19 @@ import os -import json import asyncio -from dotenv import load_dotenv -from typing import TypedDict, List, Optional +from typing import TypedDict, List, Annotated -from langchain_openai import ChatOpenAI -from langchain_tavily import TavilySearchResults +from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_core.messages import HumanMessage from langchain.tools import tool - +from langchain_tavily import TavilySearchResults +from langgraph.graph import StateGraph, START, END +from langgraph.graph.message import add_messages 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 +# ---------------------------------------------------------------------- +# Configuration +# ---------------------------------------------------------------------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -24,13 +21,32 @@ llm = ChatOpenAI( temperature=0.0, ) -# Backend for deepagents -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] +) -# TypedDict for graph state +# ---------------------------------------------------------------------- +# Tools +# ---------------------------------------------------------------------- +search_tool = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) + +@tool +def web_search(query: str) -> str: + """ + Perform a web search using Tavily and return a concise summary of the top results. + The LLM will ask for a short note based on this summary. + """ + results = search_tool.run(query) + # results is a list of dicts with 'url' and 'content' + summaries = [r.get("content", "") for r in results[:3]] + return "\n".join(summaries) if summaries else "No relevant results found." + +# ---------------------------------------------------------------------- +# State definition +# ---------------------------------------------------------------------- class BriefState(TypedDict): topic: str outline: List[str] | None @@ -38,110 +54,145 @@ class BriefState(TypedDict): 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 +# ---------------------------------------------------------------------- +# Nodes +# ---------------------------------------------------------------------- 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 + """Generate a 4-5 item outline for the given topic.""" + prompt = ( + f"Create a concise outline (4-5 bullet points) for a short research brief on the topic:\n" + f"\"{state['topic']}\"\n" + "Each bullet should be a short phrase suitable as a section heading." + ) + response = llm.invoke([HumanMessage(content=prompt)]) + outline_text = response.content.strip() + # Split on newlines and strip bullet characters + items = [line.lstrip("-• ").strip() for line in outline_text.splitlines() if line.strip()] + return { + **state, + "outline": items, + "step_index": 0, + "notes": [], + } -# 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 + """Research one outline item, produce a short note, and store it.""" + outline = state["outline"] + idx = state["step_index"] + if outline is None or idx >= len(outline): + return state # safety + + current_topic = outline[idx] + # Use web search tool + search_result = web_search(current_topic) + + # Prompt LLM to write a 5-8 sentence note based on search result + note_prompt = ( + f"Based on the following web search summary, write a short note (5-8 sentences) " + f"that could serve as a paragraph for a research brief about \"{state['topic']}\". " + f"Focus on the aspect: \"{current_topic}\".\n\n" + f"Search summary:\n{search_result}" + ) + note_response = llm.invoke([HumanMessage(content=note_prompt)]) + note = note_response.content.strip() + + new_notes = state["notes"] + [f"## {current_topic}\n{note}"] + return { + **state, + "notes": new_notes, + "step_index": idx + 1, + } -# 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}" + """Combine all notes into a coherent brief.""" + if not state["notes"]: + final = "No notes were collected." + else: + combined = "\n\n".join(state["notes"]) + synthesis_prompt = ( + f"Combine the following sections into a single cohesive research brief (about half to one page). " + f"Keep the headings, ensure logical flow, and add a brief introduction and conclusion.\n\n" + f"{combined}" ) - brief = llm.invoke(synth_prompt).content.strip() - state["final_brief"] = brief - return state + synthesis_response = llm.invoke([HumanMessage(content=synthesis_prompt)]) + final = synthesis_response.content.strip() + return { + **state, + "final_brief": final, + } -# 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 construction +# ---------------------------------------------------------------------- +workflow = StateGraph(BriefState) -graph.set_entry_point("outline") -graph.add_conditional_edges( +workflow.add_node("outline", outline_node) +workflow.add_node("research_step", research_step_node) +workflow.add_node("synthesize", synthesize_node) + +workflow.add_edge(START, "outline") +workflow.add_conditional_edges( "outline", - lambda state: "research_step" if state["outline"] is not None else END, + lambda state: "research_step" if state["outline"] else END, ) -graph.add_conditional_edges( - "research_step", - lambda state: ( - "research_step" - if state["step_index"] < len(state["outline"] or []) - else "synthesize" - ), + +def continue_condition(state: BriefState) -> str: + if state["outline"] is None: + return END + if state["step_index"] < len(state["outline"]): + return "research_step" + return "synthesize" + +workflow.add_edge("research_step", "research_step") +workflow.add_conditional_edges("research_step", continue_condition) +workflow.add_edge("synthesize", END) + +graph = workflow.compile() + +# ---------------------------------------------------------------------- +# DeepAgent wrapper (required by the course) +# ---------------------------------------------------------------------- +agent = create_deep_agent( + model=llm, + tools=[web_search], + backend=backend, + system_prompt="You are an AI research assistant that helps build short research briefs.", ) -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 = { +# ---------------------------------------------------------------------- +# Demo execution +# ---------------------------------------------------------------------- +async def main(): + topic = "Как студенту безопасно подключать MCP к LangChain" + # Initialize state + init_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"}}, + # Run the graph + result = await graph.ainvoke( + init_state, + config={"configurable": {"thread_id": "demo-1"}}, ) - print(result["messages"][-1].content) + + # Print results + print("\n--- Outline ---") + if result["outline"]: + for i, item in enumerate(result["outline"], 1): + print(f"{i}. {item}") + + print("\n--- Research Steps ---") + for i, note in enumerate(result["notes"], 1): + print(f"[Step {i}]") + print(note) + print() + + print("\n--- Final Brief ---") + print(result["final_brief"] or "No brief generated.") if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file