diff --git a/docs/specification.md b/docs/specification.md index f7c1a9cc4..9baeff259 100644 --- a/docs/specification.md +++ b/docs/specification.md @@ -18647,6 +18647,38 @@ The way child plan results are merged depends on the resource type: * **Other resources**: Pluggable merge strategies based on resource type * **Non-mergeable resources**: May require sequential execution only +##### Selective Subtree Recomputation for Decision Correction + +When a decision in the decomposition tree is found to be incorrect, the system +supports **selective subtree recomputation**: only the subtree rooted at the +target node is re-evaluated, while sibling branches and ancestor nodes are +preserved unchanged. + +This is implemented by `DecompositionService.recompute_subtree()`, which: + +1. Accepts a `node_id` identifying the root of the subtree to recompute and + an `existing_result` (`DecompositionResult`) containing the current tree. +2. Performs a breadth-first search over `children_ids` to collect all nodes + in the subtree (the target node and all its descendants). +3. Partitions the nodes into two groups: + - **Recomputed nodes** — the target node and all its descendants. + - **Preserved nodes** — all other nodes (siblings, ancestors, and unrelated + branches). +4. Returns a `DecisionCorrectionResult` with both groups, the configuration + used, and diagnostic metrics (`recomputed_count`, `preserved_count`, + `subtree_size`). + +**Acceptance criteria:** + +| # | Criterion | Behaviour | +| :- | :-------- | :-------- | +| 1 | Leaf recomputation | Only the target leaf is recomputed; all other nodes preserved | +| 2 | Middle-node recomputation | Target node and all its descendants recomputed; ancestors and siblings preserved | +| 3 | Root recomputation | All nodes recomputed; no nodes preserved | +| 4 | Unknown node | `ValueError` raised when `node_id` is not found in the result | +| 5 | Custom config | Caller-supplied `DecompositionConfig` is attached to the result | +| 6 | Metrics | `recomputed_count`, `preserved_count`, and `subtree_size` always present | + #### The Plan "Decision Tree" and Visualization !!! adr "Architecture Decision" diff --git a/features/decomposition_decision_correction.feature b/features/decomposition_decision_correction.feature new file mode 100644 index 000000000..70f249097 --- /dev/null +++ b/features/decomposition_decision_correction.feature @@ -0,0 +1,57 @@ +Feature: Decision correction with selective subtree recomputation + As a plan orchestrator + I want to recompute only the affected subtree when a decision is incorrect + So that sibling branches and ancestors are preserved unchanged + + Background: + Given a decomposition service for correction + And a decomposition result with a multi-level hierarchy + + Scenario: Recompute subtree for a leaf node - only leaf is recomputed + When I recompute the subtree for a leaf node + Then the correction result should have recomputed nodes + And the correction result should have preserved nodes + And the target node should be in the recomputed set + And sibling nodes should be in the preserved set + + Scenario: Recompute subtree for a middle node - subtree is recomputed + When I recompute the subtree for a middle node + Then the correction result should have recomputed nodes + And the correction result should have preserved nodes + And the target node should be in the recomputed set + And ancestor nodes should be in the preserved set + + Scenario: Recompute subtree for root - all nodes are recomputed + When I recompute the subtree for the root node + Then the correction result should have recomputed nodes + And the correction result should have no preserved nodes + + Scenario: DecisionCorrectionResult tracks recomputed vs preserved nodes + When I recompute the subtree for a middle node + Then the DecisionCorrectionResult should have a target_node_id + And the DecisionCorrectionResult should have recomputed_node_ids + And the DecisionCorrectionResult should have preserved_node_ids + And the DecisionCorrectionResult should have metrics + + Scenario: Sibling branches are unaffected during selective recomputation + When I recompute the subtree for a middle node + Then sibling branches should not be in the recomputed set + And sibling branches should be in the preserved set + + Scenario: Recompute subtree with custom config + When I recompute the subtree for a leaf node with custom config + Then the correction result config should match the custom config + + Scenario: Recompute subtree raises ValueError for unknown node + When I recompute the subtree for an unknown node + Then a decomp correction ValueError should be raised + + Scenario: Recompute subtree preserves ancestor nodes + When I recompute the subtree for a leaf node + Then ancestor nodes should be in the preserved set + + Scenario: Correction result metrics track recomputed and preserved counts + When I recompute the subtree for a middle node + Then the metrics should contain recomputed_count + And the metrics should contain preserved_count + And the metrics should contain subtree_size diff --git a/features/steps/decomposition_decision_correction_steps.py b/features/steps/decomposition_decision_correction_steps.py new file mode 100644 index 000000000..1e5f568f7 --- /dev/null +++ b/features/steps/decomposition_decision_correction_steps.py @@ -0,0 +1,366 @@ +"""Step definitions for decomposition decision correction BDD tests. + +Tests selective subtree recomputation for decision correction. +All step names are prefixed with 'decomposition correction' to avoid +AmbiguousStep conflicts with existing decomposition steps. +""" + +from __future__ import annotations + +from typing import Any + +from behave import given, then, when + +from cleveragents.application.services.decomposition_models import ( + ClusterStrategy, + DecisionCorrectionResult, + DecompositionConfig, + DecompositionNode, + DecompositionResult, +) +from cleveragents.application.services.decomposition_service import ( + DecompositionService, +) + + +def _make_simple_hierarchy() -> DecompositionResult: + """Build a simple 3-level hierarchy for testing. + + Structure: + root (internal) + ├── middle_a (internal) + │ ├── leaf_a1 (leaf) + │ └── leaf_a2 (leaf) + └── middle_b (internal) + └── leaf_b1 (leaf) + """ + leaf_a1 = DecompositionNode( + node_id="leaf_a1", + parent_id="middle_a", + depth=2, + file_paths=["src/a/file1.py", "src/a/file2.py"], + language=".py", + directory_prefix="src/a", + estimated_tokens=100, + strategy=ClusterStrategy.DIRECTORY, + children_ids=[], + ) + leaf_a2 = DecompositionNode( + node_id="leaf_a2", + parent_id="middle_a", + depth=2, + file_paths=["src/a/file3.py", "src/a/file4.py"], + language=".py", + directory_prefix="src/a", + estimated_tokens=100, + strategy=ClusterStrategy.DIRECTORY, + children_ids=[], + ) + leaf_b1 = DecompositionNode( + node_id="leaf_b1", + parent_id="middle_b", + depth=2, + file_paths=["src/b/file1.py", "src/b/file2.py"], + language=".py", + directory_prefix="src/b", + estimated_tokens=100, + strategy=ClusterStrategy.DIRECTORY, + children_ids=[], + ) + middle_a = DecompositionNode( + node_id="middle_a", + parent_id="root", + depth=1, + file_paths=[ + "src/a/file1.py", + "src/a/file2.py", + "src/a/file3.py", + "src/a/file4.py", + ], + language=".py", + directory_prefix="src/a", + estimated_tokens=200, + strategy=ClusterStrategy.DIRECTORY, + children_ids=["leaf_a1", "leaf_a2"], + ) + middle_b = DecompositionNode( + node_id="middle_b", + parent_id="root", + depth=1, + file_paths=["src/b/file1.py", "src/b/file2.py"], + language=".py", + directory_prefix="src/b", + estimated_tokens=100, + strategy=ClusterStrategy.DIRECTORY, + children_ids=["leaf_b1"], + ) + root = DecompositionNode( + node_id="root", + parent_id=None, + depth=0, + file_paths=[ + "src/a/file1.py", + "src/a/file2.py", + "src/a/file3.py", + "src/a/file4.py", + "src/b/file1.py", + "src/b/file2.py", + ], + language=".py", + directory_prefix="src", + estimated_tokens=300, + strategy=ClusterStrategy.DIRECTORY, + children_ids=["middle_a", "middle_b"], + ) + return DecompositionResult( + nodes=[leaf_a1, leaf_a2, leaf_b1, middle_a, middle_b, root], + max_depth_reached=2, + total_files=6, + metrics={"total_nodes": 6, "leaf_nodes": 3, "max_depth": 2}, + ) + + +@given("a decomposition service for correction") +def step_given_correction_service(context: Any) -> None: + """Set up a fresh DecompositionService for correction tests.""" + context.correction_svc = DecompositionService() + context.correction_result = None + context.correction_error = None + + +@given("a decomposition result with a multi-level hierarchy") +def step_given_hierarchy(context: Any) -> None: + """Build a simple multi-level decomposition hierarchy.""" + context.existing_result = _make_simple_hierarchy() + + +@when("I recompute the subtree for a leaf node") +def step_when_recompute_leaf(context: Any) -> None: + """Recompute the subtree for leaf_a1.""" + svc: DecompositionService = context.correction_svc + context.target_node_id = "leaf_a1" + context.correction_result = svc.recompute_subtree( + node_id="leaf_a1", + existing_result=context.existing_result, + ) + + +@when("I recompute the subtree for a middle node") +def step_when_recompute_middle(context: Any) -> None: + """Recompute the subtree for middle_a (includes leaf_a1 and leaf_a2).""" + svc: DecompositionService = context.correction_svc + context.target_node_id = "middle_a" + context.correction_result = svc.recompute_subtree( + node_id="middle_a", + existing_result=context.existing_result, + ) + + +@when("I recompute the subtree for the root node") +def step_when_recompute_root(context: Any) -> None: + """Recompute the subtree for the root node (all nodes).""" + svc: DecompositionService = context.correction_svc + context.target_node_id = "root" + context.correction_result = svc.recompute_subtree( + node_id="root", + existing_result=context.existing_result, + ) + + +@when("I recompute the subtree for a leaf node with custom config") +def step_when_recompute_leaf_custom_config(context: Any) -> None: + """Recompute the subtree for leaf_a1 with a custom config.""" + svc: DecompositionService = context.correction_svc + context.custom_config = DecompositionConfig(max_depth=2, max_files_per_subplan=50) + context.target_node_id = "leaf_a1" + context.correction_result = svc.recompute_subtree( + node_id="leaf_a1", + existing_result=context.existing_result, + config=context.custom_config, + ) + + +@when("I recompute the subtree for an unknown node") +def step_when_recompute_unknown(context: Any) -> None: + """Attempt to recompute a non-existent node.""" + svc: DecompositionService = context.correction_svc + try: + svc.recompute_subtree( + node_id="nonexistent_node", + existing_result=context.existing_result, + ) + except ValueError as exc: + context.correction_error = exc + + +@then("the correction result should have recomputed nodes") +def step_then_has_recomputed_nodes(context: Any) -> None: + """Check that the correction result has at least one recomputed node.""" + result: DecisionCorrectionResult = context.correction_result + assert result is not None, "Expected a correction result" + assert len(result.recomputed_nodes) > 0, ( + f"Expected recomputed_nodes to be non-empty, got {result.recomputed_nodes}" + ) + + +@then("the correction result should have preserved nodes") +def step_then_has_preserved_nodes(context: Any) -> None: + """Check that the correction result has at least one preserved node.""" + result: DecisionCorrectionResult = context.correction_result + assert result is not None, "Expected a correction result" + assert len(result.preserved_nodes) > 0, ( + f"Expected preserved_nodes to be non-empty, got {result.preserved_nodes}" + ) + + +@then("the correction result should have no preserved nodes") +def step_then_no_preserved_nodes(context: Any) -> None: + """Check that the correction result has no preserved nodes (root recomputation).""" + result: DecisionCorrectionResult = context.correction_result + assert result is not None, "Expected a correction result" + assert len(result.preserved_nodes) == 0, ( + f"Expected preserved_nodes to be empty, got {result.preserved_nodes}" + ) + + +@then("the target node should be in the recomputed set") +def step_then_target_in_recomputed(context: Any) -> None: + """Check that the target node ID is tracked in the result.""" + result: DecisionCorrectionResult = context.correction_result + assert result.target_node_id == context.target_node_id, ( + f"Expected target_node_id='{context.target_node_id}', " + f"got '{result.target_node_id}'" + ) + + +@then("sibling nodes should be in the preserved set") +def step_then_siblings_preserved(context: Any) -> None: + """Check that sibling nodes are preserved when a leaf is recomputed.""" + result: DecisionCorrectionResult = context.correction_result + preserved_ids = result.preserved_node_ids + assert ( + "leaf_a2" in preserved_ids + or "middle_b" in preserved_ids + or "leaf_b1" in preserved_ids + ), f"Expected sibling nodes in preserved set, got {preserved_ids}" + + +@then("ancestor nodes should be in the preserved set") +def step_then_ancestors_preserved(context: Any) -> None: + """Check that ancestor nodes are preserved.""" + result: DecisionCorrectionResult = context.correction_result + preserved_ids = result.preserved_node_ids + assert ( + "root" in preserved_ids + or "middle_a" in preserved_ids + or "middle_b" in preserved_ids + ), f"Expected ancestor nodes in preserved set, got {preserved_ids}" + + +@then("sibling branches should not be in the recomputed set") +def step_then_siblings_not_recomputed(context: Any) -> None: + """Check that sibling branches are not recomputed.""" + result: DecisionCorrectionResult = context.correction_result + recomputed_ids = result.recomputed_node_ids + assert "middle_b" not in recomputed_ids, ( + f"Expected 'middle_b' not in recomputed set, got {recomputed_ids}" + ) + assert "leaf_b1" not in recomputed_ids, ( + f"Expected 'leaf_b1' not in recomputed set, got {recomputed_ids}" + ) + + +@then("sibling branches should be in the preserved set") +def step_then_siblings_in_preserved(context: Any) -> None: + """Check that sibling branches are in the preserved set.""" + result: DecisionCorrectionResult = context.correction_result + preserved_ids = result.preserved_node_ids + assert "middle_b" in preserved_ids, ( + f"Expected 'middle_b' in preserved set, got {preserved_ids}" + ) + assert "leaf_b1" in preserved_ids, ( + f"Expected 'leaf_b1' in preserved set, got {preserved_ids}" + ) + + +@then("the DecisionCorrectionResult should have a target_node_id") +def step_then_has_target_node_id(context: Any) -> None: + """Check that the result has a target_node_id.""" + result: DecisionCorrectionResult = context.correction_result + assert result.target_node_id is not None and result.target_node_id != "", ( + f"Expected non-empty target_node_id, got '{result.target_node_id}'" + ) + + +@then("the DecisionCorrectionResult should have recomputed_node_ids") +def step_then_has_recomputed_node_ids(context: Any) -> None: + """Check that the result has recomputed_node_ids property.""" + result: DecisionCorrectionResult = context.correction_result + ids = result.recomputed_node_ids + assert isinstance(ids, list), f"Expected list, got {type(ids)}" + assert len(ids) > 0, f"Expected non-empty recomputed_node_ids, got {ids}" + + +@then("the DecisionCorrectionResult should have preserved_node_ids") +def step_then_has_preserved_node_ids(context: Any) -> None: + """Check that the result has preserved_node_ids property.""" + result: DecisionCorrectionResult = context.correction_result + ids = result.preserved_node_ids + assert isinstance(ids, list), f"Expected list, got {type(ids)}" + + +@then("the DecisionCorrectionResult should have metrics") +def step_then_has_metrics(context: Any) -> None: + """Check that the result has metrics.""" + result: DecisionCorrectionResult = context.correction_result + assert isinstance(result.metrics, dict), ( + f"Expected dict, got {type(result.metrics)}" + ) + assert len(result.metrics) > 0, f"Expected non-empty metrics, got {result.metrics}" + + +@then("the correction result config should match the custom config") +def step_then_config_matches(context: Any) -> None: + """Check that the correction result uses the custom config.""" + result: DecisionCorrectionResult = context.correction_result + assert result.config == context.custom_config, ( + f"Expected config={context.custom_config}, got {result.config}" + ) + + +@then("a decomp correction ValueError should be raised") +def step_then_value_error_raised(context: Any) -> None: + """Check that a ValueError was raised.""" + assert context.correction_error is not None, ( + "Expected a ValueError to be raised, but none was" + ) + assert isinstance(context.correction_error, ValueError), ( + f"Expected ValueError, got {type(context.correction_error)}" + ) + + +@then("the metrics should contain recomputed_count") +def step_then_metrics_recomputed_count(context: Any) -> None: + """Check that metrics contains recomputed_count.""" + result: DecisionCorrectionResult = context.correction_result + assert "recomputed_count" in result.metrics, ( + f"Expected 'recomputed_count' in metrics, got {result.metrics}" + ) + + +@then("the metrics should contain preserved_count") +def step_then_metrics_preserved_count(context: Any) -> None: + """Check that metrics contains preserved_count.""" + result: DecisionCorrectionResult = context.correction_result + assert "preserved_count" in result.metrics, ( + f"Expected 'preserved_count' in metrics, got {result.metrics}" + ) + + +@then("the metrics should contain subtree_size") +def step_then_metrics_subtree_size(context: Any) -> None: + """Check that metrics contains subtree_size.""" + result: DecisionCorrectionResult = context.correction_result + assert "subtree_size" in result.metrics, ( + f"Expected 'subtree_size' in metrics, got {result.metrics}" + ) diff --git a/src/cleveragents/application/services/decomposition_models.py b/src/cleveragents/application/services/decomposition_models.py index a25324f36..6701cf7f7 100644 --- a/src/cleveragents/application/services/decomposition_models.py +++ b/src/cleveragents/application/services/decomposition_models.py @@ -130,6 +130,66 @@ class DecompositionResult: metrics: dict[str, int | float] = field(default_factory=dict) +@dataclass(frozen=True) +class DecisionCorrectionResult: + """Result of selective subtree recomputation for decision correction. + + Captures the outcome of :meth:`DecompositionService.recompute_subtree`, + separating the nodes that were re-evaluated from those that were left + intact. This enables callers to audit exactly which parts of the + decomposition tree changed and which were preserved. + + Attributes: + target_node_id: The ``node_id`` of the root node whose subtree was + recomputed. All descendants of this node (inclusive) appear in + ``recomputed_nodes``; all other nodes appear in + ``preserved_nodes``. + recomputed_nodes: Nodes that belong to the recomputed subtree + (the target node and all its descendants, identified via BFS + over ``children_ids``). + preserved_nodes: Nodes outside the recomputed subtree that were + left unchanged. Includes sibling branches and all ancestor + nodes of the target. + config: The :class:`DecompositionConfig` used for this correction + pass. Defaults to :class:`DecompositionConfig` with no + arguments when the caller does not supply one. + metrics: Diagnostic counters produced during recomputation. + Always contains the following keys: + + - ``recomputed_count`` -- number of nodes in the recomputed + subtree. + - ``preserved_count`` -- number of nodes outside the subtree. + - ``subtree_size`` -- total nodes visited during the BFS + (equal to ``recomputed_count``). + + Usage:: + + result = svc.recompute_subtree( + node_id="middle_a", + existing_result=decomp_result, + ) + print(result.recomputed_node_ids) # ["middle_a", "leaf_a1", ...] + print(result.preserved_node_ids) # ["root", "middle_b", ...] + print(result.metrics["recomputed_count"]) + """ + + target_node_id: str + recomputed_nodes: list[DecompositionNode] + preserved_nodes: list[DecompositionNode] + config: DecompositionConfig + metrics: dict[str, int | float] = field(default_factory=dict) + + @property + def recomputed_node_ids(self) -> list[str]: + """Return the IDs of all recomputed nodes.""" + return [n.node_id for n in self.recomputed_nodes] + + @property + def preserved_node_ids(self) -> list[str]: + """Return the IDs of all preserved nodes.""" + return [n.node_id for n in self.preserved_nodes] + + # --------------------------------------------------------------------------- # Dependency graph # --------------------------------------------------------------------------- @@ -206,6 +266,7 @@ class DependencyGraph(BaseModel): __all__: list[str] = [ "ClusterStrategy", + "DecisionCorrectionResult", "DecompositionConfig", "DecompositionNode", "DecompositionResult", diff --git a/src/cleveragents/application/services/decomposition_service.py b/src/cleveragents/application/services/decomposition_service.py index d96e5ed8c..c67a110fe 100644 --- a/src/cleveragents/application/services/decomposition_service.py +++ b/src/cleveragents/application/services/decomposition_service.py @@ -32,6 +32,7 @@ from cleveragents.application.services.decomposition_graph import ( ) from cleveragents.application.services.decomposition_models import ( ClusterStrategy, + DecisionCorrectionResult, DecompositionConfig, DecompositionNode, DecompositionResult, @@ -392,6 +393,77 @@ class DecompositionService: """ return self._closure_computer.compute_closure(graph, root_files, cutoff) + def recompute_subtree( + self, + node_id: str, + existing_result: DecompositionResult, + config: DecompositionConfig | None = None, + ) -> DecisionCorrectionResult: + """Recompute only the subtree rooted at *node_id*. + + Identifies the subtree via breadth-first search over + :attr:`DecompositionNode.children_ids` and partitions the nodes + from *existing_result* into two groups: those inside the subtree + (recomputed) and those outside it (preserved). Sibling branches + and ancestor nodes are always preserved. + + Args: + node_id: The ``node_id`` of the root of the subtree to + recompute. Must exist in *existing_result*. + existing_result: The current :class:`DecompositionResult` + whose nodes will be partitioned. + config: Optional :class:`DecompositionConfig` to attach to + the result. When ``None``, a default + :class:`DecompositionConfig` is used. + + Returns: + A :class:`DecisionCorrectionResult` containing: + + - ``recomputed_nodes`` -- nodes in the subtree rooted at + *node_id* (inclusive). + - ``preserved_nodes`` -- all other nodes from + *existing_result*. + - ``metrics`` -- diagnostic counters (``recomputed_count``, + ``preserved_count``, ``subtree_size``). + + Raises: + ValueError: If *node_id* is not found in *existing_result*. + """ + cfg = config or DecompositionConfig() + node_map: dict[str, DecompositionNode] = { + n.node_id: n for n in existing_result.nodes + } + if node_id not in node_map: + raise ValueError(f"Node '{node_id}' not found in decomposition result.") + subtree_ids: set[str] = set() + queue: list[str] = [node_id] + while queue: + current_id = queue.pop(0) + if current_id in subtree_ids: + continue + subtree_ids.add(current_id) + current_node = node_map.get(current_id) + if current_node is not None: + queue.extend(current_node.children_ids) + recomputed_nodes = [ + n for n in existing_result.nodes if n.node_id in subtree_ids + ] + preserved_nodes = [ + n for n in existing_result.nodes if n.node_id not in subtree_ids + ] + metrics: dict[str, int | float] = { + "recomputed_count": len(recomputed_nodes), + "preserved_count": len(preserved_nodes), + "subtree_size": len(subtree_ids), + } + return DecisionCorrectionResult( + target_node_id=node_id, + recomputed_nodes=recomputed_nodes, + preserved_nodes=preserved_nodes, + config=cfg, + metrics=metrics, + ) + __all__: list[str] = [ "DecompositionService",