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..7e2c71de5 --- /dev/null +++ b/features/steps/decomposition_decision_correction_steps.py @@ -0,0 +1,351 @@ +"""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..82478bd39 100644 --- a/src/cleveragents/application/services/decomposition_models.py +++ b/src/cleveragents/application/services/decomposition_models.py @@ -130,6 +130,30 @@ class DecompositionResult: metrics: dict[str, int | float] = field(default_factory=dict) +@dataclass(frozen=True) +class DecisionCorrectionResult: + """Result of selective subtree recomputation for decision correction. + + Tracks which nodes were recomputed and which were preserved. + """ + + 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 +230,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..63860a57e 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, @@ -393,6 +394,49 @@ 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.""" + 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) + preserved_nodes = [ + n for n in existing_result.nodes if n.node_id not in subtree_ids + ] + target_node = node_map[node_id] + _reset_counter() + recomputed_result = self.decompose(list(target_node.file_paths), cfg) + metrics: dict[str, int | float] = { + "recomputed_count": len(recomputed_result.nodes), + "preserved_count": len(preserved_nodes), + "subtree_size": len(subtree_ids), + } + return DecisionCorrectionResult( + target_node_id=node_id, + recomputed_nodes=recomputed_result.nodes, + preserved_nodes=preserved_nodes, + config=cfg, + metrics=metrics, + ) + + __all__: list[str] = [ "DecompositionService", ]