96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
from typing import TypedDict, List, Optional
|
|
import json
|
|
import asyncio
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.prebuilt import create_conditional_node
|
|
|
|
class PlanningState(TypedDict):
|
|
task: str
|
|
plan: List[str] | None
|
|
current_step: int
|
|
results: List[str]
|
|
|
|
async def planning(state: PlanningState, llm: ChatOpenAI) -> PlanningState:
|
|
"""LLM generates a JSON plan with a list of steps."""
|
|
prompt = (
|
|
"Разбей задачу на 3–6 конкретных шагов. Возвращай JSON с ключом \"plan\".\n\n"
|
|
f"Задача: {state['task']}"
|
|
)
|
|
response = await llm.ainvoke(prompt)
|
|
# response may be a string or a ChatResult
|
|
text = response if isinstance(response, str) else response.content
|
|
plan: List[str] = []
|
|
try:
|
|
data = json.loads(text)
|
|
plan = data.get("plan", [])
|
|
except Exception:
|
|
# Fallback: parse numbered list
|
|
for line in text.splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
# Remove leading number or bullet
|
|
if line[0].isdigit():
|
|
parts = line.split('.', 1)
|
|
if len(parts) == 2:
|
|
step = parts[1].strip()
|
|
else:
|
|
step = line
|
|
plan.append(step)
|
|
elif line[0] in ('-','*'):
|
|
plan.append(line[1:].strip())
|
|
return {
|
|
**state,
|
|
"plan": plan,
|
|
"current_step": 0,
|
|
"results": []
|
|
}
|
|
|
|
async def execution(state: PlanningState) -> PlanningState:
|
|
"""Execute a single step and record the result."""
|
|
if state["plan"] is None or state["current_step"] >= len(state["plan"]):
|
|
return state
|
|
step = state["plan"][state["current_step"]]
|
|
result = f"Шаг {state['current_step']+1}: {step}"
|
|
new_results = state["results"] + [result]
|
|
return {
|
|
**state,
|
|
"results": new_results,
|
|
"current_step": state["current_step"] + 1
|
|
}
|
|
|
|
def should_continue(state: PlanningState) -> str:
|
|
"""Decide whether to finish or execute another step."""
|
|
if state["current_step"] >= len(state["plan"] or []):
|
|
return "finish"
|
|
return "execute"
|
|
|
|
async def planning_node(state: PlanningState, llm: ChatOpenAI) -> PlanningState:
|
|
return await planning(state, llm)
|
|
|
|
def create_agent(llm: ChatOpenAI):
|
|
workflow = StateGraph(PlanningState)
|
|
# Add nodes
|
|
workflow.add_node("planning", lambda state: planning_node(state, llm))
|
|
workflow.add_node("execution", execution)
|
|
workflow.add_conditional_node("should_continue", should_continue, {
|
|
"execute": "execution",
|
|
"finish": END
|
|
})
|
|
# Set entry point and edges
|
|
workflow.set_entry_point("planning")
|
|
workflow.add_edge("planning", "execution")
|
|
workflow.add_edge("execution", "should_continue")
|
|
return workflow.compile()
|
|
|
|
async def run_agent(task: str, llm: ChatOpenAI) -> PlanningState:
|
|
agent = create_agent(llm)
|
|
initial_state: PlanningState = {
|
|
"task": task,
|
|
"plan": None,
|
|
"current_step": 0,
|
|
"results": []
|
|
}
|
|
result = await agent.ainvoke(initial_state)
|
|
return result |