add: main.py
This commit is contained in:
@@ -0,0 +1,128 @@
|
|||||||
|
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 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
|
||||||
|
|
||||||
|
# ----------------- LLM -----------------
|
||||||
|
llm = ChatOpenAI(
|
||||||
|
model="openai/gpt-oss-20b:free",
|
||||||
|
base_url="https://openrouter.ai/api/v1",
|
||||||
|
api_key=os.getenv("OPENAI_API_KEY"),
|
||||||
|
temperature=0.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ----------------- Backend -----------------
|
||||||
|
backend = CompositeBackend([
|
||||||
|
LocalShellBackend(workspace_dir="./workspace"),
|
||||||
|
FilesystemBackend(),
|
||||||
|
])
|
||||||
|
|
||||||
|
# ----------------- Tavily Tool -----------------
|
||||||
|
search_tool = TavilySearchResults(max_results=3)
|
||||||
|
|
||||||
|
# ----------------- State -----------------
|
||||||
|
class BriefState(TypedDict):
|
||||||
|
topic: str
|
||||||
|
outline: List[str] | None
|
||||||
|
step_index: int
|
||||||
|
notes: List[str]
|
||||||
|
final_brief: str | None
|
||||||
|
|
||||||
|
# ----------------- Nodes -----------------
|
||||||
|
async def outline_node(state: BriefState) -> BriefState:
|
||||||
|
prompt = f"""Create a concise outline of 4–5 research steps for the topic: {state['topic']}. Return a JSON array of strings."""
|
||||||
|
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||||||
|
# Parse JSON array
|
||||||
|
import json
|
||||||
|
try:
|
||||||
|
outline = json.loads(response.content)
|
||||||
|
if not isinstance(outline, list):
|
||||||
|
raise ValueError
|
||||||
|
except Exception:
|
||||||
|
outline = ["Step 1: Define scope", "Step 2: Search", "Step 3: Analyze", "Step 4: Summarize"]
|
||||||
|
state.update(outline=outline, step_index=0, notes=[], final_brief=None)
|
||||||
|
return state
|
||||||
|
|
||||||
|
async def research_step_node(state: BriefState) -> BriefState:
|
||||||
|
step = state['outline'][state['step_index']]
|
||||||
|
# Perform web search via Tavily
|
||||||
|
results = await search_tool.ainvoke(step)
|
||||||
|
# Summarize results into 5–8 sentences
|
||||||
|
summary_prompt = f"""Summarize the following search results into 5–8 concise sentences for the research step: {step}.
|
||||||
|
|
||||||
|
Results:
|
||||||
|
{results}"""
|
||||||
|
summary = await llm.ainvoke([HumanMessage(content=summary_prompt)])
|
||||||
|
state['notes'].append(f"{step}\n{summary.content}")
|
||||||
|
state['step_index'] += 1
|
||||||
|
return state
|
||||||
|
|
||||||
|
async def synthesize_node(state: BriefState) -> BriefState:
|
||||||
|
# Combine notes into a coherent brief with headings
|
||||||
|
heading_prompt = """Combine the following notes into a ½–1 page research brief. Use the step titles as headings and write in a clear, academic style.
|
||||||
|
|
||||||
|
Notes:
|
||||||
|
""" + "\n\n".join(state['notes'])
|
||||||
|
brief = await llm.ainvoke([HumanMessage(content=heading_prompt)])
|
||||||
|
state['final_brief'] = brief.content
|
||||||
|
return state
|
||||||
|
|
||||||
|
# ----------------- Graph -----------------
|
||||||
|
graph = StateGraph(BriefState)
|
||||||
|
graph.add_node("outline", outline_node)
|
||||||
|
graph.add_node("research_step", research_step_node)
|
||||||
|
graph.add_node("synthesize", synthesize_node)
|
||||||
|
|
||||||
|
# Entry point
|
||||||
|
graph.set_entry_point("outline")
|
||||||
|
|
||||||
|
# Conditional edges
|
||||||
|
graph.add_conditional_edges(
|
||||||
|
"outline",
|
||||||
|
lambda _: "research_step",
|
||||||
|
)
|
||||||
|
|
||||||
|
graph.add_conditional_edges(
|
||||||
|
"research_step",
|
||||||
|
lambda state: "synthesize" if state['step_index'] >= len(state['outline']) else "research_step",
|
||||||
|
)
|
||||||
|
|
||||||
|
graph.add_edge("synthesize", END)
|
||||||
|
|
||||||
|
app = graph.compile()
|
||||||
|
|
||||||
|
# ----------------- DeepAgent -----------------
|
||||||
|
agent = create_deep_agent(
|
||||||
|
model=llm,
|
||||||
|
tools=[search_tool],
|
||||||
|
backend=backend,
|
||||||
|
system_prompt="You are a research assistant that builds a brief based on a topic.",
|
||||||
|
)
|
||||||
|
|
||||||
|
# ----------------- Main -----------------
|
||||||
|
async def main():
|
||||||
|
topic = os.getenv("DEFAULT_TOPIC", "Как студенту безопасно подключать MCP к LangChain")
|
||||||
|
# Run graph 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 each research step note
|
||||||
|
for i, note in enumerate(state['notes'], 1):
|
||||||
|
print(f"\n[Step {i}] {note.splitlines()[0]}")
|
||||||
|
print(note.splitlines()[1])
|
||||||
|
# Print final brief
|
||||||
|
print("\n=== Final Brief ===")
|
||||||
|
print(state['final_brief'])
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(main())
|
||||||
Reference in New Issue
Block a user