Plan-and-Execute: сначала план, потом исполнение шагами

"""plan_and_execute.py — Plan-and-Execute агент"""
import anthropic
from dataclasses import dataclass, field

client = anthropic.Anthropic()

# ── Схема плана ──────────────────────────────────────────────────────────────

# Инструмент-планировщик: модель заполняет структуру, мы читаем block.input
PLAN_TOOL = {
    "name": "create_plan",
    "description": "Создаёт структурированный план выполнения задачи шагами",
    "input_schema": {
        "type": "object",
        "properties": {
            "goal": {
                "type": "string",
                "description": "Цель задачи — одно предложение"
            },
            "steps": {
                "type": "array",
                "description": "Шаги для достижения цели, в порядке исполнения",
                "items": {
                    "type": "object",
                    "properties": {
                        "id": {
                            "type": "integer",
                            "description": "Порядковый номер шага (начиная с 1)"
                        },
                        "description": {
                            "type": "string",
                            "description": "Что конкретно нужно сделать в этом шаге"
                        },
                        "tool_hint": {
                            "type": "string",
                            "description": "Какой инструмент, вероятно, потребуется (необязательно)"
                        },
                        "depends_on": {
                            "type": "array",
                            "items": {"type": "integer"},
                            "description": "Номера шагов, результаты которых нужны для этого шага"
                        }
                    },
                    "required": ["id", "description", "depends_on"]
                }
            }
        },
        "required": ["goal", "steps"]
    }
}

@dataclass
class Step:
    id: int
    description: str
    depends_on: list[int]
    tool_hint: str = ""
    result: str = ""

@dataclass
class Plan:
    goal: str
    steps: list[Step]

# ── Планировщик ──────────────────────────────────────────────────────────────

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

Разбивай задачу на конкретные, выполнимые шаги. Каждый шаг:
- Решает одну конкретную подзадачу
- Описан как действие: «Получить...», «Рассчитать...», «Найти...»
- Имеет явные зависимости (depends_on) от других шагов

Доступные инструменты агента: search (поиск), calculate (вычисления), get_data (данные).
Указывай depends_on: [] если шаг не зависит от других — это позволит запустить его параллельно."""

def create_plan(task: str) -> Plan:
    """Планировщик: один LLM-вызов → структурированный план."""
    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=2048,
        system=PLANNER_SYSTEM,
        tools=[PLAN_TOOL],
        tool_choice={"type": "tool", "name": "create_plan"},
        messages=[{"role": "user", "content": task}],
    )

    for block in response.content:
        if block.type == "tool_use" and block.name == "create_plan":
            raw = block.input
            steps = [
                Step(
                    id=s["id"],
                    description=s["description"],
                    depends_on=s.get("depends_on", []),
                    tool_hint=s.get("tool_hint", ""),
                )
                for s in raw["steps"]
            ]
            return Plan(goal=raw["goal"], steps=steps)

    raise RuntimeError("Планировщик не вернул план")

# ── Инструменты (заглушки) ───────────────────────────────────────────────────

TOOLS = [
    {
        "name": "search",
        "description": "Поиск информации в интернете",
        "input_schema": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"]
        }
    },
    {
        "name": "calculate",
        "description": "Вычисляет математическое выражение Python",
        "input_schema": {
            "type": "object",
            "properties": {"expression": {"type": "string"}},
            "required": ["expression"]
        }
    }
]

def execute_tool(name: str, args: dict) -> str:
    if name == "calculate":
        try:
            return str(eval(args["expression"]))  # noqa: S307
        except Exception as e:
            return f"Ошибка: {e}"
    if name == "search":
        return f"[Заглушка] Результаты поиска: {args['query']}"
    return f"Неизвестный инструмент: {name}"


# ── Executor ──────────────────────────────────────────────────────────────────

EXECUTOR_SYSTEM = """Ты — исполнитель конкретного шага в многошаговой задаче.

Тебе даётся один шаг и контекст выполненных предыдущих шагов.
Выполни ТОЛЬКО этот шаг, используя инструменты при необходимости.
Когда шаг выполнен — дай чёткий, конкретный ответ."""

def execute_step(step: Step, context: str, max_iterations: int = 5) -> str:
    """
    Executor: выполняет один шаг плана.
    Это mini-ReAct: LLM ↔ инструменты до получения результата.
    """
    task_prompt = f"Выполни шаг: {step.description}"
    if context:
        task_prompt += f"\n\nКонтекст предыдущих шагов:\n{context}"

    messages = [{"role": "user", "content": task_prompt}]

    for _ in range(max_iterations):
        response = client.messages.create(
            model="claude-opus-4-6",
            max_tokens=2048,
            system=EXECUTOR_SYSTEM,
            tools=TOOLS,
            messages=messages,
        )
        messages.append({"role": "assistant", "content": response.content})

        if response.stop_reason == "end_turn":
            # Шаг выполнен — вернуть текст ответа
            for block in response.content:
                if hasattr(block, "text") and block.text:
                    return block.text.strip()
            return "(шаг выполнен)"

        # Выполнить инструменты
        tool_results = []
        for block in response.content:
            if block.type == "tool_use":
                result = execute_tool(block.name, block.input)
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": result,
                })
        messages.append({"role": "user", "content": tool_results})

    return "(достигнут лимит итераций для шага)"

def synthesize(task: str, plan: Plan) -> str:
    """Синтез: собирает результаты всех шагов в финальный ответ."""
    steps_summary = "\n\n".join(
        f"Шаг {s.id}: {s.description}\nРезультат: {s.result}"
        for s in plan.steps
    )

    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=4096,
        system="Ты — аналитик. Синтезируй финальный ответ на основе результатов всех шагов.",
        messages=[{
            "role": "user",
            "content": (
                f"Задача: {task}\n\n"
                f"Цель плана: {plan.goal}\n\n"
                f"Результаты шагов:\n{steps_summary}\n\n"
                f"Сформулируй полный финальный ответ."
            )
        }]
    )

    for block in response.content:
        if hasattr(block, "text") and block.text:
            return block.text.strip()
    return "(нет ответа)"

def plan_and_execute(task: str) -> str:
    """
    Plan-and-Execute оркестратор.
    Фаза 1: создать план → Фаза 2: выполнить шаги → Фаза 3: синтезировать ответ.
    """
    # ── Фаза 1: Планирование ─────────────────────────────────────────────────
    print("🗺️  Планирование...")
    plan = create_plan(task)
    print(f"  Цель: {plan.goal}")
    print(f"  Шагов: {len(plan.steps)}")
    for s in plan.steps:
        dep_str = f" (зависит от: {s.depends_on})" if s.depends_on else " (независимый)"
        print(f"  {s.id}. {s.description}{dep_str}")

    # ── Фаза 2: Исполнение ───────────────────────────────────────────────────
    print("\n⚙️  Исполнение шагов...")
    context_parts: list[str] = []

    for step in plan.steps:
        print(f"\n  Шаг {step.id}: {step.description}")
        # Передаём контекст: результаты всех предыдущих шагов
        context = "\n\n".join(context_parts) if context_parts else ""
        step.result = execute_step(step, context)
        context_parts.append(f"Шаг {step.id} [{step.description}]:\n{step.result}")
        print(f"  ✓ {step.result[:120]}{'...' if len(step.result) > 120 else ''}")

    # ── Фаза 3: Синтез ───────────────────────────────────────────────────────
    print("\n📝 Синтез финального ответа...")
    answer = synthesize(task, plan)
    return answer


if __name__ == "__main__":
    task = (
        "Сравни продажи компании по кварталам за 2024 и 2023 годы. "
        "Вычисли рост по каждому кварталу, выяви тренд."
    )
    print(f"Задача: {task}\n")
    result = plan_and_execute(task)
    print(f"\n{'═' * 60}")
    print(f"✅ Финальный ответ:\n{result}")

REPLAN_TOOL = {
    "name": "update_plan",
    "description": "Обновляет оставшиеся шаги плана с учётом полученных данных",
    "input_schema": {
        "type": "object",
        "properties": {
            "needs_replanning": {
                "type": "boolean",
                "description": "True если план нужно изменить"
            },
            "reason": {
                "type": "string",
                "description": "Почему нужно перепланировать"
            },
            "updated_steps": {
                "type": "array",
                "description": "Обновлённые оставшиеся шаги (если needs_replanning=True)",
                "items": {
                    "type": "object",
                    "properties": {
                        "id": {"type": "integer"},
                        "description": {"type": "string"},
                        "depends_on": {"type": "array", "items": {"type": "integer"}}
                    },
                    "required": ["id", "description", "depends_on"]
                }
            }
        },
        "required": ["needs_replanning"]
    }
}

def maybe_replan(
    original_task: str,
    plan: Plan,
    completed_steps: list[Step],
    remaining_steps: list[Step],
) -> list[Step]:
    """
    Проверяет, нужно ли обновить план.
    Возвращает либо оригинальные оставшиеся шаги, либо обновлённые.
    """
    context = "\n".join(
        f"Шаг {s.id}: {s.description} → {s.result[:300]}"
        for s in completed_steps
    )
    remaining = "\n".join(
        f"{s.id}. {s.description}" for s in remaining_steps
    )

    response = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=1024,
        system="Оцени, нужно ли изменить оставшиеся шаги плана с учётом новых данных.",
        tools=[REPLAN_TOOL],
        tool_choice={"type": "auto"},  # модель сама решает, вызывать ли инструмент
        messages=[{
            "role": "user",
            "content": (
                f"Задача: {original_task}\n\n"
                f"Выполненные шаги:\n{context}\n\n"
                f"Оставшиеся шаги:\n{remaining}\n\n"
                f"Нужно ли изменить план? Вызови update_plan если да."
            )
        }]
    )

    for block in response.content:
        if block.type == "tool_use" and block.name == "update_plan":
            data = block.input
            if data.get("needs_replanning") and data.get("updated_steps"):
                print(f"  ↻ Репланирование: {data.get('reason', '')}")
                return [
                    Step(
                        id=s["id"],
                        description=s["description"],
                        depends_on=s.get("depends_on", []),
                    )
                    for s in data["updated_steps"]
                ]

    return remaining_steps  # план не изменился

import asyncio

async def execute_step_async(step: Step, context: str) -> str:
    """Async-обёртка executor-а для параллельного запуска."""
    return await asyncio.to_thread(execute_step, step, context)


async def plan_and_execute_parallel(task: str) -> str:
    """
    Plan-and-Execute с параллельным исполнением независимых шагов.
    Использует depends_on для определения порядка.
    """
    plan = create_plan(task)
    completed: dict[int, str] = {}  # id шага → результат

    # Обрабатываем шаги волнами: каждая волна — шаги, чьи зависимости выполнены
    while len(completed) < len(plan.steps):
        # Найти все шаги, которые можно запустить сейчас
        ready = [
            s for s in plan.steps
            if s.id not in completed
            and all(dep in completed for dep in s.depends_on)
        ]

        if not ready:
            # Зависимый граф нарушен — запустить оставшиеся принудительно
            ready = [s for s in plan.steps if s.id not in completed]

        print(f"  Запуск параллельно: {[s.id for s in ready]}")

        # Сформировать контекст для каждого шага из его зависимостей
        async def run_with_context(step: Step) -> tuple[int, str]:
            ctx = "\n\n".join(
                f"Шаг {dep_id}: {completed[dep_id]}"
                for dep_id in step.depends_on
                if dep_id in completed
            )
            result = await execute_step_async(step, ctx)
            return step.id, result

        # Выполнить все готовые шаги параллельно
        results = await asyncio.gather(*[run_with_context(s) for s in ready])
        for step_id, result in results:
            completed[step_id] = result
            # Сохранить результат в шаг для синтеза
            for s in plan.steps:
                if s.id == step_id:
                    s.result = result

    return synthesize(task, plan)


# Запуск: asyncio.run(plan_and_execute_parallel(task))