134 lines
3.7 KiB
Python
134 lines
3.7 KiB
Python
#!/usr/bin/env python
|
|
# main.py
|
|
# LangGraph planning agent example
|
|
|
|
import os
|
|
from typing import TypedDict, List, Dict, Any
|
|
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.prebuilt import create_chat_agent
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import HumanMessage, AIMessage
|
|
|
|
# ---------- 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,
|
|
)
|
|
|
|
# ---------- State ----------
|
|
class PlanningState(TypedDict):
|
|
task: str
|
|
plan: List[str] | None
|
|
current_step: int
|
|
results: List[str]
|
|
|
|
# ---------- Planning node ----------
|
|
planning_prompt = (
|
|
"You are a helpful assistant. Given the following task, break it into 3-6 concrete steps. "
|
|
"Return the steps as a numbered list or JSON array. The steps should be short, actionable, and in Russian."
|
|
)
|
|
|
|
async def planning(state: PlanningState) -> PlanningState:
|
|
task = state["task"]
|
|
messages = [HumanMessage(content=f"{planning_prompt}\nTask: {task}")]
|
|
response = await llm.ainvoke(messages)
|
|
text = response.content.strip()
|
|
|
|
# Try to parse JSON first
|
|
plan: List[str] | None = None
|
|
try:
|
|
import json
|
|
data = json.loads(text)
|
|
if isinstance(data, list):
|
|
plan = [str(item).strip() for item in data]
|
|
except Exception:
|
|
pass
|
|
|
|
# If JSON parsing failed, try to extract numbered list
|
|
if plan is None:
|
|
import re
|
|
lines = re.findall(r"\d+\.\s*(.+)", text)
|
|
if lines:
|
|
plan = [line.strip() for line in lines]
|
|
|
|
if plan is None:
|
|
raise ValueError("Could not parse plan from LLM response")
|
|
|
|
return {
|
|
"task": task,
|
|
"plan": plan,
|
|
"current_step": 0,
|
|
"results": [],
|
|
}
|
|
|
|
# ---------- Execution node ----------
|
|
async def execution(state: PlanningState) -> PlanningState:
|
|
plan = state["plan"]
|
|
idx = state["current_step"]
|
|
if plan is None or idx >= len(plan):
|
|
return state
|
|
|
|
step = plan[idx]
|
|
# Execute the step: ask LLM to produce result for this step
|
|
messages = [HumanMessage(content=f"Task: {state['task']}\nStep {idx+1}: {step}\nProvide the result for this step.")]
|
|
response = await llm.ainvoke(messages)
|
|
result = response.content.strip()
|
|
|
|
new_results = state["results"].copy()
|
|
new_results.append(result)
|
|
|
|
return {
|
|
"task": state["task"],
|
|
"plan": plan,
|
|
"current_step": idx + 1,
|
|
"results": new_results,
|
|
}
|
|
|
|
# ---------- Should continue ----------
|
|
async def should_continue(state: PlanningState) -> str:
|
|
if state["current_step"] >= len(state["plan"]):
|
|
return "finish"
|
|
return "execute"
|
|
|
|
# ---------- Build graph ----------
|
|
builder = StateGraph(PlanningState)
|
|
builder.add_node("planning", planning)
|
|
builder.add_node("execution", execution)
|
|
builder.add_conditional_edges(
|
|
"planning",
|
|
lambda _: "execute",
|
|
)
|
|
builder.add_conditional_edges(
|
|
"execution",
|
|
should_continue,
|
|
{
|
|
"execute": "execution",
|
|
"finish": END,
|
|
},
|
|
)
|
|
builder.set_entry_point("planning")
|
|
graph = builder.compile()
|
|
|
|
# ---------- Run example ----------
|
|
if __name__ == "__main__":
|
|
task = "Сравни Python и JavaScript"
|
|
initial_state: PlanningState = {
|
|
"task": task,
|
|
"plan": None,
|
|
"current_step": 0,
|
|
"results": [],
|
|
}
|
|
result = graph.invoke(initial_state)
|
|
|
|
print(f"\nЗадача: {task}\n")
|
|
print("План:")
|
|
for i, step in enumerate(result["plan"], 1):
|
|
print(f"{i}. {step}")
|
|
print("\n[Шаги]")
|
|
for i, res in enumerate(result["results"], 1):
|
|
print(f"[Шаг {i}] {res}\n")
|
|
print("Итог: " + "\n".join(result["results"]))
|