cleveragents-core/features/steps/skeleton_compressor_steps.py

"""Step definitions for the Skeleton Compressor feature."""

from __future__ import annotations

from behave import given, then, when
from behave.runner import Context
from pydantic import ValidationError

from cleveragents.application.services.skeleton_compressor import (
    DEFAULT_SKELETON_RATIO,
    SkeletonCompressorService,
)
from cleveragents.domain.models.core.context_fragment import (
    ContextFragment,
    FragmentProvenance,
)
from cleveragents.domain.models.core.plan import (
    NamespacedName,
    Plan,
    PlanIdentity,
    PlanPhase,
    ProcessingState,
)
from cleveragents.domain.models.core.skeleton_metadata import SkeletonMetadata

# --- helpers ---------------------------------------------------------------

# Default provenance for skeleton test fragments.
_SKEL_PROV = FragmentProvenance(resource_uri="skeleton://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 _make_fragments(total_tokens: int, count: int = 4) -> list[ContextFragment]:
    """Create *count* fragments summing to *total_tokens*."""
    base = total_tokens // count
    remainder = total_tokens - base * count
    relevances = [0.9, 0.7, 0.5, 0.3]
    frags: list[ContextFragment] = []
    for i in range(count):
        tokens = base + (remainder if i == 0 else 0)
        frags.append(
            _make_skel_fragment(
                fragment_id=f"frag-{i:03d}",
                content=f"content-{i}" * max(1, tokens // 10),
                token_count=tokens,
                relevance_score=relevances[i % len(relevances)],
                source_decision_id=f"01HX{'A' * 22}{i}" if i < 3 else None,
            )
        )
    return frags


def _make_plan_id() -> str:
    return "01HXAAAAAAAAAAAAAAAAAAAAAA"


# --- Background ------------------------------------------------------------


@given("a skeleton compressor service")
def step_create_service(context: Context) -> None:
    context.service = SkeletonCompressorService()
    context.compressor_error = None


# --- Fragment setup --------------------------------------------------------


@given("context fragments with total tokens {total:d}")
def step_fragments_total(context: Context, total: int) -> None:
    context.fragments = _make_fragments(total)


@given("three fragments with equal relevance {rel:g}")
def step_equal_relevance(context: Context, rel: float) -> None:
    context.fragments = [
        _make_skel_fragment(
            fragment_id=f"frag-{chr(ord('c') - i)}",
            content="x" * 50,
            token_count=100,
            relevance_score=rel,
        )
        for i in range(3)
    ]


@given("fragments with relevances 0.9, 0.3, and 0.7")
def step_varied_relevances(context: Context) -> None:
    context.fragments = [
        _make_skel_fragment(
            fragment_id="f-1", content="a", token_count=100, relevance_score=0.9
        ),
        _make_skel_fragment(
            fragment_id="f-2", content="b", token_count=100, relevance_score=0.3
        ),
        _make_skel_fragment(
            fragment_id="f-3", content="c", token_count=100, relevance_score=0.7
        ),
    ]


@given("fragments with known decision IDs")
def step_known_ids(context: Context) -> None:
    context.fragments = [
        _make_skel_fragment(
            fragment_id="f-1",
            content="a",
            token_count=100,
            relevance_score=0.9,
            source_decision_id="01HXDECISION00000000000001",
        ),
        _make_skel_fragment(
            fragment_id="f-2",
            content="b",
            token_count=100,
            relevance_score=0.5,
            source_decision_id="01HXDECISION00000000000002",
        ),
    ]


@given("an empty fragment list")
def step_empty_frags(context: Context) -> None:
    context.fragments = []


@given("a single fragment with {tokens:d} tokens")
def step_single_frag(context: Context, tokens: int) -> None:
    context.fragments = [
        _make_skel_fragment(
            fragment_id="only",
            content="x" * tokens,
            token_count=tokens,
            relevance_score=0.8,
        ),
    ]


@given("a fragment with negative token count")
def step_neg_tokens(context: Context) -> None:
    try:
        context.fragments = [
            _make_skel_fragment(
                fragment_id="bad",
                content="x",
                token_count=-10,
                relevance_score=0.5,
            ),
        ]
    except (ValueError, ValidationError) as exc:
        # Model-level validation now catches this at construction time.
        context.compressor_error = exc
        context.fragments = []


@given("a fragment with empty fragment_id")
def step_empty_id(context: Context) -> None:
    context.fragments = [
        ContextFragment(
            fragment_id="",
            uko_node="skeleton://empty-id",
            content="x",
            token_count=10,
            relevance_score=0.5,
            provenance=_SKEL_PROV,
        ),
    ]


@given("a fragment with relevance {rel:g}")
def step_bad_relevance(context: Context, rel: float) -> None:
    try:
        context.fragments = [
            _make_skel_fragment(
                fragment_id="bad",
                content="x",
                token_count=10,
                relevance_score=rel,
            ),
        ]
    except (ValueError, ValidationError) as exc:
        # Model-level validation now catches this at construction time.
        context.compressor_error = exc
        context.fragments = []


@given(
    "a skeleton metadata with ratio {ratio:g} and {orig:d} original tokens and {comp:d} compressed"
)
def step_make_metadata(context: Context, ratio: float, orig: int, comp: int) -> None:
    context.skel_meta = SkeletonMetadata(
        ratio=ratio,
        original_tokens=orig,
        compressed_tokens=comp,
        source_decision_ids=("01HXDECISION00000000000001",),
    )


# --- When clauses ----------------------------------------------------------


@when("I compress with skeleton_ratio {ratio:g}")
def step_compress_ratio(context: Context, ratio: float) -> None:
    # If model-level validation already caught an error during fragment
    # construction (Given step), propagate it as the compress error.
    if getattr(context, "compressor_error", None) is not None:
        context.error = context.compressor_error
        context.result = None
        return
    try:
        context.result = context.service.compress(
            context.fragments, skeleton_ratio=ratio
        )
        context.error = None
    except (ValueError, TypeError) as exc:
        context.error = exc
        context.result = None


@when("I compress with skeleton_ratio not specified")
def step_compress_default(context: Context) -> None:
    context.result = context.service.compress(context.fragments)
    context.error = None


@when("I compress with a non-list fragments argument")
def step_compress_non_list(context: Context) -> None:
    try:
        context.result = context.service.compress("not-a-list", skeleton_ratio=0.5)  # type: ignore[arg-type]
        context.error = None
    except TypeError as exc:
        context.error = exc
        context.result = None


@when("I compress with a list containing a non-fragment item")
def step_compress_non_fragment_item(context: Context) -> None:
    try:
        context.result = context.service.compress(
            [
                _make_skel_fragment(
                    fragment_id="ok",
                    content="x",
                    token_count=10,
                    relevance_score=0.5,
                ),
                "not-a-fragment",
            ],  # type: ignore[list-item]
            skeleton_ratio=0.5,
        )
        context.error = None
    except TypeError as exc:
        context.error = exc
        context.result = None


@when("I compress with a non-numeric skeleton_ratio")
def step_compress_non_numeric_ratio(context: Context) -> None:
    try:
        context.result = context.service.compress(
            context.fragments,
            skeleton_ratio="bad",  # type: ignore[arg-type]
        )
        context.error = None
    except TypeError as exc:
        context.error = exc
        context.result = None


@when("I create skeleton metadata with compressed exceeding original")
def step_create_bad_metadata(context: Context) -> None:
    try:
        context.bad_meta = SkeletonMetadata(
            ratio=0.5,
            original_tokens=100,
            compressed_tokens=200,
        )
        context.meta_error = None
    except Exception as exc:
        context.meta_error = exc
        context.bad_meta = None


@when("I attach skeleton_metadata to a plan")
def step_attach_to_plan(context: Context) -> None:
    context.plan = Plan(
        identity=PlanIdentity(plan_id=_make_plan_id()),
        namespaced_name=NamespacedName.parse("local/test-plan"),
        description="Test plan",
        action_name="local/test-action",
        phase=PlanPhase.STRATEGIZE,
        processing_state=ProcessingState.QUEUED,
        skeleton_metadata=context.skel_meta,
    )


# --- Then clauses ----------------------------------------------------------


@then("the compressor should raise a ValueError for invalid ratio")
def step_check_value_error_ratio(context: Context) -> None:
    assert context.error is not None, "Expected ValueError"
    assert isinstance(context.error, ValueError), (
        f"Expected ValueError, got {type(context.error)}"
    )


@then("all fragments should be returned unchanged")
def step_all_returned(context: Context) -> None:
    assert context.result is not None
    assert len(context.result.fragments) == len(context.fragments)


@then("only the highest-relevance fragment should be returned")
def step_top_one(context: Context) -> None:
    assert context.result is not None
    assert len(context.result.fragments) == 1
    top_relevance = max(f.relevance_score for f in context.fragments)
    assert context.result.fragments[0].relevance_score == top_relevance


@then("compressed tokens should be at most {limit:d}")
def step_tokens_limit(context: Context, limit: int) -> None:
    assert context.result is not None
    assert context.result.metadata.compressed_tokens <= limit


@then("the metadata ratio should equal the default {expected:g}")
def step_default_ratio(context: Context, expected: float) -> None:
    assert context.result is not None
    assert context.result.metadata.ratio == expected
    assert expected == DEFAULT_SKELETON_RATIO


@then("fragments should be ordered by fragment_id ascending")
def step_ordered_by_id(context: Context) -> None:
    assert context.result is not None
    ids = [f.fragment_id for f in context.result.fragments]
    assert ids == sorted(ids), f"Expected sorted IDs, got {ids}"


@then("the first fragment should have relevance {rel:g}")
def step_first_relevance(context: Context, rel: float) -> None:
    assert context.result is not None
    assert context.result.fragments[0].relevance_score == rel


@then("the last fragment should have relevance {rel:g}")
def step_last_relevance(context: Context, rel: float) -> None:
    assert context.result is not None
    assert context.result.fragments[-1].relevance_score == rel


@then("metadata original_tokens should be {expected:d}")
def step_original_tokens(context: Context, expected: int) -> None:
    assert context.result is not None
    assert context.result.metadata.original_tokens == expected


@then("metadata compressed_tokens should equal {expected:d}")
def step_compressed_equals(context: Context, expected: int) -> None:
    assert context.result is not None
    assert context.result.metadata.compressed_tokens == expected


@then("metadata compressed_tokens should be at most {limit:d}")
def step_compressed_at_most(context: Context, limit: int) -> None:
    assert context.result is not None
    assert context.result.metadata.compressed_tokens <= limit


@then("metadata should contain all source decision IDs")
def step_all_decision_ids(context: Context) -> None:
    assert context.result is not None
    expected_ids = {
        f.metadata["source_decision_id"]
        for f in context.fragments
        if "source_decision_id" in f.metadata
    }
    actual_ids = set(context.result.metadata.source_decision_ids)
    assert expected_ids == actual_ids


@then("metadata ratio should be {expected:g}")
def step_ratio_value(context: Context, expected: float) -> None:
    assert context.result is not None
    assert context.result.metadata.ratio == expected


@then("the result should contain zero fragments")
def step_zero_frags(context: Context) -> None:
    assert context.result is not None
    assert len(context.result.fragments) == 0


@then("the result should contain one fragment")
def step_one_frag(context: Context) -> None:
    assert context.result is not None
    assert len(context.result.fragments) == 1


@then("the compressor should raise a TypeError")
def step_type_error(context: Context) -> None:
    assert context.error is not None, "Expected TypeError"
    assert isinstance(context.error, TypeError), (
        f"Expected TypeError, got {type(context.error)}"
    )


@then("the compressor should raise a ValueError for invalid fragment")
def step_frag_value_error(context: Context) -> None:
    assert context.error is not None, "Expected ValueError"
    assert isinstance(context.error, ValueError), (
        f"Expected ValueError, got {type(context.error)}"
    )


@then("the compressor should raise a TypeError for invalid item")
def step_type_error_item(context: Context) -> None:
    assert context.error is not None, "Expected TypeError"
    assert isinstance(context.error, TypeError), (
        f"Expected TypeError, got {type(context.error)}"
    )


@then("the compressor should raise a TypeError for invalid ratio type")
def step_type_error_ratio_type(context: Context) -> None:
    assert context.error is not None, "Expected TypeError"
    assert isinstance(context.error, TypeError), (
        f"Expected TypeError, got {type(context.error)}"
    )


@then("a validation error should be raised for compressed exceeding original")
def step_validation_error_compressed(context: Context) -> None:
    assert context.meta_error is not None, "Expected validation error"


@then("compressed_tokens should be less than original_tokens")
def step_less_tokens(context: Context) -> None:
    assert context.result is not None
    assert (
        context.result.metadata.compressed_tokens
        < context.result.metadata.original_tokens
    )


@then("the plan should expose skeleton metadata in cli dict")
def step_plan_cli_dict(context: Context) -> None:
    cli_dict = context.plan.as_cli_dict()
    assert "skeleton" in cli_dict
    skel = cli_dict["skeleton"]
    assert skel["ratio"] == context.skel_meta.ratio
    assert skel["original_tokens"] == context.skel_meta.original_tokens
    assert skel["compressed_tokens"] == context.skel_meta.compressed_tokens


# --- _validate_fragments edge-case steps --------------------------------


@given("a context fragment constructed with negative token_count")
def step_frag_negative_token(context: Context) -> None:
    frag = ContextFragment.model_construct(
        fragment_id="neg-tc",
        uko_node="skeleton://neg-tc",
        content="x",
        token_count=-5,
        relevance_score=0.5,
        provenance=_SKEL_PROV,
    )
    context.invalid_fragments = [frag]


@given("a context fragment constructed with relevance_score {score}")
def step_frag_bad_relevance(context: Context, score: str) -> None:
    frag = ContextFragment.model_construct(
        fragment_id="bad-rel",
        uko_node="skeleton://bad-rel",
        content="x",
        token_count=10,
        relevance_score=float(score),
        provenance=_SKEL_PROV,
    )
    context.invalid_fragments = [frag]


@given("a context fragment constructed with empty fragment_id")
def step_frag_empty_id(context: Context) -> None:
    frag = ContextFragment.model_construct(
        fragment_id="",
        uko_node="skeleton://empty",
        content="x",
        token_count=10,
        relevance_score=0.5,
        provenance=_SKEL_PROV,
    )
    context.invalid_fragments = [frag]


@when("I validate the invalid fragments")
def step_validate_invalid(context: Context) -> None:
    try:
        SkeletonCompressorService._validate_fragments(context.invalid_fragments)
        context.validation_error = None
    except (TypeError, ValueError) as exc:
        context.validation_error = exc


@then('the compressor should raise a ValueError mentioning "{text}"')
def step_check_validation_error(context: Context, text: str) -> None:
    assert context.validation_error is not None, (
        "Expected a ValueError but none was raised"
    )
    assert isinstance(context.validation_error, ValueError), (
        f"Expected ValueError, got {type(context.validation_error).__name__}"
    )
    assert text in str(context.validation_error), (
        f"Expected '{text}' in error message: {context.validation_error}"
    )