Планирование: как агент строит и пересматривает планы

from dataclasses import dataclass, field
from enum import Enum
from typing import Any
import time
import json


class StepStatus(str, Enum):
    PENDING  = "pending"   # ожидает исполнения
    RUNNING  = "running"   # выполняется прямо сейчас
    DONE     = "done"      # успешно завершён
    FAILED   = "failed"    # завершился с ошибкой
    SKIPPED  = "skipped"   # пропущен (например, недостижимый шаг)


@dataclass
class PlanStep:
    """Один шаг плана агента."""
    id: str                          # уникальный идентификатор: "step_1", "step_2a"
    title: str                       # короткое название для логов и UI
    description: str                 # что именно нужно сделать
    tool: str | None = None          # инструмент для вызова (или None если LLM-решает)
    tool_args: dict[str, Any] = field(default_factory=dict)
    depends_on: list[str] = field(default_factory=list)  # id шагов-предшественников

    # Состояние (заполняется в ходе исполнения)
    status: StepStatus = StepStatus.PENDING
    result: str | None = None        # результат успешного выполнения
    error: str | None = None         # сообщение об ошибке при сбое
    attempts: int = 0                # сколько раз пытались выполнить
    started_at: float | None = None
    finished_at: float | None = None

    @property
    def is_ready(self) -> bool:
        """Шаг готов к выполнению: все зависимости завершены."""
        return self.status == StepStatus.PENDING

    @property
    def duration(self) -> float | None:
        if self.started_at and self.finished_at:
            return self.finished_at - self.started_at
        return None


@dataclass
class Plan:
    """Полный план агента для выполнения задачи."""
    task: str                        # исходная задача пользователя
    goal: str                        # конечная цель (что считать успехом)
    steps: list[PlanStep] = field(default_factory=list)
    created_at: float = field(default_factory=time.time)
    revision: int = 0                # 0 = первоначальный план, 1+ = после перепланирования

    # ── Методы навигации по плану ──

    def current_step(self) -> PlanStep | None:
        """Следующий шаг, готовый к исполнению."""
        for step in self.steps:
            if step.status == StepStatus.PENDING:
                return step
        return None

    def failed_steps(self) -> list[PlanStep]:
        return [s for s in self.steps if s.status == StepStatus.FAILED]

    def done_steps(self) -> list[PlanStep]:
        return [s for s in self.steps if s.status == StepStatus.DONE]

    def is_complete(self) -> bool:
        """Все шаги завершены (успешно или пропущены)."""
        return all(s.status in (StepStatus.DONE, StepStatus.SKIPPED)
                   for s in self.steps)

    def is_stuck(self) -> bool:
        """Нет доступных шагов, но план не завершён — нужно перепланирование."""
        return self.current_step() is None and not self.is_complete()

    def summary(self) -> str:
        """Краткий текстовый отчёт о состоянии плана."""
        total = len(self.steps)
        done = len(self.done_steps())
        failed = len(self.failed_steps())
        return (f"Plan v{self.revision}: {done}/{total} done, "
                f"{failed} failed, revision={self.revision}")

    def to_context_string(self) -> str:
        """Форматируем план для вставки в промпт перепланировщика."""
        lines = [f"ЗАДАЧА: {self.task}", f"ЦЕЛЬ: {self.goal}", ""]
        for s in self.steps:
            icon = {"done": "✓", "failed": "✗", "running": "→",
                    "pending": "…", "skipped": "~"}[s.status]
            lines.append(f"  [{icon}] {s.id}: {s.title}")
            if s.result:
                lines.append(f"       результат: {s.result[:100]}")
            if s.error:
                lines.append(f"       ошибка: {s.error[:100]}")
        return "\n".join(lines)

import anthropic
import json
import re

client = anthropic.Anthropic()

PLANNER_SYSTEM = """Ты — планировщик задач для AI-агента. Разбивай задачу на конкретные шаги.

Правила:
- Каждый шаг выполняет ОДНО действие (не «сделай A и B»)
- Порядок шагов учитывает зависимости
- Шаги должны быть проверяемыми: у каждого есть чёткий критерий успеха
- Максимум 8 шагов на задачу. Если нужно больше — укрупни шаги
- Не добавляй «запасных» шагов без явной необходимости

Доступные инструменты: bash, read_file, write_file, http_get, http_post, python_exec

Верни JSON без пояснений:
{
  "goal": "Что является конечным результатом",
  "steps": [
    {
      "id": "step_1",
      "title": "Короткое название",
      "description": "Что конкретно делает этот шаг",
      "tool": "bash",
      "tool_args": {"command": "docker run ..."},
      "depends_on": []
    }
  ]
}"""


def generate_plan(task: str) -> Plan:
    """Генерирует план выполнения задачи через LLM."""
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        system=PLANNER_SYSTEM,
        messages=[{"role": "user", "content": f"Задача: {task}"}],
    )

    text = response.content[0].text.strip()

    # Вырезаем JSON даже если модель добавила markdown-блок
    match = re.search(r'\{.*\}', text, re.DOTALL)
    if not match:
        raise ValueError(f"Planner did not return valid JSON: {text[:200]}")

    data = json.loads(match.group())

    steps = []
    for raw in data["steps"]:
        steps.append(PlanStep(
            id=raw["id"],
            title=raw["title"],
            description=raw["description"],
            tool=raw.get("tool"),
            tool_args=raw.get("tool_args", {}),
            depends_on=raw.get("depends_on", []),
        ))

    return Plan(
        task=task,
        goal=data["goal"],
        steps=steps,
    )

import subprocess
import time


# ── Реестр инструментов ──

def tool_bash(command: str, **_) -> str:
    result = subprocess.run(
        command, shell=True, capture_output=True, text=True, timeout=30
    )
    if result.returncode != 0:
        raise RuntimeError(result.stderr.strip() or f"exit code {result.returncode}")
    return result.stdout.strip()


def tool_python_exec(code: str, **_) -> str:
    namespace = {}
    exec(code, namespace)  # noqa: S102 — только для доверенных команд
    return str(namespace.get("result", "executed"))


TOOLS = {
    "bash":        tool_bash,
    "python_exec": tool_python_exec,
}


# ── Исполнитель одного шага ──

MAX_ATTEMPTS = 2


def execute_step(step: PlanStep, context: dict[str, str]) -> None:
    """
    Выполняет один шаг плана.
    context — словарь {step_id: result} из уже выполненных шагов.
    Изменяет step.status, step.result или step.error in-place.
    """
    step.status = StepStatus.RUNNING
    step.started_at = time.time()

    for attempt in range(1, MAX_ATTEMPTS + 1):
        step.attempts = attempt
        try:
            # Подставляем результаты предыдущих шагов в аргументы
            args = resolve_args(step.tool_args, context)

            if step.tool and step.tool in TOOLS:
                result = TOOLS[step.tool](**args)
            else:
                # Если инструмент не указан — используем LLM для исполнения
                result = llm_execute_step(step, context)

            step.status  = StepStatus.DONE
            step.result  = str(result)[:500]   # обрезаем длинные выводы
            step.error   = None
            break

        except Exception as exc:
            step.error = str(exc)
            if attempt == MAX_ATTEMPTS:
                step.status = StepStatus.FAILED
            else:
                time.sleep(2 ** attempt)   # exponential backoff

    step.finished_at = time.time()


def resolve_args(args: dict, context: dict[str, str]) -> dict:
    """
    Подставляет {{step_1.result}} шаблоны в аргументы шага.
    Позволяет передавать вывод одного шага как вход другого.
    """
    resolved = {}
    for k, v in args.items():
        if isinstance(v, str):
            for step_id, result in context.items():
                v = v.replace(f"{{{{{step_id}.result}}}}", result)
        resolved[k] = v
    return resolved


def llm_execute_step(step: PlanStep, context: dict[str, str]) -> str:
    """Просим LLM выполнить шаг без явного инструмента."""
    ctx_str = "\n".join(f"{k}: {v}" for k, v in context.items())
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": (
                f"Выполни следующий шаг плана:\n"
                f"Шаг: {step.title}\n"
                f"Описание: {step.description}\n"
                f"Контекст предыдущих шагов:\n{ctx_str}\n\n"
                f"Верни только результат шага."
            )
        }]
    )
    return response.content[0].text.strip()


# ── Главный цикл исполнения ──

def execute_plan(plan: Plan, on_failure=None) -> bool:
    """
    Исполняет план шаг за шагом.
    on_failure(plan, step) — колбэк при сбое (например, перепланировщик).
    Возвращает True если план завершился успешно.
    """
    context: dict[str, str] = {}   # накапливаем результаты

    while True:
        step = plan.current_step()
        if step is None:
            break

        print(f"[{step.id}] {step.title}...")
        execute_step(step, context)

        if step.status == StepStatus.DONE:
            context[step.id] = step.result
            print(f"  ✓ {step.result[:80]}")
        else:
            print(f"  ✗ {step.error}")
            if on_failure:
                # Передаём управление перепланировщику
                should_continue = on_failure(plan, step)
                if not should_continue:
                    return False
            else:
                return False

    return plan.is_complete()

REPLANNER_SYSTEM = """Ты — перепланировщик задач. Тебе дан текущий план с ошибкой на одном из шагов.
Твоя задача — исправить план так, чтобы задача всё равно была выполнена.

Стратегии (выбери одну):
- RETRY: повторить тот же шаг (если транзиентная ошибка)
- SUBSTITUTE: заменить проваленный шаг альтернативным подходом
- SKIP: пропустить шаг (если он некритичный)
- REPLAN: перегенерировать оставшуюся часть плана

Верни JSON:
{
  "strategy": "SUBSTITUTE",
  "reasoning": "Объяснение выбора стратегии",
  "steps": [ ... полный обновлённый список шагов ... ]
}

Правила:
- Сохраняй уже выполненные шаги (status: "done") без изменений
- Для новых/изменённых шагов — начинай id с оригинального + суффикс (step_2 → step_2a)
- Не меняй задачу и цель, только маршрут к ней"""


def replan(plan: Plan, failed_step: PlanStep) -> Plan | None:
    """
    Перепланирует задачу после сбоя на конкретном шаге.
    Возвращает обновлённый план или None если восстановление невозможно.
    """
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        system=REPLANNER_SYSTEM,
        messages=[{
            "role": "user",
            "content": (
                f"Текущее состояние плана:\n{plan.to_context_string()}\n\n"
                f"Провалившийся шаг: {failed_step.id} — {failed_step.title}\n"
                f"Ошибка: {failed_step.error}\n"
                f"Попыток: {failed_step.attempts}\n\n"
                f"Исправь план."
            )
        }]
    )

    text = response.content[0].text.strip()
    match = re.search(r'\{.*\}', text, re.DOTALL)
    if not match:
        return None

    data = json.loads(match.group())
    strategy = data.get("strategy", "REPLAN")

    if strategy == "ABORT":
        print(f"Replanner decided to abort: {data.get('reasoning', '')}")
        return None

    # Обновляем шаги
    new_steps = []
    for raw in data["steps"]:
        # Ищем существующий шаг (чтобы не потерять статус DONE)
        existing = next((s for s in plan.steps if s.id == raw["id"]), None)
        if existing and existing.status == StepStatus.DONE:
            new_steps.append(existing)   # сохраняем выполненный шаг как есть
        else:
            new_steps.append(PlanStep(
                id=raw["id"],
                title=raw["title"],
                description=raw["description"],
                tool=raw.get("tool"),
                tool_args=raw.get("tool_args", {}),
                depends_on=raw.get("depends_on", []),
                status=StepStatus.PENDING,
            ))

    plan.steps = new_steps
    plan.revision += 1
    print(f"  Replanned (strategy={strategy}): {data.get('reasoning', '')[:80]}")
    return plan

MAX_REPLAN_ATTEMPTS = 3


class PlanningAgent:
    def __init__(self):
        self.client = anthropic.Anthropic()
        self.plan: Plan | None = None
        self._replan_count = 0

    def run(self, task: str) -> dict:
        """
        Полный цикл: план → исполнение → перепланирование при сбоях → отчёт.
        """
        print(f"\n=== Planning Agent: {task[:60]} ===\n")

        # Генерируем начальный план
        self.plan = generate_plan(task)
        print(f"Generated plan with {len(self.plan.steps)} steps:\n")
        for s in self.plan.steps:
            print(f"  {s.id}: {s.title}")
        print()

        # Запускаем исполнение с перепланировщиком как колбэком
        success = execute_plan(self.plan, on_failure=self._handle_failure)

        return self._build_report(success)

    def _handle_failure(self, plan: Plan, failed_step: PlanStep) -> bool:
        """
        Колбэк при сбое шага. Возвращает True если исполнение можно продолжить.
        """
        self._replan_count += 1
        if self._replan_count > MAX_REPLAN_ATTEMPTS:
            print(f"\nMax replan attempts ({MAX_REPLAN_ATTEMPTS}) exceeded. Aborting.")
            return False

        print(f"\n  [replan #{self._replan_count}] Failed: {failed_step.title}")
        updated = replan(plan, failed_step)

        if updated is None:
            print("  Replanner decided to abort.")
            return False

        # Продолжаем исполнение обновлённого плана
        return True

    def _build_report(self, success: bool) -> dict:
        """Итоговый отчёт о выполнении задачи."""
        done = self.plan.done_steps()
        failed = self.plan.failed_steps()
        total_duration = sum(
            s.duration for s in self.plan.steps
            if s.duration is not None
        )

        return {
            "success":      success,
            "task":         self.plan.task,
            "plan_revision": self.plan.revision,
            "steps_done":   len(done),
            "steps_failed": len(failed),
            "duration_sec": round(total_duration, 1),
            "results": {s.id: s.result for s in done},
            "errors":  {s.id: s.error for s in failed},
        }


# ── Использование ──

agent = PlanningAgent()
report = agent.run(
    "Подготовь тестовое окружение: подними PostgreSQL в Docker, "
    "создай БД test_db, загрузи фикстуры из fixtures/test_data.sql, "
    "проверь что сервис отвечает на порту 5432"
)

print("\n=== Report ===")
print(f"Success: {report['success']}")
print(f"Plan revisions: {report['plan_revision']}")
print(f"Steps done: {report['steps_done']}, failed: {report['steps_failed']}")
print(f"Total time: {report['duration_sec']}s")

# Минимальный рабочий паттерн планировщика

agent = PlanningAgent()
report = agent.run("Задача на естественном языке")

# Внутри происходит:
# 1. plan = generate_plan(task)         — LLM генерирует JSON-план
# 2. execute_plan(plan, on_failure=...) — цикл исполнения
#    a. step = plan.current_step()
#    b. execute_step(step, context)     — retry встроен
#    c. if failed → replan(plan, step)  — LLM перепланирует
# 3. build_report()                     — итог