cleveragents-core/benchmarks/acms_fusion_bench.py

"""ASV benchmarks for ACMS StrategyCoordinator and FusionEngine.

Measures the performance of:
- StrategyCoordinator coordination round (selection + allocation + execution)
- StrategyCoordinator with per-strategy max caps
- FusionEngine full fusion pipeline (dedup + depth + score + pack)
- FusionEngine with budget overage guard
- Integration: Coordinator → FusionEngine pipeline
"""

from __future__ import annotations

import importlib
import sys
from collections.abc import Sequence
from pathlib import Path
from typing import Any

_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

import cleveragents  # noqa: E402

importlib.reload(cleveragents)

from cleveragents.application.services.acms_service import (  # noqa: E402
    StrategyCapabilities,
)
from cleveragents.application.services.fusion_engine import (  # noqa: E402
    FusionConfig,
    FusionEngine,
)
from cleveragents.application.services.strategy_coordinator import (  # noqa: E402
    CoordinatorConfig,
    StrategyCoordinator,
)
from cleveragents.domain.models.core.context_fragment import (  # noqa: E402
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)

_DEFAULT_PROV = FragmentProvenance(resource_uri="bench://fusion")


def _make_frag(**kwargs: Any) -> ContextFragment:
    kwargs.setdefault("uko_node", "bench://fusion")
    kwargs.setdefault("token_count", 10)
    kwargs.setdefault("provenance", _DEFAULT_PROV)
    return ContextFragment(**kwargs)


class _BenchStrategy:
    """Minimal strategy for benchmarks."""

    def __init__(self, name: str = "bench", confidence: float = 0.7) -> None:
        self._name = name
        self._confidence = confidence

    @property
    def name(self) -> str:
        return self._name

    @property
    def capabilities(self) -> StrategyCapabilities:
        return StrategyCapabilities()

    def can_handle(self, request: dict[str, Any]) -> float:
        return self._confidence

    def assemble(
        self,
        fragments: Sequence[ContextFragment],
        budget: ContextBudget,
    ) -> Sequence[ContextFragment]:
        result: list[ContextFragment] = []
        total = 0
        for frag in fragments:
            if total + frag.token_count <= budget.available_tokens:
                result.append(frag)
                total += frag.token_count
        return result

    def explain(self) -> str:
        return "Benchmark strategy."


class StrategyCoordinatorSuite:
    """Benchmark StrategyCoordinator throughput."""

    def setup(self) -> None:
        self.coordinator = StrategyCoordinator()
        self.capped_coordinator = StrategyCoordinator(
            config=CoordinatorConfig(per_strategy_max_cap=500),
        )
        self.strategies = [
            _BenchStrategy("fast", 0.8),
            _BenchStrategy("slow", 0.4),
        ]
        self.fragments = [
            _make_frag(
                uko_node=f"bench://f{i}.py",
                content=f"content_{i}",
                token_count=10,
                relevance_score=0.5,
            )
            for i in range(20)
        ]
        self.budget = ContextBudget(max_tokens=500, reserved_tokens=0)

    def time_coordinate_2_strategies(self) -> None:
        """Coordinate 2 strategies with 20 fragments."""
        self.coordinator.coordinate(
            request={},
            strategies=self.strategies,
            budget=self.budget,
            fragments=self.fragments,
        )

    def time_coordinate_with_caps(self) -> None:
        """Coordinate with per-strategy max caps."""
        self.capped_coordinator.coordinate(
            request={},
            strategies=self.strategies,
            budget=self.budget,
            fragments=self.fragments,
        )


class FusionEngineSuite:
    """Benchmark FusionEngine throughput."""

    def setup(self) -> None:
        self.engine = FusionEngine(
            config=FusionConfig(min_fragment_tokens=1),
        )
        self.budget = ContextBudget(max_tokens=300, reserved_tokens=0)
        self.fragments_small = [
            _make_frag(
                uko_node=f"bench://s{i}.py",
                content=f"small_{i}",
                token_count=10,
                relevance_score=round(0.9 - i * 0.05, 2),
            )
            for i in range(10)
        ]
        self.fragments_with_dups = [
            _make_frag(
                uko_node="bench://dup.py",
                content="same content",
                token_count=10,
                relevance_score=0.9,
            ),
            _make_frag(
                uko_node="bench://dup.py",
                content="same content",
                token_count=10,
                relevance_score=0.7,
            ),
            *[
                _make_frag(
                    uko_node=f"bench://u{i}.py",
                    content=f"unique_{i}",
                    token_count=10,
                    relevance_score=round(0.8 - i * 0.05, 2),
                )
                for i in range(8)
            ],
        ]

    def time_fuse_10_fragments(self) -> None:
        """Fuse 10 unique fragments."""
        self.engine.fuse(self.fragments_small, self.budget)

    def time_fuse_with_duplicates(self) -> None:
        """Fuse fragments with duplicates requiring dedup."""
        self.engine.fuse(self.fragments_with_dups, self.budget)


class IntegrationSuite:
    """Benchmark end-to-end Coordinator → FusionEngine pipeline."""

    def setup(self) -> None:
        self.coordinator = StrategyCoordinator()
        self.engine = FusionEngine(
            config=FusionConfig(min_fragment_tokens=1),
        )
        self.strategies = [_BenchStrategy("main", 0.8)]
        self.fragments = [
            _make_frag(
                uko_node=f"bench://i{i}.py",
                content=f"integration_{i}",
                token_count=10,
                relevance_score=round(0.9 - i * 0.02, 2),
            )
            for i in range(30)
        ]
        self.budget = ContextBudget(max_tokens=200, reserved_tokens=0)

    def time_coordinate_then_fuse(self) -> None:
        """Full pipeline: coordinate then fuse 30 fragments."""
        coord = self.coordinator.coordinate(
            request={},
            strategies=self.strategies,
            budget=self.budget,
            fragments=self.fragments,
        )
        self.engine.fuse(coord.fragments, self.budget)