cleveragents-core/robot/helper_depth_breadth_projection.py

"""Robot Framework helper library for depth/breadth projection tests."""

from __future__ import annotations

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

_PROV = FragmentProvenance(resource_uri="test://robot")
_GRAPH = {"A": ["B"], "B": ["C"]}


def _make_frag(
    uko_node: str = "test://node",
    content: str = "test",
    token_count: int = 100,
    detail_depth: int = 3,
) -> ContextFragment:
    return ContextFragment(
        uko_node=uko_node,
        content=content,
        token_count=token_count,
        relevance_score=0.8,
        detail_depth=detail_depth,
        provenance=_PROV,
    )


class _FakeContext:
    def __init__(self, frags: list[ContextFragment], ctx_hash: str = "h123") -> None:
        self.fragments = frags
        self.context_hash = ctx_hash


class _MockCompressor:
    def __init__(self, n: int = 2) -> None:
        self._n = n

    def compress(
        self,
        fragments: tuple[ContextFragment, ...],
        skeleton_budget: int,
    ) -> tuple[ContextFragment, ...]:
        return tuple(fragments[: self._n])


def create_projection_spec(focus: str, breadth: int, depth: int) -> bool:
    spec = ProjectionSpec(focus=(focus,), breadth=breadth, depth=depth)
    return spec.focus == (focus,) and spec.breadth == breadth


def create_projected_node(uri: str, distance: int, depth: int) -> bool:
    node = ProjectedNode(uri=uri, distance=distance, resolved_depth=depth)
    return node.uri == uri and node.distance == distance


def project_and_count_nodes(focus: str, breadth: int, depth: int) -> int:
    projector = DepthBreadthProjector()
    spec = ProjectionSpec(focus=(focus,), breadth=breadth, depth=depth)
    return len(projector.project(spec, _GRAPH))


def project_with_gradient(focus: str, breadth: int, depth: int) -> bool:
    projector = DepthBreadthProjector()
    spec = ProjectionSpec(
        focus=(focus,), breadth=breadth, depth=depth, depth_gradient=True
    )
    nodes = projector.project(spec, _GRAPH)
    by_uri = {n.uri: n for n in nodes}
    # Focus at full depth, farthest at 0
    return by_uri[focus].resolved_depth == depth and by_uri["C"].resolved_depth == 0


def project_to_fragments_count(focus: str, breadth: int, depth: int) -> int:
    projector = DepthBreadthProjector()
    spec = ProjectionSpec(focus=(focus,), breadth=breadth, depth=depth)
    return len(projector.project_to_fragments(spec, _GRAPH))


def resolve_code_level(level: str) -> int:
    return code_detail_map().resolve(level)


def compute_child_context_default(focus: str, budget: int) -> bool:
    frags = [_make_frag(uko_node=f"mod://m{i}") for i in range(3)]
    ctx = _FakeContext(frags)
    svc = PlanContextInheritance()
    result = svc.compute_child_context(
        parent_context=ctx, child_focus=[focus], child_token_budget=budget
    )
    return isinstance(result, ChildContextResult) and focus in result.request.focus


def compute_child_with_skeleton(focus: str, budget: int) -> int:
    frags = [_make_frag(uko_node=f"mod://m{i}") for i in range(3)]
    ctx = _FakeContext(frags)
    comp = _MockCompressor(n=2)
    svc = PlanContextInheritance(skeleton_compressor=comp)
    result = svc.compute_child_context(
        parent_context=ctx, child_focus=[focus], child_token_budget=budget
    )
    return len(result.skeleton_fragments)


def extract_focus(d1: str, d2: str) -> bool:
    result = PlanContextInheritance.extract_child_focus([d1, d2])
    return d1 in result and d2 in result