Plan-and-Execute
с перепланированием

from typing import Annotated, TypedDict
import operator
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

llm = ChatOpenAI(model="gpt-4o-mini")

class State(TypedDict):
    input: str
    plan: list[str]                                  # оставшиеся шаги
    past_steps: Annotated[list[tuple], operator.add] # (шаг, результат) — копится
    response: str                                    # финальный ответ

# --- Planner: структурированный список шагов ---
class Plan(BaseModel):
    steps: list[str]

planner = llm.with_structured_output(Plan)

def plan_step(state: State) -> dict:
    plan = planner.invoke(f"Составь пошаговый план для задачи:\n{state['input']}")
    return {"plan": plan.steps}

# --- Executor: исполняет ПЕРВЫЙ шаг (здесь — мини ReAct-агент) ---
executor = create_react_agent(llm, tools=[])      # сюда дают инструменты

def execute_step(state: State) -> dict:
    task = state["plan"][0]
    result = executor.invoke({"messages": [{"role": "user", "content": task}]})
    answer = result["messages"][-1].content
    return {"past_steps": [(task, answer)]}        # reducer допишет

from typing import Union
from langgraph.graph import StateGraph, START, END

# Replanner возвращает ЛИБО новый план, ЛИБО финальный ответ
class Response(BaseModel):
    response: str

class Act(BaseModel):
    action: Union[Plan, Response]      # план = продолжить, ответ = закончить

replanner = llm.with_structured_output(Act)

def replan_step(state: State) -> dict:
    out = replanner.invoke(
        f"Задача: {state['input']}\n"
        f"Исходный план: {state['plan']}\n"
        f"Уже сделано: {state['past_steps']}\n"
        f"Обнови план оставшихся шагов. Если задача решена — верни ответ."
    )
    if isinstance(out.action, Response):
        return {"response": out.action.response}   # закончили
    return {"plan": out.action.steps}              # новый план

def should_end(state: State) -> str:
    return "end" if state.get("response") else "execute"

builder = StateGraph(State)
builder.add_node("planner", plan_step)
builder.add_node("execute", execute_step)
builder.add_node("replan", replan_step)

builder.add_edge(START, "planner")
builder.add_edge("planner", "execute")
builder.add_edge("execute", "replan")
builder.add_conditional_edges("replan", should_end,
    {"execute": "execute", "end": END})            # план → исполнять / ответ → END

graph = builder.compile()
result = graph.invoke({"input": "Кто тренер победителя Australian Open 2024 у мужчин?"})
print(result["response"])

# Состояние: задача + план + сделанное + ответ
class State(TypedDict):
    input: str
    plan: list[str]
    past_steps: Annotated[list[tuple], operator.add]
    response: str

# 3 роли:
#  planner  → строит план  (структурированный вывод: Plan.steps)
#  execute  → исполняет plan[0]  (часто create_react_agent с инструментами)
#  replan   → План ИЛИ Response  (Union — продолжить или закончить)

def should_end(state):
    return "end" if state.get("response") else "execute"

# Граф: planner → execute → replan → (execute | END)
builder.add_edge(START, "planner")
builder.add_edge("planner", "execute")
builder.add_edge("execute", "replan")
builder.add_conditional_edges("replan", should_end,
    {"execute": "execute", "end": END})

# Правила:
#  • ОБЯЗАТЕЛЬНО replanning после каждого шага
#  • replanner умеет вернуть финальный ответ (иначе вечный цикл)
#  • план — структурированный (Pydantic), не свободный текст
#  • planner/replan — сильная модель, executor — дешёвая
#  • паттерн для ДЛИННЫХ задач; для простых — ReAct