Building a task-decomposition layer for your AI agent

To build task decomposition layer AI agent systems that survive real workloads, you need a planning module that turns vague goals into structured subtasks. This tutorial implements one in Python with an LLM-based planner and a lightweight execution loop you can drop into an existing service.

Prerequisites

• Python 3.11 or newer
• openai and pydantic v2 (pip install openai pydantic)
• An API key for any OpenAI-compatible chat endpoint
• Comfort with dict, basic graph algorithms, and CLI env vars

Create a clean environment and export your key:

python -m venv .venv && source .venv/bin/activate
pip install openai pydantic
export OPENAI_API_KEY="sk-..."

If you later route through a gateway, only the client constructor changes.

Define the plan schema

A decomposition layer is useless without a contract. Force the model to emit a directed acyclic graph of subtasks. Each node carries an id, a natural-language description, and a list of prerequisite ids.

from pydantic import BaseModel, Field
from typing import List

class Subtask(BaseModel):
    id: str = Field(pattern=r"^[a-z0-9_]+$")
    depends_on: List[str] = Field(default_factory=list)

class Plan(BaseModel):
    goal: str
    subtasks: List[Subtask]

The id regex keeps downstream routing and logging sane. depends_on references other subtask ids; empty means the node is a root.

Call the LLM for decomposition

Use function calling to bind the schema. The model invokes emit_plan instead of free-forming JSON, which removes a parsing step.

from openai import OpenAI
import os, json

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

def decompose(goal: str) -> Plan:
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role":"system","content":"You decompose goals into subtasks."},
            {"role":"user","content":goal}
        ],
        tools=[{
            "type":"function",
            "function":{
                "name":"emit_plan",
                "description":"Return a task graph",
                "parameters":Plan.model_json_schema()
            }
        }],
        tool_choice={"type":"function","function":{"name":"emit_plan"}}
    )
    tool_call = resp.choices[0].message.tool_calls[0]
    return Plan.model_validate(json.loads(tool_call.function.arguments))

Checkpoint: run decompose("Ship a CLI that queries weather and writes a report"). Expected shape:

{
  "goal": "Ship a CLI that queries weather and writes a report",
  "subtasks": [
    {"id":"fetch_weather","description":"Call weather API for city","depends_on":[]},
    {"id":"format_report","description":"Render markdown report","depends_on":["fetch_weather"]},
    {"id":"write_file","description":"Save report to disk","depends_on":["format_report"]}
  ]
}

The tool_choice lock guarantees the model does not chat back; you get a single structured call.

Validate and topologically sort

Never trust a planner’s graph. Sort before execution with Kahn’s algorithm so you catch cycles at plan time, not at runtime.

from collections import defaultdict, deque

def topo_order(plan: Plan) -> List[Subtask]:
    indeg = {t.id: 0 for t in plan.subtasks}
    adj = defaultdict(list)
    id_map = {t.id: t for t in plan.subtasks}
    for t in plan.subtasks:
        for dep in t.depends_on:
            adj[dep].append(t.id)
            indeg[t.id] += 1
    q = deque([i for i, d in indeg.items() if d == 0])
    out = []
    while q:
        n = q.popleft()
        out.append(id_map[n])
        for m in adj[n]:
            indeg[m] -= 1
            if indeg[m] == 0:
                q.append(m)
    if len(out) != len(plan.subtasks):
        raise ValueError("Cycle detected in plan")
    return out

Run topo_order(plan) on the sample above. Output order: fetch_weather → format_report → write_file.

Execute subtasks

Map ids to callables. For the tutorial, stub with side effects and a shared context object.

def exec_fetch_weather(ctx):
    print("FETCH weather")
    ctx["weather"] = {"temp": 12}

def exec_format_report(ctx):
    print("FORMAT report")
    ctx["report"] = f"Temp {ctx['weather']['temp']}C"

def exec_write_file(ctx):
    print("WRITE", ctx["report"])

REGISTRY = {
    "fetch_weather": exec_fetch_weather,
    "format_report": exec_format_report,
    "write_file": exec_write_file,
}

def run_plan(plan: Plan):
    ctx = {}
    for task in topo_order(plan):
        REGISTRY[task.id](ctx)

Expected console:

FETCH weather
FORMAT report
WRITE Temp 12C

Add resilience without rewrites

The decomposition call is the fragile part. If you point the OpenAI client at n4n.ai’s OpenAI-compatible endpoint, you get automatic fallback when a provider is rate-limited or degraded, and per-token usage metering, without touching the decompose function.

client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url="https://api.n4n.ai/v1"
)

Everything else stays identical. The gateway honors your model field and forwards cache-control hints if you set them.

Handle failures per subtask

Production agents need retry and skip logic. Wrap execution so a flaky node does not kill the whole run.

import time

def run_plan_safe(plan: Plan, retries=2):
    ctx = {}
    for task in topo_order(plan):
        for attempt in range(retries+1):
            try:
                REGISTRY[task.id](ctx)
                break
            except Exception as e:
                if attempt == retries:
                    raise RuntimeError(f"Subtask {task.id} failed") from e
                time.sleep(0.5 * (attempt+1))

This isolates a failure to its subtree. Dependent tasks naturally skip because ctx lacks required keys; you can extend with explicit cancellation.

Parallelize independent subtasks

When the graph fans out, sequential execution wastes cycles. Swap the registry call for an async executor that gathers roots.

import asyncio

async def run_async(plan: Plan):
    ctx = {}
    ordered = topo_order(plan)
    completed = set()
    while len(completed) < len(ordered):
        ready = [t for t in ordered if t.id not in completed and all(d in completed for d in t.depends_on)]
        await asyncio.gather(*[asyncio.to_thread(REGISTRY[t.id], ctx) for t in ready])
        completed.update(t.id for t in ready)

This runs fetch_weather alone, then format_report, then write_file—but if you added a second root, both roots fire concurrently.

Test the planner offline

You should not call the LLM in unit tests. Record one response and replay it.

import json

def test_decompose_offline():
    canned = {
        "goal":"test",
        "subtasks":[{"id":"a","description":"do a","depends_on":[]}]
    }
    plan = Plan.model_validate(canned)
    assert topo_order(plan)[0].id == "a"

Mock the client.chat.completions.create call with pytest-mock to assert the tool schema is sent correctly.

Compose into a service

To build task decomposition layer AI agent features into a larger system, expose decompose and run_plan_safe behind a queue. The planner can run on a cheap model; execution workers scale independently.

def handle_goal(goal: str):
    plan = decompose(goal)
    run_plan_safe(plan)
    return "done"

That is the core. From here, add streaming progress, human-in-the-loop approval on subtasks with depends_on length zero, and persistent state for long plans.

Where to go next

Swap the linear registry for an async executor using asyncio.gather on independent subtasks. Add a validator that rejects subtasks referencing unknown ids. Instrument the planner with logging of token usage from resp.usage. The pattern holds whether you run three subtasks or three hundred. When you need to build task decomposition layer AI agent logic that spans multiple providers, keep the schema strict and the executor dumb.