cleveragents-core/benchmarks/context_assembly_scaling_bench.py

"""ASV benchmarks for ACMS context assembly at production scale.

Extends the existing 1K-fragment ceiling from :mod:`acms_pipeline_bench`
to 5K and 10K fragments, exercising the full 10-stage pipeline,
deduplication + scoring, and knapsack budget packing.
"""

from __future__ import annotations

import importlib
import sys
import time
from pathlib import Path
from typing import ClassVar

# Ensure the local *source* tree is importable even when ASV has an
# older build of the package installed.
_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 ACMSPipeline  # noqa: E402
from cleveragents.domain.models.core.context_fragment import (  # noqa: E402
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)

# Default provenance for benchmark fragments.
_DEFAULT_PROV = FragmentProvenance(resource_uri="bench://scale")

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_fragments(count: int) -> list[ContextFragment]:
    """Build *count* benchmark fragments with varied scores and tiers."""
    tiers = ("hot", "warm", "cold")
    return [
        ContextFragment(
            uko_node=f"bench://scale/{i}",
            content=f"fragment content block {i} " * 5,
            relevance_score=round((i % 10) / 10, 1),
            token_count=50,
            tier=tiers[i % 3],
            provenance=_DEFAULT_PROV,
        )
        for i in range(count)
    ]


# ---------------------------------------------------------------------------
# Parameterized context assembly suite
# ---------------------------------------------------------------------------


class ContextAssemblyScalingSuite:
    """Benchmark ACMS context assembly at production fragment counts."""

    params: ClassVar[list[int]] = [100, 1_000, 5_000, 10_000]
    param_names: ClassVar[list[str]] = ["fragment_count"]
    timeout = 300  # 5 min for 10K fragments

    _pipeline: ACMSPipeline
    _budget: ContextBudget
    _fragments: list[ContextFragment]

    def setup(self, fragment_count: int) -> None:
        """Create *fragment_count* ContextFragment objects."""
        self._pipeline = ACMSPipeline()
        self._budget = ContextBudget(max_tokens=500_000, reserved_tokens=0)
        self._fragments = _make_fragments(fragment_count)

    # -- timing methods -----------------------------------------------------

    def time_full_pipeline(self, fragment_count: int) -> None:
        """Time the full ACMS pipeline with *fragment_count* fragments."""
        self._pipeline.assemble(
            plan_id="01JQBENCHM00000000000000AA",
            fragments=self._fragments,
            budget=self._budget,
        )

    def time_tiered_strategy(self, fragment_count: int) -> None:
        """Time tiered strategy assembly."""
        self._pipeline.assemble(
            plan_id="01JQBENCHM00000000000000AA",
            fragments=self._fragments,
            budget=self._budget,
            strategy="tiered",
        )

    def time_recency_strategy(self, fragment_count: int) -> None:
        """Time recency strategy assembly."""
        self._pipeline.assemble(
            plan_id="01JQBENCHM00000000000000AA",
            fragments=self._fragments,
            budget=self._budget,
            strategy="recency",
        )

    # -- tracking methods ---------------------------------------------------

    def track_assembled_tokens(self, fragment_count: int) -> int:
        """Track total tokens in assembled context."""
        result = self._pipeline.assemble(
            plan_id="01JQBENCHM00000000000000AA",
            fragments=self._fragments,
            budget=self._budget,
        )
        return result.total_tokens

    def track_fragments_per_second(self, fragment_count: int) -> float:
        """Track assembly throughput in fragments / second."""
        t0 = time.perf_counter()
        self._pipeline.assemble(
            plan_id="01JQBENCHM00000000000000AA",
            fragments=self._fragments,
            budget=self._budget,
        )
        elapsed = time.perf_counter() - t0
        if elapsed <= 0:
            return float(fragment_count)
        return fragment_count / elapsed


# Attach ASV unit metadata without ``# type: ignore``.
setattr(  # noqa: B010
    ContextAssemblyScalingSuite.track_assembled_tokens, "unit", "tokens"
)
setattr(  # noqa: B010
    ContextAssemblyScalingSuite.track_fragments_per_second, "unit", "frags/s"
)