Skip to content

feat: implement task decomposition and routing engine (#14) #165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Aureliolo merged 4 commits into from
Mar 8, 2026
Merged

feat: implement task decomposition and routing engine (#14) #165

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c091212
feat: implement task decomposition and routing engine (#14)
Aureliolo Mar 8, 2026
dfc99f8
fix: incorporate pre-PR review fixes for task decomposition and routing
Aureliolo Mar 8, 2026
93f602e
fix: incorporate pre-PR review fixes for task decomposition and routing
Aureliolo Mar 8, 2026
5e539d2
fix: address CodeRabbit round-2 feedback (log before raise, missing r…
Aureliolo Mar 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ src/ai_company/
communication/ # Message bus, dispatcher, messenger, channels, delegation, loop prevention
config/ # YAML company config loading and validation
core/ # Shared domain models and base classes
engine/ # Agent orchestration, execution loops, parallel execution, task lifecycle, recovery, and shutdown
engine/ # Agent orchestration, execution loops, parallel execution, task decomposition, routing, task lifecycle, recovery, and shutdown
memory/ # Persistent agent memory (memory layer TBD)
observability/ # Structured logging, correlation tracking, log sinks
providers/ # LLM provider abstraction (LiteLLM adapter)
Expand Down
17 changes: 17 additions & 0 deletions DESIGN_SPEC.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2385,6 +2385,21 @@ ai-company/
│ │ ├── parallel_models.py # AgentAssignment, ParallelExecutionGroup, AgentOutcome, ParallelExecutionResult, ParallelProgress
│ │ ├── resource_lock.py # ResourceLock protocol + InMemoryResourceLock
│ │ ├── shutdown.py # Graceful shutdown strategy & manager
│ │ ├── decomposition/ # Task decomposition subsystem
│ │ │ ├── __init__.py # Package exports
│ │ │ ├── classifier.py # TaskStructureClassifier (sequential/parallel/mixed)
│ │ │ ├── dag.py # DependencyGraph (validation, topo sort, parallel groups)
│ │ │ ├── manual.py # ManualDecompositionStrategy
│ │ │ ├── models.py # SubtaskDefinition, DecompositionPlan, DecompositionResult, SubtaskStatusRollup, DecompositionContext
│ │ │ ├── protocol.py # DecompositionStrategy protocol
│ │ │ ├── rollup.py # StatusRollup (compute subtask status aggregation)
│ │ │ └── service.py # DecompositionService (orchestrates strategy + classifier + DAG)
│ │ ├── routing/ # Task routing subsystem
│ │ │ ├── __init__.py # Package exports
│ │ │ ├── models.py # RoutingCandidate, RoutingDecision, RoutingResult, AutoTopologyConfig
│ │ │ ├── scorer.py # AgentTaskScorer (skill/role/seniority matching)
│ │ │ ├── service.py # TaskRoutingService (routes subtasks to agents)
│ │ │ └── topology_selector.py # TopologySelector (auto coordination topology)
│ │ ├── task_engine.py # Task routing & scheduling (M4)
│ │ ├── workflow_engine.py # Workflow orchestration (M4)
│ │ ├── meeting_engine.py # Meeting coordination (M4)
Expand Down Expand Up @@ -2436,6 +2451,7 @@ ai-company/
│ │ │ ├── config.py # CONFIG_* constants
│ │ │ ├── delegation.py # DELEGATION_* constants
│ │ │ ├── correlation.py # CORRELATION_* constants
│ │ │ ├── decomposition.py # DECOMPOSITION_* constants
│ │ │ ├── execution.py # EXECUTION_* constants
│ │ │ ├── git.py # GIT_* constants
│ │ │ ├── parallel.py # PARALLEL_* constants
Expand All @@ -2446,6 +2462,7 @@ ai-company/
│ │ │ ├── routing.py # ROUTING_* constants
│ │ │ ├── sandbox.py # SANDBOX_* constants
│ │ │ ├── task.py # TASK_* constants
│ │ │ ├── task_routing.py # TASK_ROUTING_* constants
│ │ │ ├── template.py # TEMPLATE_* constants
│ │ │ └── tool.py # TOOL_* constants
│ │ ├── processors.py # Log processors
Expand Down
1 change: 1 addition & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ AI Company lets you spin up a virtual organization staffed entirely by AI agents
- **Multi-Provider** - Any LLM via LiteLLM — cloud APIs, OpenRouter (400+ models), local Ollama, and more
- **Smart Cost Management** - Per-agent budget tracking, auto model routing, CFO agent optimization
- **Hierarchical Delegation** - Chain-of-command task delegation with five-mechanism loop prevention
- **Task Decomposition & Routing** - DAG-based subtask decomposition, structure classification, and agent-task scoring
- **Configurable Autonomy** - From fully autonomous to human-approves-everything, with a Security Ops agent in between
- **Persistent Memory** - Agents remember past decisions, code, relationships (memory layer TBD)
- **HR System** - Hire, fire, promote agents. HR agent analyzes skill gaps and proposes candidates
Expand Down
4 changes: 4 additions & 0 deletions src/ai_company/core/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
CompanyType,
Complexity,
ConflictApproach,
CoordinationTopology,
CostTier,
CreativityLevel,
DecisionMakingStyle,
Expand All @@ -43,6 +44,7 @@
SeniorityLevel,
SkillCategory,
TaskStatus,
TaskStructure,
TaskType,
ToolAccessLevel,
ToolCategory,
Expand Down Expand Up @@ -84,6 +86,7 @@
"CompanyType",
"Complexity",
"ConflictApproach",
"CoordinationTopology",
"CostTier",
"CreativityLevel",
"CustomRole",
Expand Down Expand Up @@ -114,6 +117,7 @@
"SkillSet",
"Task",
"TaskStatus",
"TaskStructure",
"TaskType",
"Team",
"ToolAccessLevel",
Expand Down
26 changes: 26 additions & 0 deletions src/ai_company/core/enums.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,32 @@ class ConflictApproach(StrEnum):
COLLABORATE = "collaborate"


class TaskStructure(StrEnum):
"""Classification of how a task's subtasks relate to each other.

Used by the decomposition engine to determine coordination topology
and execution ordering. See DESIGN_SPEC Section 6.9.
"""

SEQUENTIAL = "sequential"
PARALLEL = "parallel"
MIXED = "mixed"


class CoordinationTopology(StrEnum):
"""Coordination topology for multi-agent task execution.

Determines how agents coordinate when executing decomposed subtasks.
See DESIGN_SPEC Section 6.9.
"""

SAS = "sas"
CENTRALIZED = "centralized"
DECENTRALIZED = "decentralized"
CONTEXT_DEPENDENT = "context_dependent"
AUTO = "auto"


class ActionType(StrEnum):
"""Convenience constants for common approval action types.

Expand Down
17 changes: 16 additions & 1 deletion src/ai_company/core/task.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,14 @@
from pydantic import BaseModel, ConfigDict, Field, model_validator

from ai_company.core.artifact import ExpectedArtifact # noqa: TC001
from ai_company.core.enums import Complexity, Priority, TaskStatus, TaskType
from ai_company.core.enums import (
Complexity,
CoordinationTopology,
Priority,
TaskStatus,
TaskStructure,
TaskType,
)
from ai_company.core.task_transitions import validate_transition
from ai_company.core.types import NotBlankStr # noqa: TC001
from ai_company.observability import get_logger
Expand Down Expand Up @@ -133,6 +140,14 @@ class Task(BaseModel):
default=(),
description="Ordered agent IDs of delegators (root first)",
)
task_structure: TaskStructure | None = Field(
default=None,
description="Classification of subtask relationships (None = not classified)",
)
coordination_topology: CoordinationTopology = Field(
default=CoordinationTopology.AUTO,
description="Coordination topology for multi-agent execution",
)
Comment on lines +143 to +150

@coderabbitai coderabbitai Bot Mar 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Update class docstring to document new fields.

The Task class docstring (lines 52-73) lists all attributes but doesn't include the newly added task_structure and coordination_topology fields. Add them for consistency.

📝 Proposed docstring addition

Add these entries to the Attributes section in the class docstring:

         delegation_chain: Ordered agent IDs of delegators (root first).
+        task_structure: Classification of subtask relationships
+            (``None`` if not classified).
+        coordination_topology: Coordination topology for multi-agent
+            execution (default ``AUTO``).
     """
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/ai_company/core/task.py` around lines 143 - 150, The Task class docstring
is missing documentation for the two new fields; update the Attributes section
of the Task class docstring to include brief descriptions for task_structure
(TaskStructure | None — classification of subtask relationships; None means not
classified) and coordination_topology (CoordinationTopology — coordination
topology for multi-agent execution, default AUTO). Edit the Task class docstring
near the existing attribute list so the entries match the style/format of the
other attributes and reference the exact field names task_structure and
coordination_topology.


@model_validator(mode="after")
def _validate_deadline_format(self) -> Self:
Expand Down
48 changes: 48 additions & 0 deletions src/ai_company/engine/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,32 @@
AgentContext,
AgentContextSnapshot,
)
from ai_company.engine.decomposition import (
DecompositionContext,
DecompositionPlan,
DecompositionResult,
DecompositionService,
DecompositionStrategy,
DependencyGraph,
ManualDecompositionStrategy,
StatusRollup,
SubtaskDefinition,
SubtaskStatusRollup,
TaskStructureClassifier,
)
from ai_company.engine.errors import (
BudgetExhaustedError,
DecompositionCycleError,
DecompositionDepthError,
DecompositionError,
EngineError,
ExecutionStateError,
LoopExecutionError,
MaxTurnsExceededError,
ParallelExecutionError,
PromptBuildError,
ResourceConflictError,
TaskRoutingError,
)
from ai_company.engine.loop_protocol import (
BudgetChecker,
Expand Down Expand Up @@ -58,6 +75,15 @@
RecoveryStrategy,
)
from ai_company.engine.resource_lock import InMemoryResourceLock, ResourceLock
from ai_company.engine.routing import (
AgentTaskScorer,
AutoTopologyConfig,
RoutingCandidate,
RoutingDecision,
RoutingResult,
TaskRoutingService,
TopologySelector,
)
from ai_company.engine.run_result import AgentRunResult
from ai_company.engine.shutdown import (
CleanupCallback,
Expand All @@ -78,11 +104,22 @@
"AgentEngine",
"AgentOutcome",
"AgentRunResult",
"AgentTaskScorer",
"AutoTopologyConfig",
"BudgetChecker",
"BudgetExhaustedError",
"CleanupCallback",
"CooperativeTimeoutStrategy",
"DecompositionContext",
"DecompositionCycleError",
"DecompositionDepthError",
"DecompositionError",
"DecompositionPlan",
"DecompositionResult",
"DecompositionService",
"DecompositionStrategy",
"DefaultTokenEstimator",
"DependencyGraph",
"EngineError",
"ExecutionLoop",
"ExecutionPlan",
Expand All @@ -91,6 +128,7 @@
"FailAndReassignStrategy",
"InMemoryResourceLock",
"LoopExecutionError",
"ManualDecompositionStrategy",
"MaxTurnsExceededError",
"ParallelExecutionError",
"ParallelExecutionGroup",
Expand All @@ -108,16 +146,26 @@
"RecoveryStrategy",
"ResourceConflictError",
"ResourceLock",
"RoutingCandidate",
"RoutingDecision",
"RoutingResult",
"ShutdownChecker",
"ShutdownManager",
"ShutdownResult",
"ShutdownStrategy",
"StatusRollup",
"StatusTransition",
"StepStatus",
"SubtaskDefinition",
"SubtaskStatusRollup",
"SystemPrompt",
"TaskCompletionMetrics",
"TaskExecution",
"TaskRoutingError",
"TaskRoutingService",
"TaskStructureClassifier",
"TerminationReason",
"TopologySelector",
"TurnRecord",
"add_token_usage",
"build_system_prompt",
Expand Down
33 changes: 33 additions & 0 deletions src/ai_company/engine/decomposition/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
"""Task decomposition engine.

Breaks complex tasks into subtasks with dependency tracking,
classifies task structure, and manages status rollup.
"""

from ai_company.engine.decomposition.classifier import TaskStructureClassifier
from ai_company.engine.decomposition.dag import DependencyGraph
from ai_company.engine.decomposition.manual import ManualDecompositionStrategy
from ai_company.engine.decomposition.models import (
DecompositionContext,
DecompositionPlan,
DecompositionResult,
SubtaskDefinition,
SubtaskStatusRollup,
)
from ai_company.engine.decomposition.protocol import DecompositionStrategy
from ai_company.engine.decomposition.rollup import StatusRollup
from ai_company.engine.decomposition.service import DecompositionService

__all__ = [
"DecompositionContext",
"DecompositionPlan",
"DecompositionResult",
"DecompositionService",
"DecompositionStrategy",
"DependencyGraph",
"ManualDecompositionStrategy",
"StatusRollup",
"SubtaskDefinition",
"SubtaskStatusRollup",
"TaskStructureClassifier",
]
Loading