add: main.py
This commit is contained in:
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user