Planning & Task Decomposition

import anthropic
import json

client = anthropic.Anthropic()  # reads ANTHROPIC_API_KEY env var

CLASSIFY_PROMPT = """Classify the user's request into one of these intents:
- simple_question: Can be answered directly without tools
- single_tool: Requires one tool call (search, calculate, etc.)
- multi_step_task: Requires multiple steps with dependencies
- creative_generation: Writing, brainstorming, or content creation

Return JSON with this exact format:
{
  "intent": "simple_question|single_tool|multi_step_task|creative_generation",
  "complexity": "low|medium|high",
  "entities": ["extracted", "key", "terms"],
  "needs_planning": true/false,
  "reasoning": "One sentence explaining your classification"
}"""


def classify_intent(user_request: str) -> dict:
    """Classify the user's intent and determine if planning is needed."""
    try:
        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=256,
            system=CLASSIFY_PROMPT,
            messages=[{"role": "user", "content": user_request}],
        )
        text = response.content[0].text

        # Extract JSON from the response
        # Claude may wrap it in ```json ... ``` or return plain JSON
        if "```" in text:
            text = text.split("```")[1]
            if text.startswith("json"):
                text = text[4:]
        return json.loads(text.strip())

    except (json.JSONDecodeError, IndexError, KeyError) as e:
        # Fallback: treat as multi-step to be safe
        return {
            "intent": "multi_step_task",
            "complexity": "medium",
            "entities": [],
            "needs_planning": True,
            "reasoning": f"Classification failed ({e}), defaulting to planning",
        }


# --- Usage ---
result = classify_intent(
    "Research the top 3 project management tools, compare their "
    "pricing, and draft a summary report for our team"
)
print(json.dumps(result, indent=2))

DECOMPOSE_PROMPT = """Break the user's goal into a task graph.
Return JSON with this structure:
{
  "tasks": [
    {
      "id": "t1",
      "description": "What this task does",
      "type": "tool_call|generation|aggregation",
      "dependencies": [],
      "tools_hint": ["suggested_tool_names"]
    }
  ]
}

Rules:
- Each leaf task should be achievable in 1-2 tool calls
- Use dependencies to indicate which tasks must finish first
- tasks with no dependencies can run in parallel
- IDs must be unique strings (t1, t2, etc.)
- No circular dependencies allowed"""


def decompose_task(goal: str, entities: list[str]) -> dict:
    """Decompose a complex goal into a DAG of sub-tasks."""
    context = f"Key entities: {', '.join(entities)}" if entities else ""

    try:
        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=1024,
            system=DECOMPOSE_PROMPT,
            messages=[{
                "role": "user",
                "content": f"Goal: {goal}\n{context}",
            }],
        )
        text = response.content[0].text
        if "```" in text:
            text = text.split("```")[1]
            if text.startswith("json"):
                text = text[4:]
        plan = json.loads(text.strip())

        # Validate: check for cycles using topological sort
        if not _validate_dag(plan["tasks"]):
            raise ValueError("Task graph contains cycles")

        return plan

    except Exception as e:
        # Fallback: simple sequential plan
        return {
            "tasks": [
                {
                    "id": "t1",
                    "description": goal,
                    "type": "tool_call",
                    "dependencies": [],
                    "tools_hint": [],
                }
            ],
            "error": f"Decomposition failed: {e}",
        }


def _validate_dag(tasks: list[dict]) -> bool:
    """Validate that the task graph is acyclic using topological sort."""
    task_ids = {t["id"] for t in tasks}
    in_degree = {t["id"]: 0 for t in tasks}
    adj = {t["id"]: [] for t in tasks}

    for t in tasks:
        for dep in t.get("dependencies", []):
            if dep not in task_ids:
                return False  # dependency references non-existent task
            adj[dep].append(t["id"])
            in_degree[t["id"]] += 1

    # Kahn's algorithm
    queue = [tid for tid, deg in in_degree.items() if deg == 0]
    visited = 0
    while queue:
        node = queue.pop(0)
        visited += 1
        for neighbor in adj[node]:
            in_degree[neighbor] -= 1
            if in_degree[neighbor] == 0:
                queue.append(neighbor)

    return visited == len(tasks)  # All nodes visited = no cycles


# --- Usage ---
plan = decompose_task(
    "Research top 3 competitors, compare pricing, draft report",
    ["competitors", "pricing", "report"],
)
for task in plan["tasks"]:
    deps = ", ".join(task["dependencies"]) or "none"
    print(f"  {task['id']}: {task['description']} (deps: {deps})")

import asyncio
from enum import Enum

class TaskStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    DONE = "done"
    FAILED = "failed"


class DAGExecutor:
    """Execute a DAG of tasks with parallel execution."""

    def __init__(self, tasks: list[dict], execute_fn):
        self.tasks = {t["id"]: t for t in tasks}
        self.status = {t["id"]: TaskStatus.PENDING for t in tasks}
        self.results = {}
        self.execute_fn = execute_fn  # async function(task) -> result

    def _get_ready_tasks(self) -> list[str]:
        """Find tasks whose dependencies are all satisfied."""
        ready = []
        for tid, task in self.tasks.items():
            if self.status[tid] != TaskStatus.PENDING:
                continue
            deps = task.get("dependencies", [])
            if all(self.status.get(d) == TaskStatus.DONE for d in deps):
                ready.append(tid)
        return ready

    async def execute(self, on_progress=None) -> dict:
        """Run the DAG to completion, executing parallel tasks concurrently."""
        while True:
            ready = self._get_ready_tasks()
            if not ready:
                # No ready tasks — either all done or deadlocked
                break

            # Execute all ready tasks in parallel
            self.status.update({tid: TaskStatus.RUNNING for tid in ready})
            if on_progress:
                on_progress(self.status, self.results)

            results = await asyncio.gather(*[
                self._run_task(tid) for tid in ready
            ], return_exceptions=True)

            for tid, result in zip(ready, results):
                if isinstance(result, Exception):
                    self.status[tid] = TaskStatus.FAILED
                    self.results[tid] = f"Error: {result}"
                else:
                    self.status[tid] = TaskStatus.DONE
                    self.results[tid] = result

            if on_progress:
                on_progress(self.status, self.results)

        return {
            "results": self.results,
            "status": {k: v.value for k, v in self.status.items()},
            "completed": sum(
                1 for v in self.status.values() if v == TaskStatus.DONE
            ),
            "total": len(self.tasks),
        }

    async def _run_task(self, task_id: str):
        """Execute a single task with error handling."""
        task = self.tasks[task_id]
        # Gather results from dependencies as context
        dep_context = {
            dep: self.results.get(dep, "")
            for dep in task.get("dependencies", [])
        }
        try:
            return await self.execute_fn(task, dep_context)
        except Exception as e:
            raise RuntimeError(
                f"Task {task_id} failed: {e}"
            ) from e


# --- Usage with a simple executor function ---
async def execute_task(task: dict, dep_context: dict) -> str:
    """Execute a single sub-task using Claude."""
    context_str = "\n".join(
        f"Result of {k}: {v}" for k, v in dep_context.items()
    )
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=512,
        messages=[{
            "role": "user",
            "content": (
                f"Task: {task['description']}\n"
                f"Previous results:\n{context_str}"
            ),
        }],
    )
    return response.content[0].text


async def main():
    plan = decompose_task(
        "Research top 3 competitors, compare pricing, draft report",
        ["competitors", "pricing", "report"],
    )
    executor = DAGExecutor(plan["tasks"], execute_task)
    result = await executor.execute(
        on_progress=lambda s, r: print(
            f"  Progress: {sum(1 for v in s.values() if v == TaskStatus.DONE)}"
            f"/{len(s)} done"
        )
    )
    print(f"\nCompleted: {result['completed']}/{result['total']}")

asyncio.run(main())

mkdir planning-lab && cd planning-lab
python -m venv venv && source venv/bin/activate   # Windows: venv\Scripts\activate
pip install "anthropic>=0.30.0"
export ANTHROPIC_API_KEY=your-key-here             # Windows: set ANTHROPIC_API_KEY=your-key-here

import anthropic
import json
import asyncio
import time

client = anthropic.Anthropic()

# ── Intent Classifier ────────────────────────────────────────
def classify_intent(request: str) -> dict:
    """Classify whether a request needs planning or direct answer."""
    response = client.messages.create(
        model="claude-sonnet-4-6", max_tokens=256,
        system=(
            "Classify the user request. Respond with JSON only:\n"
            '{"intent": "direct|research|multi_step", '
            '"complexity": "simple|moderate|complex", '
            '"needs_planning": true/false, '
            '"reason": "one sentence why"}'
        ),
        messages=[{"role": "user", "content": request}],
    )
    try:
        return json.loads(response.content[0].text)
    except json.JSONDecodeError:
        return {"intent": "direct", "complexity": "simple",
                "needs_planning": False, "reason": "Parse error, defaulting to direct"}

# ── Task Decomposer ─────────────────────────────────────────
def decompose_task(goal: str) -> list[dict]:
    """Break a complex goal into a DAG of sub-tasks."""
    response = client.messages.create(
        model="claude-sonnet-4-6", max_tokens=1024,
        system=(
            "Decompose this goal into 3-7 sub-tasks. Respond with JSON only:\n"
            '[{"id": "task_1", "description": "...", "depends_on": [], '
            '"tools_needed": ["search", "analyze", etc]}, ...]'
            "\n\nRules:\n"
            "- Each task should be achievable in 1-2 tool calls\n"
            "- depends_on lists task IDs that must complete first\n"
            "- Independent tasks should have empty depends_on (they run in parallel)\n"
            "- NO circular dependencies"
        ),
        messages=[{"role": "user", "content": f"Goal: {goal}"}],
    )
    try:
        tasks = json.loads(response.content[0].text)
        # Validate: check for cycles using topological sort
        if not _validate_dag(tasks):
            print("  ⚠ WARNING: Cycle detected in task DAG, removing problematic edges")
        return tasks
    except json.JSONDecodeError:
        return [{"id": "task_1", "description": goal, "depends_on": [], "tools_needed": []}]

def _validate_dag(tasks: list[dict]) -> bool:
    """Check for circular dependencies via topological sort."""
    task_ids = {t["id"] for t in tasks}
    in_degree = {t["id"]: len(t.get("depends_on", [])) for t in tasks}
    deps = {t["id"]: set(t.get("depends_on", [])) for t in tasks}
    queue = [tid for tid, deg in in_degree.items() if deg == 0]
    visited = 0
    while queue:
        tid = queue.pop(0)
        visited += 1
        for other_id, other_deps in deps.items():
            if tid in other_deps:
                in_degree[other_id] -= 1
                if in_degree[other_id] == 0:
                    queue.append(other_id)
    return visited == len(tasks)  # True if no cycle

# ── Task Executor (simulated) ────────────────────────────────
async def execute_task(task: dict) -> dict:
    """Execute a single sub-task (simulated)."""
    await asyncio.sleep(0.5)  # simulate work
    return {
        "task_id": task["id"],
        "status": "completed",
        "result": f"Completed: {task['description'][:60]}",
    }

# ── DAG Executor ─────────────────────────────────────────────
async def execute_dag(tasks: list[dict], verbose: bool = True) -> list[dict]:
    """Execute tasks respecting dependencies, parallelizing where possible."""
    completed = {}
    failed = set()
    results = []
    remaining = {t["id"]: t for t in tasks}
    wave = 0

    while remaining:
        wave += 1
        # Find tasks whose dependencies are all satisfied
        ready = [
            t for tid, t in remaining.items()
            if all(d in completed for d in t.get("depends_on", []))
            and not any(d in failed for d in t.get("depends_on", []))
        ]

        if not ready:
            # Remaining tasks are blocked by failures
            for tid in remaining:
                results.append({"task_id": tid, "status": "blocked", "result": "Blocked by failed dependency"})
            break

        if verbose:
            task_names = [t["id"] for t in ready]
            mode = "PARALLEL" if len(ready) > 1 else "SEQUENTIAL"
            print(f"    Wave {wave} [{mode}]: {task_names}")

        # Execute ready tasks in parallel
        wave_results = await asyncio.gather(
            *[execute_task(t) for t in ready],
            return_exceptions=True,
        )

        for task, result in zip(ready, wave_results):
            if isinstance(result, Exception):
                failed.add(task["id"])
                results.append({"task_id": task["id"], "status": "failed", "result": str(result)})
            else:
                completed[task["id"]] = result
                results.append(result)
            del remaining[task["id"]]

    return results

# ── Progress Visualization ───────────────────────────────────
def print_plan(tasks: list[dict], results: list[dict] = None) -> None:
    """Print a visual task tree with status indicators."""
    status_map = {}
    if results:
        for r in results:
            status_map[r["task_id"]] = r["status"]

    print("\n  📋 Execution Plan:")
    for t in tasks:
        status = status_map.get(t["id"], "pending")
        icon = {"completed": "✅", "failed": "❌", "blocked": "🚫", "pending": "⏳"}.get(status, "⏳")
        deps = f" (after: {', '.join(t.get('depends_on', []))})" if t.get("depends_on") else " (no deps)"
        print(f"    {icon} {t['id']}: {t['description'][:50]}{deps}")

# ── Full Pipeline ────────────────────────────────────────────
async def planning_pipeline(request: str, verbose: bool = True) -> str:
    """Classify → (Plan → Execute) or Direct Answer."""
    if verbose:
        print(f"\n{'═' * 55}")
        print(f"  Request: {request}")
        print(f"{'═' * 55}")

    # Step 1: Classify
    classification = classify_intent(request)
    if verbose:
        print(f"\n  📊 Classification: {classification['intent']} "
              f"({classification['complexity']})")
        print(f"     Needs planning: {classification['needs_planning']}")
        print(f"     Reason: {classification['reason']}")

    if not classification["needs_planning"]:
        # Direct answer — skip planning overhead
        if verbose:
            print("  → Routing to direct answer (no planning needed)")
        response = client.messages.create(
            model="claude-sonnet-4-6", max_tokens=1024,
            messages=[{"role": "user", "content": request}],
        )
        return response.content[0].text

    # Step 2: Decompose into DAG
    if verbose:
        print("\n  🔨 Decomposing into sub-tasks...")
    tasks = decompose_task(request)
    if verbose:
        print_plan(tasks)

    # Step 3: Execute DAG
    if verbose:
        print(f"\n  ▶ Executing {len(tasks)} tasks...")
    start = time.time()
    results = await execute_dag(tasks, verbose=verbose)
    elapsed = time.time() - start

    if verbose:
        print_plan(tasks, results)
        completed = sum(1 for r in results if r["status"] == "completed")
        print(f"\n  ⏱ Completed {completed}/{len(tasks)} tasks in {elapsed:.1f}s")

    # Step 4: Synthesize final answer
    result_text = "\n".join(f"- {r['task_id']}: {r['result']}" for r in results)
    response = client.messages.create(
        model="claude-sonnet-4-6", max_tokens=1024,
        system="Synthesize a final answer from these sub-task results.",
        messages=[{"role": "user", "content":
            f"Original request: {request}\n\nSub-task results:\n{result_text}"}],
    )
    return response.content[0].text


# ── Tests ────────────────────────────────────────────────────
async def main():
    # Test 1: Simple task (should skip planning)
    print("\n▶ TEST 1: Simple task (direct answer)")
    r1 = await planning_pipeline("What is 2 + 2?")
    print(f"\n  Answer: {r1[:120]}")

    # Test 2: Complex task (should trigger planning)
    print("\n▶ TEST 2: Complex task (planning + execution)")
    r2 = await planning_pipeline(
        "Research the top 3 AI agent frameworks, compare their features, "
        "and draft a recommendation for a startup team."
    )
    print(f"\n  Answer: {r2[:200]}...")

    # Test 3: DAG validation
    print("\n▶ TEST 3: DAG validation")
    valid = _validate_dag([
        {"id": "a", "depends_on": []},
        {"id": "b", "depends_on": ["a"]},
        {"id": "c", "depends_on": ["a"]},
        {"id": "d", "depends_on": ["b", "c"]},
    ])
    print(f"  Valid DAG (no cycles): {valid}")

    cyclic = _validate_dag([
        {"id": "a", "depends_on": ["c"]},
        {"id": "b", "depends_on": ["a"]},
        {"id": "c", "depends_on": ["b"]},
    ])
    print(f"  Cyclic DAG (should be False): {cyclic}")

if __name__ == "__main__":
    asyncio.run(main())