add: main.py
This commit is contained in:
@@ -0,0 +1,130 @@
|
||||
import os
|
||||
import asyncio
|
||||
from typing import TypedDict, Annotated
|
||||
|
||||
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 CompositeBackend, LocalShellBackend, FilesystemBackend
|
||||
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# ---------- 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(),
|
||||
])
|
||||
|
||||
# ---------- Planning State ----------
|
||||
class PlanningState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
plan: list[str]
|
||||
current_step: int
|
||||
results: list[str]
|
||||
|
||||
# ---------- Planner Node ----------
|
||||
async def planner_node(state: PlanningState):
|
||||
task = state["messages"][-1].content
|
||||
prompt = (
|
||||
"You are a task planner.\n"
|
||||
"Task: {task}\n"
|
||||
"Break the task into 3–6 concrete steps.\n"
|
||||
"Return the steps as a numbered list or JSON array.\n"
|
||||
"Do not include any other text."
|
||||
).format(task=task)
|
||||
plan_text = await llm.ainvoke([HumanMessage(content=prompt)])
|
||||
plan_str = plan_text.content.strip()
|
||||
# Try to parse JSON
|
||||
plan = []
|
||||
try:
|
||||
import json
|
||||
plan = json.loads(plan_str)
|
||||
if not isinstance(plan, list):
|
||||
raise ValueError
|
||||
except Exception:
|
||||
# Fallback to numbered list parsing
|
||||
import re
|
||||
plan = [line.strip() for line in plan_str.splitlines() if re.match(r"^\s*\d+\.", line)]
|
||||
return {
|
||||
"plan": plan,
|
||||
"current_step": 0,
|
||||
"results": [],
|
||||
}
|
||||
|
||||
# ---------- Executor Node ----------
|
||||
async def executor_node(state: PlanningState):
|
||||
step = state["plan"][state["current_step"]]
|
||||
prompt = (
|
||||
"You are an executor.\n"
|
||||
"Step: {step}\n"
|
||||
"Provide a concise result for this step."
|
||||
).format(step=step)
|
||||
result_text = await llm.ainvoke([HumanMessage(content=prompt)])
|
||||
result = result_text.content.strip()
|
||||
new_results = state["results"] + [result]
|
||||
return {
|
||||
"results": new_results,
|
||||
"current_step": state["current_step"] + 1,
|
||||
}
|
||||
|
||||
# ---------- Should Continue ----------
|
||||
def should_continue(state: PlanningState):
|
||||
if state["current_step"] >= len(state["plan"]):
|
||||
return "finish"
|
||||
return "execute"
|
||||
|
||||
# ---------- Graph ----------
|
||||
graph = StateGraph(PlanningState)
|
||||
graph.add_node("planner", planner_node)
|
||||
graph.add_node("executor", executor_node)
|
||||
graph.add_conditional_edges("planner", lambda _: "execute")
|
||||
graph.add_conditional_edges("executor", should_continue)
|
||||
graph.set_entry_point("planner")
|
||||
graph.add_edge("execute", "executor")
|
||||
graph.add_edge("finish", END)
|
||||
planner_chain = graph.compile()
|
||||
|
||||
# ---------- Tool that runs the planner graph ----------
|
||||
@tool
|
||||
def run_planner(task: str) -> str:
|
||||
"""Run the planning graph on the given task and return the final summary."""
|
||||
# Initialize state with the task as a message
|
||||
init_state = {"messages": [HumanMessage(content=task)], "plan": [], "current_step": 0, "results": []}
|
||||
final_state = planner_chain.invoke(init_state)
|
||||
# Build final output
|
||||
plan_lines = [f"{i+1}. {step}" for i, step in enumerate(final_state["plan"])]
|
||||
results = final_state["results"]
|
||||
summary = "\n".join(results)
|
||||
return (
|
||||
f"План:\n" + "\n".join(plan_lines) + "\n\n[Шаги]" + "\n".join([f"[Шаг {i+1}] {r}" for i, r in enumerate(results)]) + "\n\nИтог: " + summary
|
||||
)
|
||||
|
||||
# ---------- Deep Agent ----------
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[run_planner],
|
||||
backend=backend,
|
||||
system_prompt="You are a helpful agent that can plan and execute tasks.",
|
||||
)
|
||||
|
||||
# ---------- Demo ----------
|
||||
async def main():
|
||||
task = "Сравни Python и JavaScript"
|
||||
result = await agent.ainvoke(
|
||||
{"messages": [HumanMessage(content=task)]},
|
||||
{"configurable": {"thread_id": "session-1"}},
|
||||
)
|
||||
print(result["messages"][-1].content)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user