Task Decomposition Algorithms for AI Agents
Source: mortalapps.com- Task decomposition is the algorithmic process of breaking down an ambiguous, high-level prompt into a structured, executable Directed Acyclic Graph (DAG) of sub-tasks.
- It solves the critical problems of planning failures, context window exhaustion, and execution drift in complex, multi-step agent workflows.
- In production, structured decomposition enables parallel execution, precise cost attribution, and localized error recovery without restarting the entire run.
- This tutorial provides a complete, type-safe Python implementation of a dependency-aware decomposition engine using Pydantic and topological sorting.
Why This Matters
In complex agentic systems, executing a high-level, ambiguous user prompt directly through a single LLM call leads to high failure rates, hallucinated steps, and state drift. Algorithmic ai agent task decomposition resolves this by translating unstructured natural language into a structured Directed Acyclic Graph (DAG) of discrete, executable sub-tasks before any execution begins. This approach was invented to mirror human cognitive planning (such as Chain of Thought and Least-to-Most prompting) but formalized into software engineering constructs like dependency trees and topological execution queues.
If you ignore structured task decomposition in production, your agentic workflows will suffer from catastrophic cascading failures. A single API timeout or minor validation error in a late-stage step will force you to re-run the entire sequence, wasting expensive tokens and introducing unpredictable latencies. Furthermore, without a clear DAG, parallel execution of independent sub-tasks is impossible, severely limiting system throughput.
Use structured task decomposition when the user prompt requires multi-step orchestration, external tool integrations, or parallel processing paths (e.g., generating a comprehensive market report with parallel data fetching, synthesis, and formatting). Avoid it for simple, single-turn tasks or linear workflows where a lightweight state machine or simple router is more cost-effective and introduces less latency overhead. By decoupling planning from execution, engineering teams can implement deterministic validation, localized retries, and precise cost attribution at the individual sub-task level.
Core Concepts
To build a robust task decomposition engine, you must understand the following core concepts:
- Directed Acyclic Graph (DAG): A directed graph with no directed cycles. In task decomposition, nodes represent sub-tasks (e.g., "Fetch API Data") and directed edges represent execution dependencies (e.g., Task B cannot start until Task A completes).
- Granularity Heuristics: Rules or metrics used by the planner to determine when a sub-task is sufficiently atomic. Over-decomposition leads to token waste and orchestrator latency; under-decomposition leads to monolithic execution failures.
- Topological Sorting: An algorithmic ordering of DAG vertices such that for every directed edge $u \to v$, vertex $u$ comes before $v$ in the ordering. This determines the exact execution sequence.
- Dynamic Re-planning: The process of modifying the remaining DAG at runtime when a sub-task fails, returns unexpected data, or reveals new execution requirements.
- Sub-task State: The execution context of a single node, containing its inputs, outputs, status (Pending, Running, Completed, Failed), and execution metadata.
| Granularity Level | Description | Pros | Cons |
|---|---|---|---|
| Monolithic | Single LLM call handles all steps | Zero orchestration overhead | High failure rate, no parallelization |
| Compound | Coarse-grained steps (e.g., "Research & Write") | Moderate token cost, simpler graphs | Hard to recover from localized errors |
| Atomic | Fine-grained steps (e.g., "Fetch URL", "Parse HTML") | High parallelization, robust recovery | High orchestration latency and token cost |
How It Works
The task decomposition lifecycle consists of four distinct phases that transform an ambiguous prompt into a successfully executed set of sub-tasks.
Phase 1: Prompt Parsing and Schema Alignment
The user input is sent to the Planner LLM along with a strict JSON schema representing the target DAG. The system prompt instructs the LLM to analyze the request, identify necessary external tools, and break the request down into discrete steps. The output must conform to a structured schema containing an array of tasks, each with a unique ID, a task type, input parameters, and an array of dependency IDs.
Phase 2: Dependency Mapping and DAG Validation
Once the raw plan is generated, the orchestrator validates the graph structure. It parses the dependencies to ensure there are no circular references (e.g., Task A depends on Task B, which depends on Task A). If a cycle is detected, the plan is rejected, and a repair prompt is sent back to the LLM. The orchestrator also validates that all task inputs mapped from parent outputs are type-compatible.
Phase 3: Topological Sorting and Execution Queueing
The validated DAG is passed to a topological sorting algorithm (such as Kahn's algorithm). This sort produces a linear execution queue or, more commonly, a tiered execution plan where independent tasks are grouped together. Tasks with zero dependencies are marked as "Ready" and pushed to an active execution queue.
Phase 4: Runtime Execution and Failure Handling
The orchestrator pops "Ready" tasks from the queue and dispatches them to worker threads or async runners. When a task completes, its output is stored in the shared state context. The orchestrator then checks all downstream tasks that depended on the completed task. If all of a downstream task's dependencies are resolved, its state transitions to "Ready" and it is queued for execution. If a task fails, the orchestrator halts downstream execution of dependent nodes, triggers localized retries, or invokes the Planner LLM to dynamically re-plan the remaining graph based on the failure context.
Architecture
The Task Decomposition Engine architecture consists of six core components. Execution starts when a user submits a prompt to the API Gateway and ends when the final aggregator node returns the synthesized result.
- API Gateway: Receives the unstructured user prompt and forwards it to the orchestrator.
- Planner LLM (Decomposer): A high-capability LLM configured with structured output schemas. It parses the prompt and generates the raw task nodes and dependency edges.
- DAG Validator: A deterministic software component that validates the graph structure, checking for cycles using Depth-First Search (DFS) and verifying type-safety of input/output mappings.
- State Store: A persistent or in-memory database (e.g., Redis or PostgreSQL) that tracks the state, inputs, outputs, and execution logs of every sub-task.
- Execution Orchestrator: An asynchronous engine that performs topological sorting, manages the execution queue, dispatches tasks to workers, and handles state transitions.
- Worker Pool: A set of concurrent execution environments (threads, processes, or sandboxed containers) that run the actual tools or sub-agents associated with each task node.
Data flows from the API Gateway to the Planner, then through the Validator to the State Store. The Orchestrator continuously reads from the State Store, dispatches work to the Worker Pool, and writes execution results back to the State Store until the graph is exhausted.
Granularity Heuristics: Finding the Atomic Sweet Spot
Determining the optimal size of a sub-task is the most challenging aspect of task decomposition. If tasks are too large, the agent behaves like a monolithic system, losing the benefits of parallelization and localized recovery. If tasks are too small, the system suffers from "orchestration bloat," where the latency and cost of LLM planning calls exceed the actual execution time of the tasks.
To programmatically guide the Planner LLM, you must inject granularity heuristics into the system prompt. These heuristics should define:
- Single Responsibility Principle: A sub-task must perform exactly one action (e.g., call one API, write one file, or synthesize one set of data).
- Context Window Boundaries: A sub-task's input and output data must fit comfortably within a localized context window (typically under 4,000 tokens) to ensure high-quality processing.
- Tool Isolation: Any step requiring an external tool call must be isolated into its own sub-task. This allows the orchestrator to handle tool-specific timeouts, rate limits, and authentication failures independently.
Algorithmic DAG Validation: Cycle Detection and Schema Enforcement
Before executing any LLM-generated plan, you must programmatically guarantee that the graph is a valid DAG. LLMs frequently generate circular dependencies when dealing with complex, multi-step prompts.
To detect cycles, implement a Depth-First Search (DFS) algorithm using three-color marking:
- White (Unvisited): Nodes that have not been processed yet.
- Gray (Visiting): Nodes currently in the recursion stack. If you encounter a Gray node during traversal, a cycle exists.
- Black (Visited): Nodes that have been fully processed along with all their descendants.
If the validator detects a cycle, it must raise a validation error. In production, this error should trigger a self-correction loop, feeding the cyclic path back to the Planner LLM with instructions to break the loop.
Topological Sort and Parallel Execution Orchestration
Once validated, the DAG must be executed efficiently. A simple linear execution of a topological sort wastes the opportunity for parallel processing. Instead, use an event-driven orchestrator that tracks the in-degree (number of incoming dependencies) of each node.
- Calculate the in-degree for all nodes in the DAG.
- Identify all nodes with an in-degree of 0 and push them to the active execution queue.
- As each task completes, decrement the in-degree of its downstream neighbor nodes.
- If a neighbor node's in-degree reaches 0, push it to the active queue.
- Repeat until all nodes are executed or a terminal failure occurs.
This approach naturally maximizes concurrency, allowing independent branches of the graph to run in parallel across the Worker Pool.
Mid-Execution Failure Recovery and Dynamic Re-planning
In real-world environments, sub-tasks will fail due to network timeouts, API rate limits, or unexpected data formats. A robust task decomposition engine must handle these failures gracefully without restarting the entire DAG.
There are three primary recovery strategies:
- Localized Retries with Exponential Backoff: For transient errors (e.g., HTTP 503), the orchestrator retries the specific node up to a configured limit. Downstream nodes remain paused in the "Pending" state.
- Alternative Branch Execution: If a node fails permanently, the orchestrator can execute a pre-planned fallback branch (e.g., switching from a premium search API to a local database query).
- Dynamic Re-planning: For structural failures, the orchestrator halts execution, packages the current state of the DAG (including completed task outputs and the failure error message), and sends it back to the Planner LLM. The LLM then generates a *new* partial DAG to replace the remaining unexecuted nodes.
Code Example
import asyncio
import logging
import os
from enum import Enum
from typing import List, Dict, Any, Set
from pydantic import BaseModel, Field
from openai import OpenAI
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger("TaskDecomposer")
# Define the schema for structured output
class TaskType(str, Enum):
FETCH_DATA = "fetch_data"
ANALYZE_DATA = "analyze_data"
GENERATE_REPORT = "generate_report"
VALIDATE_OUTPUT = "validate_output"
class SubTask(BaseModel):
id: str = Field(..., description="Unique identifier for the task, e.g., 'task_01'")
type: TaskType = Field(..., description="The type of action to perform")
description: str = Field(..., description="Detailed description of what this sub-task does")
dependencies: List[str] = Field(default_factory=list, description="List of task IDs that must complete before this task starts")
input_parameters: Dict[str, Any] = Field(default_factory=dict, description="Parameters required for execution")
class TaskDAG(BaseModel):
tasks: List[SubTask] = Field(..., description="The complete list of sub-tasks making up the DAG")
# Cycle Detection using DFS (Three-Color Marking)
class CycleDetector:
def __init__(self, tasks: List[SubTask]):
self.adj_list: Dict[str, List[str]] = {t.id: t.dependencies for t in tasks}
self.visited: Dict[str, str] = {t.id: "WHITE" for t in tasks} # WHITE, GRAY, BLACK
def has_cycle(self) -> bool:
for node in self.adj_list:
if self.visited[node] == "WHITE":
if self._dfs(node):
return True
return False
def _dfs(self, node: str) -> bool:
self.visited[node] = "GRAY"
for neighbor in self.adj_list.get(node, []):
# If neighbor is not in graph, log error and skip (or treat as cycle/invalid)
if neighbor not in self.visited:
logger.warning(f"Task {node} depends on non-existent task {neighbor}")
continue
if self.visited[neighbor] == "GRAY":
return True # Cycle detected
if self.visited[neighbor] == "WHITE":
if self._dfs(neighbor):
return True
self.visited[node] = "BLACK"
return False
# Mock Tool Executor
async def execute_tool(task: SubTask, context: Dict[str, Any]) -> Dict[str, Any]:
logger.info(f"Starting execution of task: {task.id} ({task.type.value})")
# Simulate network latency
await asyncio.sleep(1.0)
# Resolve inputs that depend on parent outputs
resolved_inputs = {}
for k, v in task.input_parameters.items():
if isinstance(v, str) and v.startswith("$"):
parent_key = v[1:]
resolved_inputs[k] = context.get(parent_key, f"[Resolved {parent_key}]")
else:
resolved_inputs[k] = v
# Execute mock logic based on task type
if task.type == TaskType.FETCH_DATA:
result = {"status": "success", "data": f"Raw data fetched for {resolved_inputs.get('source', 'unknown')}"}
elif task.type == TaskType.ANALYZE_DATA:
result = {"status": "success", "analysis": f"Analyzed: {resolved_inputs.get('data', '')}"}
elif task.type == TaskType.GENERATE_REPORT:
result = {"status": "success", "report": f"Report generated based on: {resolved_inputs.get('analysis', '')}"}
else:
result = {"status": "success", "message": "Task executed successfully"}
logger.info(f"Completed task: {task.id}")
return result
# Orchestrator Engine
class DAGOrchestrator:
def __init__(self, dag: TaskDAG):
self.dag = dag
self.tasks_by_id = {t.id: t for t in dag.tasks}
self.context: Dict[str, Any] = {}
self.completed_tasks: Set[str] = set()
self.running_tasks: Set[str] = set()
async def run(self) -> Dict[str, Any]:
# Validate DAG before running
detector = CycleDetector(self.dag.tasks)
if detector.has_cycle():
raise ValueError("Catastrophic Error: Circular dependency detected in generated task plan.")
while len(self.completed_tasks) < len(self.dag.tasks):
# Find tasks ready to execute (all dependencies are in completed_tasks)
ready_tasks = []
for task in self.dag.tasks:
if task.id in self.completed_tasks or task.id in self.running_tasks:
continue
if all(dep in self.completed_tasks for dep in task.dependencies):
ready_tasks.append(task)
if not ready_tasks and not self.running_tasks:
raise RuntimeError("Execution stalled. Unresolved dependencies or disconnected graph.")
# Launch ready tasks concurrently
futures = []
for task in ready_tasks:
self.running_tasks.add(task.id)
futures.append(self._run_task_wrapper(task))
if futures:
await asyncio.gather(*futures)
else:
# Wait briefly for running tasks to complete
await asyncio.sleep(0.1)
return self.context
async def _run_task_wrapper(self, task: SubTask):
try:
output = await execute_tool(task, self.context)
self.context[task.id] = output
self.completed_tasks.add(task.id)
except Exception as e:
logger.error(f"Task {task.id} failed: {str(e)}")
raise e
finally:
self.running_tasks.remove(task.id)
# Planner / Decomposer using OpenAI Structured Outputs
def decompose_prompt(user_prompt: str) -> TaskDAG:
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
raise ValueError("OPENAI_API_KEY environment variable is not set.")
client = OpenAI(api_key=api_key)
system_instruction = (
"You are an expert AI planner. Break down the user's request into a set of discrete sub-tasks "
"represented as a Directed Acyclic Graph (DAG). Ensure tasks have clear dependencies. "
"Use the format '$task_id' in input_parameters to reference outputs of parent tasks."
)
completion = client.beta.chat.completions.parse(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": system_instruction},
{"role": "user", "content": user_prompt}
],
response_format=TaskDAG,
)
return completion.choices[0].message.parsed
# Execution Entrypoint
async def main():
# Example complex prompt requiring parallel fetching and sequential analysis
prompt = "Fetch data from Google and Bing, analyze both datasets, and generate a final market report."
try:
logger.info(f"Decomposing prompt: '{prompt}'")
dag = decompose_prompt(prompt)
logger.info(f"Decomposition successful. Generated {len(dag.tasks)} tasks.")
for t in dag.tasks:
logger.info(f" - Task {t.id}: {t.description} (Depends on: {t.dependencies})")
orchestrator = DAGOrchestrator(dag)
final_context = await orchestrator.run()
logger.info("All tasks executed successfully!")
logger.info(f"Final Context Keys: {list(final_context.keys())}")
except Exception as e:
logger.error(f"Execution failed: {str(e)}")
if __name__ == "__main__":
# Ensure OPENAI_API_KEY is set in your environment before running
asyncio.run(main())2026-03-30 10:00:00 - INFO - Decomposing prompt: 'Fetch data from Google and Bing, analyze both datasets, and generate a final market report.' 2026-03-30 10:00:02 - INFO - Decomposition successful. Generated 4 tasks. 2026-03-30 10:00:02 - INFO - - Task task_01: Fetch data from Google (Depends on: []) 2026-03-30 10:00:02 - INFO - - Task task_02: Fetch data from Bing (Depends on: []) 2026-03-30 10:00:02 - INFO - - Task task_03: Analyze Google and Bing data (Depends on: ['task_01', 'task_02']) 2026-03-30 10:00:02 - INFO - - Task task_04: Generate final market report (Depends on: ['task_03']) 2026-03-30 10:00:02 - INFO - Starting execution of task: task_01 (fetch_data) 2026-03-30 10:00:02 - INFO - Starting execution of task: task_02 (fetch_data) 2026-03-30 10:00:03 - INFO - Completed task: task_01 2026-03-30 10:00:03 - INFO - Completed task: task_02 2026-03-30 10:00:03 - INFO - Starting execution of task: task_03 (analyze_data) 2026-03-30 10:00:04 - INFO - Completed task: task_03 2026-03-30 10:00:04 - INFO - Starting execution of task: task_04 (generate_report) 2026-03-30 10:00:05 - INFO - Completed task: task_04 2026-03-30 10:00:05 - INFO - All tasks executed successfully! 2026-03-30 10:00:05 - INFO - Final Context Keys: ['task_01', 'task_02', 'task_03', 'task_04']