cleveragents-core/benchmarks/tool_wrapping_bench.py

"""ASV benchmarks for tool wrapping delegation overhead.

Measures the performance of:
- ArgumentMapper.apply() with identity and configured mappings
- TransformExecutor.execute() with simple transform functions
- WrappedToolExecutor.execute() for single and chained delegation
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path
from typing import Any

# 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.domain.models.core.tool import (  # noqa: E402
    Tool,
    ToolCapability,
    ToolSource,
    ToolType,
    Validation,
    ValidationMode,
)
from cleveragents.tool.wrapping import (  # noqa: E402
    ArgumentMapper,
    TransformExecutor,
    WrappedToolExecutor,
)

_SIMPLE_TRANSFORM = """\
def transform(tool_output):
    passed = tool_output.get("returncode") == 0
    return {"passed": passed, "message": "done", "data": tool_output}
"""

_PASSTHROUGH_TRANSFORM = """\
def transform(tool_output):
    return {"passed": True, "message": "ok", "data": tool_output}
"""


def _make_tool(name: str) -> Tool:
    return Tool(
        name=name,
        description=f"Benchmark tool {name}",
        source=ToolSource.BUILTIN,
        capability=ToolCapability(read_only=True),
    )


def _make_validation(
    name: str,
    wraps: str,
    transform_code: str,
    argument_mapping: dict[str, Any] | None = None,
) -> Validation:
    return Validation(
        name=name,
        description=f"Benchmark validation wrapping {wraps}",
        source=ToolSource.WRAPPED,
        tool_type=ToolType.VALIDATION,
        mode=ValidationMode.REQUIRED,
        wraps=wraps,
        transform=transform_code,
        argument_mapping=argument_mapping,
    )


class ArgumentMapperSuite:
    """Benchmarks for ArgumentMapper."""

    def time_identity_mapping(self) -> None:
        """Measure identity (passthrough) mapping overhead."""
        mapper = ArgumentMapper(None)
        for _ in range(1000):
            mapper.apply({"path": "/src", "verbose": True, "count": 42})

    def time_configured_mapping(self) -> None:
        """Measure mapping with configured argument translation."""
        mapper = ArgumentMapper(
            {
                "test_directory": "source_dir",
                "coverage_enabled": True,
                "verbose": False,
            }
        )
        for _ in range(1000):
            mapper.apply({"source_dir": "/tests", "extra": "ignored"})


class TransformExecutorSuite:
    """Benchmarks for TransformExecutor."""

    def setup(self) -> None:
        self._executor = TransformExecutor(_SIMPLE_TRANSFORM, "bench/tool")
        self._output = {"returncode": 0, "tests_run": 100}

    def time_execute_transform(self) -> None:
        """Measure transform function execution."""
        for _ in range(1000):
            self._executor.execute(self._output)

    def time_create_and_execute(self) -> None:
        """Measure creation + execution (cold path)."""
        for _ in range(100):
            executor = TransformExecutor(_SIMPLE_TRANSFORM, "bench/tool")
            executor.execute(self._output)


class WrappedToolExecutorSuite:
    """Benchmarks for WrappedToolExecutor delegation."""

    def setup(self) -> None:
        self._tool = _make_tool("local/base-tool")
        self._output: dict[str, Any] = {"returncode": 0, "data": "ok"}
        self._validation = _make_validation(
            "local/bench-wrap",
            "local/base-tool",
            _SIMPLE_TRANSFORM,
        )

        tools: dict[str, tuple[Tool, dict[str, Any]]] = {
            "local/base-tool": (self._tool, self._output),
        }

        def lookup(name: str) -> Tool | Validation | None:
            if name in tools:
                return tools[name][0]
            return None

        def executor(name: str, args: dict[str, Any]) -> dict[str, Any]:
            if name in tools:
                return tools[name][1]
            raise RuntimeError(f"Not found: {name}")

        self._executor = WrappedToolExecutor(lookup, executor)

    def time_single_delegation(self) -> None:
        """Measure single-level wrapping delegation."""
        for _ in range(1000):
            self._executor.execute(self._validation, {"path": "/src"})

    def time_chained_delegation(self) -> None:
        """Measure two-level wrapping delegation chain."""
        inner_val = _make_validation(
            "local/inner-wrap",
            "local/base-tool",
            _PASSTHROUGH_TRANSFORM,
        )
        outer_val = _make_validation(
            "local/outer-wrap",
            "local/inner-wrap",
            _SIMPLE_TRANSFORM,
        )

        tools: dict[str, tuple[Tool, dict[str, Any]]] = {
            "local/base-tool": (self._tool, self._output),
        }

        all_items: dict[str, Tool | Validation] = {
            "local/inner-wrap": inner_val,
            "local/base-tool": self._tool,
        }

        def lookup(name: str) -> Tool | Validation | None:
            return all_items.get(name)

        def executor(name: str, args: dict[str, Any]) -> dict[str, Any]:
            if name in tools:
                return tools[name][1]
            raise RuntimeError(f"Not found: {name}")

        chained_executor = WrappedToolExecutor(lookup, executor)
        for _ in range(500):
            chained_executor.execute(outer_val, {"path": "/src"})