add: main.py

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."""
+    """Search the web using Tavily and return a short summary."""
-    from tavily import TavilyClient
+    tavily = TavilySearchResults(max_results=3)
-    client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+    results = tavily.run(query)
-    results = client.search(query, max_results=1)
+    # Concatenate titles and snippets
-    if not results:
+    return "\n".join(f"{r['title']}: {r['content']}" for r in results)
-        return "No relevant information found."
-    # Return the first result's content (title + snippet)
-    first = results[0]
-    return f"{first.title}\n{first.snippet}"
 
-# --- 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()
+    # Extract JSON array
-    # Assume the LLM returns a numbered list; split into lines
+    import json, re
-    outline = [line.strip() for line in outline_text.splitlines() if line.strip()]
+    try:
-    state['outline'] = outline
+        array_text = re.search(r"\[.*\]", response.content, re.S).group(0)
-    state['step_index'] = 0
+        outline = json.loads(array_text)
-    state['notes'] = []
+    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']]
+    step = state["outline"][state["step_index"]]
-    # Use Tavily to get info
+    # Use Tavily to gather info
-    search_query = f"{current_point}"
+    search_query = f"{state['topic']} {step}"
     search_result = tavily_search(search_query)
-    # Summarize the result with LLM
+    # Summarize with LLM
-    prompt = f"Summarize the following information in 5–8 sentences, focusing on the key points relevant to the research brief: {search_result}"
+    prompt = (
-    summary = await llm.ainvoke([HumanMessage(content=prompt)])
+        f"Using the following search results, write a concise note (5–8 sentences) for the step: {step}\n"
-    state['notes'].append(f"**{current_point}**\n{summary.content.strip()}")
+        f"Search results:\n{search_result}\n"
-    state['step_index'] += 1
+        "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'])
+    notes = state["notes"]
-    prompt = f"Using the following notes, write a cohesive ½–1 page research brief. Include headings for each section.\n\n{notes_text}"
+    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",
 )
-
+# After each research_step, decide whether to continue or synthesize
-def research_cond(state: BriefState):
+graph.add_conditional_edges(
-    return "research_step" if state['step_index'] < len(state['outline']) else "synthesize"
+    "research_step",
-
+    lambda state: "synthesize" if state["step_index"] >= len(state["outline"]) else "research_step",
-graph.add_conditional_edges("research_step", research_cond)
+)
+# Final node
 graph.add_edge("synthesize", END)
 
-# Compile the graph into a tool
+app = graph.compile()
-from langgraph.prebuilt import create_react_agent
 
-# The graph will be used as a tool inside deepagents
+# DeepAgent wrapper
-@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 ---
 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"
+    # Default topic
-    result = await agent.ainvoke(
+    topic = os.getenv("DEFAULT_TOPIC", "Как студенту безопасно подключать MCP к LangChain")
-        {"messages": [HumanMessage(content=f"Create a research brief on: {default_topic}")]},
+    # Run LangGraph to get outline and notes
-        {"configurable": {"thread_id": "session-1"}},
+    state = await app.ainvoke({"topic": topic, "outline": None, "step_index": 0, "notes": [], "final_brief": None})
-    )
+    # Print outline
-    print("\n--- Outline ---")
+    print("\n=== Outline ===")
-    # The outline is part of the first tool call; extract it
+    for i, step in enumerate(state["outline"], 1):
-    for msg in result["messages"]:
+        print(f"{i}. {step}")
-        if msg.type == "tool":
+    # Print notes per step
-            if "outline" in msg.content.lower():
+    print("\n=== Notes ===")
-                print(msg.content)
+    for i, note in enumerate(state["notes"], 1):
-    print("\n--- Final Brief ---")
+        print(f"[Step {i}] {note}\n")
-    print(result["messages"][-1].content)
+    # Print final brief
+    print("\n=== Final Brief ===")
+    print(state["final_brief"])
 
 if __name__ == "__main__":
     asyncio.run(main())