cleveragents-core/benchmarks/acms_context_policy_bench.py

"""ASV benchmarks for ACMS context policy loader and plan execution integration.

Measures the performance of:
- ContextPolicyConfigurationLoader.load_from_string (YAML and TOML)
- ACMSContextAssembler.assemble_context with varying policy counts
- PlanExecutionACMSIntegration.prepare_llm_context with and without policies
- PolicyScope.matches with scalar and list values
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path

_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.acms.context_policy_loader import (  # noqa: E402
    ContextPolicyConfigurationLoader,
    PolicyScope,
)
from cleveragents.acms.plan_execution_integration import (  # noqa: E402
    ACMSContextAssembler,
    PlanExecutionACMSIntegration,
)

_YAML_CONFIG_1 = """
view_name: bench_view
policies:
  - name: policy1
    priority_weight: 1.0
    scopes:
      - name: file_type
        value: python
"""

_YAML_CONFIG_10 = """
view_name: bench_view
policies:
""" + "\n".join(f"  - name: policy{i}\n    priority_weight: {i}.0" for i in range(10))

_YAML_CONFIG_100 = """
view_name: bench_view
policies:
""" + "\n".join(f"  - name: policy{i}\n    priority_weight: {i}.0" for i in range(100))

_TOML_CONFIG_1 = """
view_name = "bench_view"
[[policies]]
name = "policy1"
priority_weight = 1.0
"""


class LoaderSuite:
    """Benchmark ContextPolicyConfigurationLoader throughput."""

    def setup(self) -> None:
        """Set up loader instance."""
        self._loader = ContextPolicyConfigurationLoader()

    def time_load_yaml_1_policy(self) -> None:
        """Benchmark loading a YAML config with 1 policy."""
        self._loader.load_from_string(_YAML_CONFIG_1, "yaml")

    def time_load_yaml_10_policies(self) -> None:
        """Benchmark loading a YAML config with 10 policies."""
        self._loader.load_from_string(_YAML_CONFIG_10, "yaml")

    def time_load_yaml_100_policies(self) -> None:
        """Benchmark loading a YAML config with 100 policies."""
        self._loader.load_from_string(_YAML_CONFIG_100, "yaml")

    def time_load_toml_1_policy(self) -> None:
        """Benchmark loading a TOML config with 1 policy."""
        self._loader.load_from_string(_TOML_CONFIG_1, "toml")


class AssemblerSuite:
    """Benchmark ACMSContextAssembler.assemble_context throughput."""

    def setup(self) -> None:
        """Set up assembler instances with varying policy counts."""
        loader = ContextPolicyConfigurationLoader()

        config_1 = loader.load_from_string(_YAML_CONFIG_1, "yaml")
        config_10 = loader.load_from_string(_YAML_CONFIG_10, "yaml")
        config_100 = loader.load_from_string(_YAML_CONFIG_100, "yaml")

        self._assembler_1 = ACMSContextAssembler(config_1)
        self._assembler_10 = ACMSContextAssembler(config_10)
        self._assembler_100 = ACMSContextAssembler(config_100)

        self._raw_context = {
            "file_type": "python",
            "path": "src/module.py",
            "content": "def hello(): pass",
        }

    def time_assemble_1_policy(self) -> None:
        """Benchmark assembling context with 1 policy."""
        self._assembler_1.assemble_context(self._raw_context)

    def time_assemble_10_policies(self) -> None:
        """Benchmark assembling context with 10 policies."""
        self._assembler_10.assemble_context(self._raw_context)

    def time_assemble_100_policies(self) -> None:
        """Benchmark assembling context with 100 policies."""
        self._assembler_100.assemble_context(self._raw_context)


class IntegrationSuite:
    """Benchmark PlanExecutionACMSIntegration.prepare_llm_context throughput."""

    def setup(self) -> None:
        """Set up integration instances."""
        loader = ContextPolicyConfigurationLoader()
        config = loader.load_from_string(_YAML_CONFIG_10, "yaml")

        self._integration_no_policy = PlanExecutionACMSIntegration()
        self._integration_with_policy = PlanExecutionACMSIntegration(
            policy_config=config
        )
        self._raw_context = {
            "file_type": "python",
            "path": "src/module.py",
        }

    def time_prepare_context_no_policy(self) -> None:
        """Benchmark prepare_llm_context with no policy (passthrough)."""
        self._integration_no_policy.prepare_llm_context(self._raw_context)

    def time_prepare_context_with_policy(self) -> None:
        """Benchmark prepare_llm_context with ACMS policy assembly."""
        self._integration_with_policy.prepare_llm_context(self._raw_context)


class PolicyScopeSuite:
    """Benchmark PolicyScope.matches throughput."""

    def setup(self) -> None:
        """Set up scope instances."""
        self._scope_scalar = PolicyScope(name="file_type", value="python")
        self._scope_list = PolicyScope(
            name="file_type", value=["python", "javascript", "typescript"]
        )
        self._context_match = {"file_type": "python"}
        self._context_no_match = {"file_type": "java"}
        self._context_list = {"file_type": ["python", "rust"]}

    def time_scope_scalar_match(self) -> None:
        """Benchmark scalar scope match."""
        self._scope_scalar.matches(self._context_match)

    def time_scope_scalar_no_match(self) -> None:
        """Benchmark scalar scope no-match."""
        self._scope_scalar.matches(self._context_no_match)

    def time_scope_list_match(self) -> None:
        """Benchmark list scope match."""
        self._scope_list.matches(self._context_match)

    def time_scope_context_list_match(self) -> None:
        """Benchmark scope match against list context value."""
        self._scope_scalar.matches(self._context_list)