cleveragents-core/robot/helper_skeleton_compressor.py

"""Helper script for skeleton compressor Robot Framework tests.

Usage:
    python helper_skeleton_compressor.py compress <budget>
    python helper_skeleton_compressor.py validate-budget-bounds
    python helper_skeleton_compressor.py stable-ordering
    python helper_skeleton_compressor.py protocol-check
"""

from __future__ import annotations

import sys
from pathlib import Path

# Ensure source tree is importable
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from cleveragents.application.services.acms_service import (  # noqa: E402
    SkeletonCompressor,
)
from cleveragents.application.services.skeleton_compressor import (  # noqa: E402
    SkeletonCompressorService,
)
from cleveragents.domain.models.core.context_fragment import (  # noqa: E402
    ContextFragment,
    FragmentProvenance,
)

# Default provenance for skeleton test fragments.
_SKEL_PROV = FragmentProvenance(resource_uri="skeleton://robot-test")


def _make_skel_fragment(
    fragment_id: str,
    content: str,
    token_count: int,
    relevance_score: float,
    source_decision_id: str | None = None,
) -> ContextFragment:
    """Create a ContextFragment suitable for skeleton compression tests."""
    metadata: dict[str, str] = {}
    if source_decision_id is not None:
        metadata["source_decision_id"] = source_decision_id
    return ContextFragment(
        fragment_id=fragment_id,
        uko_node=f"skeleton://{fragment_id}",
        content=content,
        token_count=token_count,
        relevance_score=relevance_score,
        provenance=_SKEL_PROV,
        metadata=metadata,
    )


def _sample_fragments() -> tuple[ContextFragment, ...]:
    """Build a repeatable set of sample fragments."""
    return (
        _make_skel_fragment(
            fragment_id="frag-001",
            content="High relevance content " * 20,
            token_count=400,
            relevance_score=0.9,
            source_decision_id="01HXDECISION00000000000001",
        ),
        _make_skel_fragment(
            fragment_id="frag-002",
            content="Medium relevance content " * 15,
            token_count=300,
            relevance_score=0.6,
            source_decision_id="01HXDECISION00000000000002",
        ),
        _make_skel_fragment(
            fragment_id="frag-003",
            content="Low relevance content " * 15,
            token_count=300,
            relevance_score=0.3,
            source_decision_id="01HXDECISION00000000000003",
        ),
    )


def cmd_compress(budget_str: str) -> None:
    """Compress sample fragments and print summary."""
    budget = int(budget_str)
    svc = SkeletonCompressorService()
    result = svc.compress(_sample_fragments(), budget)
    total_tokens = sum(f.token_count for f in result)
    print(f"skeleton-compress-ok budget={budget}")
    print(f"fragment_count={len(result)}")
    print(f"total_tokens={total_tokens}")


def cmd_validate_budget_bounds() -> None:
    """Verify that negative budgets raise ValueError."""
    svc = SkeletonCompressorService()
    frags = _sample_fragments()

    for bad in (-1, -10, -100):
        try:
            svc.compress(frags, bad)
            print(f"FAIL: budget {bad} did not raise")
            sys.exit(1)
        except ValueError:
            pass

    # Valid bounds
    for good in (0, 500, 1000):
        svc.compress(frags, good)

    print("skeleton-budget-bounds-ok")


def cmd_stable_ordering() -> None:
    """Verify fragments come out in deterministic order."""
    svc = SkeletonCompressorService()
    frags = tuple(
        _make_skel_fragment(
            fragment_id=f"frag-{chr(ord('c') - i)}",
            content="x",
            token_count=10,
            relevance_score=0.5,
        )
        for i in range(3)
    )

    r1 = svc.compress(frags, 10000)
    r2 = svc.compress(tuple(reversed(frags)), 10000)

    ids1 = [f.fragment_id for f in r1]
    ids2 = [f.fragment_id for f in r2]

    if ids1 == ids2:
        print("skeleton-stable-ordering-ok")
    else:
        print(f"FAIL: {ids1} != {ids2}")
        sys.exit(1)


def cmd_protocol_check() -> None:
    """Verify SkeletonCompressorService satisfies the SkeletonCompressor protocol."""
    svc = SkeletonCompressorService()
    if isinstance(svc, SkeletonCompressor):
        print("skeleton-protocol-check-ok")
    else:
        print(
            "FAIL: SkeletonCompressorService does not satisfy"
            " SkeletonCompressor protocol"
        )
        sys.exit(1)


def main() -> None:
    """Dispatch subcommand."""
    if len(sys.argv) < 2:
        print("Usage: helper_skeleton_compressor.py <command> [args]")
        sys.exit(1)

    cmd = sys.argv[1]
    if cmd == "compress":
        if len(sys.argv) < 3:
            print("Usage: compress <budget>")
            sys.exit(1)
        cmd_compress(sys.argv[2])
    elif cmd == "validate-budget-bounds":
        cmd_validate_budget_bounds()
    elif cmd == "stable-ordering":
        cmd_stable_ordering()
    elif cmd == "protocol-check":
        cmd_protocol_check()
    else:
        print(f"Unknown command: {cmd}")
        sys.exit(1)


if __name__ == "__main__":
    main()