cleveragents-core/benchmarks/decision_correction_bench.py

"""Airspeed Velocity benchmarks for decision correction flows.

Measures BFS subtree traversal, impact analysis, dry-run report
generation, and revert execution overhead across varying tree sizes.
"""

from __future__ import annotations

from cleveragents.application.services.correction_service import CorrectionService
from cleveragents.domain.models.core.correction import (
    CorrectionMode,
    CorrectionStatus,
)

# ---------------------------------------------------------------------------
# Tree generators
# ---------------------------------------------------------------------------

_PLAN_ID = "01HGZ6FE0AQDYTR4BXVQZ6EA00"


def _linear_tree(root: str, depth: int) -> dict[str, list[str]]:
    """Build a linear chain of *depth* nodes starting at *root*."""
    tree: dict[str, list[str]] = {}
    current = root
    for i in range(1, depth):
        child = f"{root}_c{i}"
        tree[current] = [child]
        current = child
    return tree


def _binary_tree(root: str, depth: int) -> dict[str, list[str]]:
    """Build a complete binary tree of *depth* levels rooted at *root*."""
    tree: dict[str, list[str]] = {}
    level = [root]
    counter = 0
    for _ in range(depth - 1):
        next_level: list[str] = []
        for node in level:
            left = f"{root}_n{counter}"
            right = f"{root}_n{counter + 1}"
            counter += 2
            tree[node] = [left, right]
            next_level.extend([left, right])
        level = next_level
    return tree


# ---------------------------------------------------------------------------
# BFS traversal benchmarks
# ---------------------------------------------------------------------------


class BfsTraversalSuite:
    """Benchmark BFS subtree computation at varying tree sizes."""

    def setup(self) -> None:
        """Prepare trees and service instances."""
        self.svc = CorrectionService()
        self.small_tree = _linear_tree("S", 5)  # 5 nodes
        self.medium_tree = _binary_tree("M", 5)  # 31 nodes
        self.large_tree = _binary_tree("L", 8)  # 255 nodes

    def time_bfs_small_tree(self) -> None:
        """BFS on a 5-node linear chain."""
        self.svc._compute_affected_subtree("S", self.small_tree)

    def time_bfs_medium_tree(self) -> None:
        """BFS on a 31-node binary tree (depth 5)."""
        self.svc._compute_affected_subtree("M", self.medium_tree)

    def time_bfs_large_tree(self) -> None:
        """BFS on a 255-node binary tree (depth 8)."""
        self.svc._compute_affected_subtree("L", self.large_tree)


# ---------------------------------------------------------------------------
# Impact analysis benchmarks
# ---------------------------------------------------------------------------


class ImpactAnalysisSuite:
    """Benchmark full impact analysis including risk classification."""

    def setup(self) -> None:
        """Prepare service and correction requests."""
        self.svc = CorrectionService()
        self.small_tree = _linear_tree("S", 3)
        self.medium_tree = _binary_tree("M", 4)
        self.large_tree = _binary_tree("L", 6)

        self.req_small = self.svc.request_correction(
            _PLAN_ID, "S", CorrectionMode.REVERT
        )
        self.req_medium = self.svc.request_correction(
            _PLAN_ID, "M", CorrectionMode.REVERT
        )
        self.req_large = self.svc.request_correction(
            _PLAN_ID, "L", CorrectionMode.REVERT
        )

    def time_impact_small(self) -> None:
        """Impact analysis on a 3-node chain."""
        # Reset status so analyze_impact can run again
        self.svc._corrections[
            self.req_small.correction_id
        ].status = CorrectionStatus.PENDING
        self.svc.analyze_impact(self.req_small.correction_id, self.small_tree)

    def time_impact_medium(self) -> None:
        """Impact analysis on a 15-node binary tree."""
        self.svc._corrections[
            self.req_medium.correction_id
        ].status = CorrectionStatus.PENDING
        self.svc.analyze_impact(self.req_medium.correction_id, self.medium_tree)

    def time_impact_large(self) -> None:
        """Impact analysis on a 63-node binary tree."""
        self.svc._corrections[
            self.req_large.correction_id
        ].status = CorrectionStatus.PENDING
        self.svc.analyze_impact(self.req_large.correction_id, self.large_tree)


# ---------------------------------------------------------------------------
# Dry-run report benchmarks
# ---------------------------------------------------------------------------


class DryRunReportSuite:
    """Benchmark dry-run report generation."""

    def setup(self) -> None:
        """Prepare service and requests."""
        self.svc = CorrectionService()
        self.tree = _binary_tree("DR", 5)
        self.req = self.svc.request_correction(_PLAN_ID, "DR", CorrectionMode.REVERT)

    def time_dry_run_report(self) -> None:
        """Generate dry-run report for a 31-node tree."""
        self.svc._corrections[self.req.correction_id].status = CorrectionStatus.PENDING
        self.svc.generate_dry_run_report(self.req.correction_id, self.tree)


# ---------------------------------------------------------------------------
# Revert execution benchmarks
# ---------------------------------------------------------------------------


class RevertExecutionSuite:
    """Benchmark revert execution overhead."""

    def _fresh_request(self, target: str) -> str:
        """Create a fresh correction request and return its ID."""
        req = self.svc.request_correction(_PLAN_ID, target, CorrectionMode.REVERT)
        return req.correction_id

    def setup(self) -> None:
        """Prepare service and trees."""
        self.svc = CorrectionService()
        self.small_tree = _linear_tree("RS", 3)
        self.medium_tree = _binary_tree("RM", 4)

    def time_revert_small(self) -> None:
        """Revert execution on a 3-node chain."""
        cid = self._fresh_request("RS")
        self.svc.execute_revert(cid, self.small_tree)

    def time_revert_medium(self) -> None:
        """Revert execution on a 15-node binary tree."""
        cid = self._fresh_request("RM")
        self.svc.execute_revert(cid, self.medium_tree)