cleveragents-core/benchmarks/depth_breadth_projection_bench.py

"""ASV benchmarks for depth/breadth projection system.

Measures latency of projection, gradient computation, and
skeleton context inheritance.
"""

from __future__ import annotations

from cleveragents.application.services.depth_breadth_projection import (
    DepthBreadthProjector,
    PlanContextInheritance,
    ProjectedNode,
    ProjectionSpec,
    code_detail_map,
)
from cleveragents.domain.models.acms.crp import (
    ContextFragment,
    FragmentProvenance,
)

_PROV = FragmentProvenance(resource_uri="bench://projection")


def _make_frag(i: int, depth: int = 3) -> ContextFragment:
    return ContextFragment(
        uko_node=f"bench://node_{i}",
        content=f"node_{i}",
        token_count=100,
        relevance_score=0.8,
        detail_depth=depth,
        provenance=_PROV,
    )


class _FakeContext:
    def __init__(self, n: int = 10) -> None:
        self.fragments = [_make_frag(i) for i in range(n)]
        self.context_hash = "benchhash"


def _build_linear_graph(n: int) -> dict[str, list[str]]:
    """Build a linear graph: node_0 -> node_1 -> ... -> node_(n-1)."""
    adj: dict[str, list[str]] = {}
    for i in range(n - 1):
        adj[f"node_{i}"] = [f"node_{i + 1}"]
    return adj


class TimeProjectionSpec:
    """Benchmark ProjectionSpec creation."""

    def time_create_spec(self) -> None:
        ProjectionSpec(focus=("class://Auth",), breadth=2, depth=9)

    def time_create_spec_named_depth(self) -> None:
        ProjectionSpec(focus=("class://Auth",), depth="FULL_SOURCE")


class TimeProjectedNode:
    """Benchmark ProjectedNode creation."""

    def time_create_node(self) -> None:
        ProjectedNode(uri="class://Foo", distance=1, resolved_depth=4)


class TimeProjector:
    """Benchmark DepthBreadthProjector."""

    def setup(self) -> None:
        self.projector = DepthBreadthProjector()
        self.projector.register_detail_map("uko-code:", code_detail_map())
        self.graph_10 = _build_linear_graph(10)
        self.graph_100 = _build_linear_graph(100)

    def time_project_small_graph(self) -> None:
        spec = ProjectionSpec(focus=("node_0",), breadth=5, depth=9)
        self.projector.project(spec, self.graph_10)

    def time_project_medium_graph(self) -> None:
        spec = ProjectionSpec(focus=("node_0",), breadth=10, depth=9)
        self.projector.project(spec, self.graph_100)

    def time_project_no_gradient(self) -> None:
        spec = ProjectionSpec(
            focus=("node_0",), breadth=5, depth=9, depth_gradient=False
        )
        self.projector.project(spec, self.graph_10)

    def time_project_to_fragments(self) -> None:
        spec = ProjectionSpec(focus=("node_0",), breadth=5, depth=9)
        self.projector.project_to_fragments(spec, self.graph_10)

    def time_project_named_depth(self) -> None:
        spec = ProjectionSpec(
            focus=("node_0",),
            breadth=3,
            depth="FULL_SOURCE",
            domain="uko-code:",
        )
        self.projector.project(spec, self.graph_10)


class TimeInheritance:
    """Benchmark PlanContextInheritance."""

    def setup(self) -> None:
        self.service = PlanContextInheritance()
        self.parent_ctx = _FakeContext(n=10)

    def time_compute_child_context(self) -> None:
        self.service.compute_child_context(
            parent_context=self.parent_ctx,
            child_focus=["class://Child"],
            child_token_budget=4096,
        )

    def time_extract_focus(self) -> None:
        PlanContextInheritance.extract_child_focus(
            ["class://A", "class://B", "class://C"]
        )