cleveragents-core/robot/helper_acms_pipeline.py

"""Robot Framework helper for ACMS v1 pipeline integration tests.

Provides a CLI-style interface for Robot to invoke ACMS pipeline
operations and verify the results.  Exit code 0 = success, 1 = failure.

Usage:
    python robot/helper_acms_pipeline.py fragment-create
    python robot/helper_acms_pipeline.py budget-calc
    python robot/helper_acms_pipeline.py assemble-relevance
    python robot/helper_acms_pipeline.py assemble-recency
    python robot/helper_acms_pipeline.py assemble-tiered
    python robot/helper_acms_pipeline.py payload-budget-check
    python robot/helper_acms_pipeline.py assemble-size-budget
"""

from __future__ import annotations

import sys
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path

_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC in sys.path:
    sys.path.remove(_SRC)
sys.path.insert(0, _SRC)

from cleveragents.application.services.acms_service import ACMSPipeline  # noqa: E402
from cleveragents.domain.models.core.context_fragment import (  # noqa: E402
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)
from cleveragents.domain.models.core.context_policy import ContextView  # noqa: E402

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


def _cmd_fragment_create() -> int:
    """Create a ContextFragment and verify defaults."""
    frag = ContextFragment(
        uko_node="project://app/main.py",
        content="hello world",
        token_count=10,
        provenance=FragmentProvenance(resource_uri="project://app/main.py"),
    )
    if frag.uko_node != "project://app/main.py":
        print(
            f"acms-fail: expected uko_node=project://app/main.py, got {frag.uko_node}"
        )
        return 1
    if frag.relevance_score != 0.5:
        print(f"acms-fail: expected score=0.5, got {frag.relevance_score}")
        return 1
    if not frag.fragment_id:
        print("acms-fail: fragment_id is empty")
        return 1
    if not isinstance(frag.created_at, datetime):
        print(f"acms-fail: created_at is not datetime, got {type(frag.created_at)}")
        return 1
    if frag.token_count != 10:
        print(f"acms-fail: expected token_count=10, got {frag.token_count}")
        return 1
    print(f"acms-fragment-ok: id={frag.fragment_id}")
    return 0


def _cmd_budget_calc() -> int:
    """Create ContextBudget and verify available_tokens."""
    budget = ContextBudget(max_tokens=4096, reserved_tokens=512)
    expected = 3584
    if budget.available_tokens != expected:
        print(
            f"acms-fail: expected available={expected}, got {budget.available_tokens}"
        )
        return 1
    print(f"acms-budget-ok: available={budget.available_tokens}")
    return 0


def _cmd_assemble_relevance() -> int:
    """Assemble fragments with relevance strategy."""
    frags = [
        ContextFragment(
            uko_node="project://app/high.py",
            content="high",
            relevance_score=0.9,
            token_count=100,
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/low.py",
            content="low",
            relevance_score=0.2,
            token_count=100,
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/mid.py",
            content="mid",
            relevance_score=0.6,
            token_count=100,
            provenance=_DEFAULT_PROV,
        ),
    ]
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    pipeline = ACMSPipeline()
    payload = pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="relevance",
    )
    if len(payload.fragments) != 2:
        print(f"acms-fail: expected 2 fragments, got {len(payload.fragments)}")
        return 1
    if payload.fragments[0].content != "high":
        print(f"acms-fail: expected first=high, got {payload.fragments[0].content}")
        return 1
    if not payload.strategies_used:
        print("acms-fail: strategies_used is empty")
        return 1
    print(f"acms-assemble-relevance-ok: fragments={len(payload.fragments)}")
    return 0


def _cmd_assemble_recency() -> int:
    """Assemble fragments with recency strategy."""
    frags = [
        ContextFragment(
            uko_node="project://app/old.py",
            content="old",
            token_count=100,
            created_at=datetime(2024, 1, 1, tzinfo=UTC),
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/mid.py",
            content="mid",
            token_count=100,
            created_at=datetime(2024, 6, 1, tzinfo=UTC),
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/new.py",
            content="new",
            token_count=100,
            created_at=datetime(2025, 1, 1, tzinfo=UTC),
            provenance=_DEFAULT_PROV,
        ),
    ]
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    pipeline = ACMSPipeline()
    payload = pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="recency",
    )
    if len(payload.fragments) != 2:
        print(f"acms-fail: expected 2 fragments, got {len(payload.fragments)}")
        return 1
    if payload.fragments[0].content != "new":
        print(f"acms-fail: expected first=new, got {payload.fragments[0].content}")
        return 1
    if payload.fragments[1].content != "mid":
        print(f"acms-fail: expected second=mid, got {payload.fragments[1].content}")
        return 1
    print(f"acms-assemble-recency-ok: fragments={len(payload.fragments)}")
    return 0


def _cmd_assemble_tiered() -> int:
    """Assemble fragments with tiered strategy."""
    frags = [
        ContextFragment(
            uko_node="project://app/cold.py",
            content="cold-item",
            relevance_score=0.9,
            token_count=100,
            tier="cold",
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/hot.py",
            content="hot-item",
            relevance_score=0.85,
            token_count=100,
            tier="hot",
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/warm.py",
            content="warm-item",
            relevance_score=0.8,
            token_count=100,
            tier="warm",
            provenance=_DEFAULT_PROV,
        ),
    ]
    budget = ContextBudget(max_tokens=200, reserved_tokens=0)
    pipeline = ACMSPipeline()
    payload = pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="tiered",
    )
    if len(payload.fragments) != 2:
        print(f"acms-fail: expected 2 fragments, got {len(payload.fragments)}")
        return 1
    if payload.fragments[0].content != "hot-item":
        print(f"acms-fail: expected first=hot-item, got {payload.fragments[0].content}")
        return 1
    print(f"acms-assemble-tiered-ok: fragments={len(payload.fragments)}")
    return 0


def _cmd_payload_budget_check() -> int:
    """Verify is_within_budget property on assembled payload."""
    frags = [
        ContextFragment(
            uko_node="project://app/item.py",
            content="item",
            relevance_score=0.9,
            token_count=100,
            provenance=_DEFAULT_PROV,
        ),
    ]
    budget = ContextBudget(max_tokens=500, reserved_tokens=0)
    pipeline = ACMSPipeline()
    payload = pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="relevance",
    )
    if not payload.is_within_budget:
        print("acms-fail: expected is_within_budget=True")
        return 1
    if payload.remaining_tokens != 400:
        print(f"acms-fail: expected remaining=400, got {payload.remaining_tokens}")
        return 1
    if not payload.context_hash:
        print("acms-fail: expected non-empty context_hash")
        return 1
    if len(payload.context_hash) != 64:
        print(
            f"acms-fail: expected SHA-256 hex (64 chars), "
            f"got {len(payload.context_hash)}"
        )
        return 1
    try:
        int(payload.context_hash, 16)
    except ValueError:
        print(f"acms-fail: context_hash is not valid hex: {payload.context_hash!r}")
        return 1
    print(
        f"acms-payload-budget-ok: within_budget={payload.is_within_budget} "
        f"remaining={payload.remaining_tokens}"
    )
    return 0


def _cmd_assemble_size_budget() -> int:
    """Verify ACMSPipeline budget pre-filter with ContextView limits."""
    frags = [
        ContextFragment(
            uko_node="project://app/small.py",
            content="a" * 50,
            relevance_score=0.9,
            token_count=60,
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/too_big.py",
            content="b" * 200,
            relevance_score=0.8,
            token_count=60,
            provenance=_DEFAULT_PROV,
        ),
        ContextFragment(
            uko_node="project://app/fit.py",
            content="c" * 80,
            relevance_score=0.7,
            token_count=60,
            provenance=_DEFAULT_PROV,
        ),
    ]
    budget = ContextBudget(max_tokens=1000, reserved_tokens=0)
    pipeline = ACMSPipeline()
    payload = pipeline.assemble(
        plan_id="01JQTESTPN00000000000000AA",
        fragments=frags,
        budget=budget,
        strategy="relevance",
        context_view=ContextView(max_file_size=100),
    )

    result = pipeline.last_enforcement_result
    if result is None:
        print("acms-fail: expected non-empty enforcement result")
        return 1
    if len(payload.fragments) != 2:
        print(f"acms-fail: expected 2 fragments, got {len(payload.fragments)}")
        return 1
    if len(result.accepted) != 2:
        print(f"acms-fail: expected 2 accepted ids, got {len(result.accepted)}")
        return 1
    if len(result.violations) != 1:
        print(f"acms-fail: expected 1 violation, got {len(result.violations)}")
        return 1
    if result.violations[0].violation_type != "max_file_size":
        print(
            "acms-fail: expected violation_type=max_file_size, "
            f"got {result.violations[0].violation_type}"
        )
        return 1

    payload_ids = {fragment.fragment_id for fragment in payload.fragments}
    if payload_ids != set(result.accepted):
        print(
            "acms-fail: payload fragment ids do not match accepted ids "
            f"payload={sorted(payload_ids)} accepted={sorted(result.accepted)}"
        )
        return 1

    if result.total_size != 130:
        print(f"acms-fail: expected total_size=130, got {result.total_size}")
        return 1

    print("acms-assemble-size-budget-ok: accepted=2 violations=1 total_size=130")
    return 0


_COMMANDS: dict[str, Callable[[], int]] = {
    "fragment-create": _cmd_fragment_create,
    "budget-calc": _cmd_budget_calc,
    "assemble-relevance": _cmd_assemble_relevance,
    "assemble-recency": _cmd_assemble_recency,
    "assemble-tiered": _cmd_assemble_tiered,
    "payload-budget-check": _cmd_payload_budget_check,
    "assemble-size-budget": _cmd_assemble_size_budget,
}


def main() -> int:
    """Entry point called by Robot Framework ``Run Process``."""
    if len(sys.argv) < 2:
        print(
            "Usage: helper_acms_pipeline.py "
            "<fragment-create|budget-calc|assemble-relevance"
            "|assemble-recency|assemble-tiered|payload-budget-check"
            "|assemble-size-budget>"
        )
        return 1

    command = sys.argv[1]
    handler = _COMMANDS.get(command)
    if handler is None:
        print(f"Unknown command: {command}")
        return 1

    return handler()


if __name__ == "__main__":
    sys.exit(main())