cleveragents-core/features/steps/acms_pipeline_phase2_steps.py

"""Step definitions for features/acms_pipeline_phase2.feature.

Tests the ACMS pipeline Phase 2 (Fragment Fusion) components directly
in-memory — no database required.
"""

from __future__ import annotations

from typing import Any

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

from cleveragents.application.services.acms_phase2 import (
    ContentHashDeduplicator,
    GreedyKnapsackPacker,
    MaxDepthResolver,
    ScorerWeights,
    WeightedCompositeScorer,
)
from cleveragents.application.services.acms_service import ACMSPipeline
from cleveragents.domain.models.core.context_fragment import (
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
    ScoredFragment,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_DEFAULT_PROVENANCE = FragmentProvenance(resource_uri="test://phase2")


def _make_phase2_fragment(**kwargs: Any) -> ContextFragment:
    """Create a ContextFragment with sensible Phase 2 test defaults."""
    kwargs.setdefault("uko_node", "test://default")
    kwargs.setdefault("token_count", 0)
    kwargs.setdefault("provenance", _DEFAULT_PROVENANCE)
    return ContextFragment(**kwargs)


# ---------------------------------------------------------------------------
# ScoredFragment — Given / When / Then
# ---------------------------------------------------------------------------


@given(
    'a context fragment with uko_node "{uko}" and content "{content}" '
    "and token_count {tokens:d} for phase2"
)
def step_given_fragment_for_phase2(
    context: Context, uko: str, content: str, tokens: int
) -> None:
    context.phase2_fragment = _make_phase2_fragment(
        uko_node=uko, content=content, token_count=tokens
    )


@when("I create a ScoredFragment with composite_score {score:g} and components:")
def step_create_scored_fragment(context: Context, score: float) -> None:
    components: dict[str, float] = {}
    for row in context.table:
        components[row["component"]] = float(row["value"])
    context.scored_fragment = ScoredFragment(
        fragment=context.phase2_fragment,
        composite_score=score,
        score_components=components,
    )


@when("I create a ScoredFragment with invalid composite_score {score:g}")
def step_create_scored_fragment_invalid(context: Context, score: float) -> None:
    context.phase2_error = None
    try:
        ScoredFragment(
            fragment=context.phase2_fragment,
            composite_score=score,
            score_components={},
        )
    except (ValidationError, ValueError) as exc:
        context.phase2_error = exc


@then("the ScoredFragment composite_score should be {expected:g}")
def step_scored_fragment_composite(context: Context, expected: float) -> None:
    assert context.scored_fragment.composite_score == expected


@then('the ScoredFragment should have component "{name}" with value {val:g}')
def step_scored_fragment_component(context: Context, name: str, val: float) -> None:
    assert context.scored_fragment.score_components[name] == val


@then("the ScoredFragment should reference the original fragment")
def step_scored_fragment_ref(context: Context) -> None:
    assert context.scored_fragment.fragment is context.phase2_fragment


@then("a phase2 validation error should be raised")
def step_phase2_validation_error(context: Context) -> None:
    assert context.phase2_error is not None


# ---------------------------------------------------------------------------
# Fragment list construction — Given steps
# ---------------------------------------------------------------------------


@given("the following phase2 fragments:")
def step_given_phase2_fragments(context: Context) -> None:
    frags: list[ContextFragment] = []
    for row in context.table:
        frags.append(
            _make_phase2_fragment(
                uko_node=row["uko_node"],
                content=row["content"],
                relevance_score=float(row["score"]),
                token_count=int(row["tokens"]),
                detail_depth=int(row["depth"]),
            )
        )
    context.phase2_fragments = frags


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


@given("a phase2 fragment with metadata:")
def step_given_phase2_fragment_with_metadata(context: Context) -> None:
    row = context.table[0]
    meta: dict[str, str] = {}
    if "hierarchy_weight" in row.headings:
        meta["hierarchy_weight"] = row["hierarchy_weight"]
    if "strategy_quality" in row.headings:
        meta["strategy_quality"] = row["strategy_quality"]
    if "recency_bonus" in row.headings:
        meta["recency_bonus"] = row["recency_bonus"]

    context.phase2_fragments = [
        _make_phase2_fragment(
            uko_node=row["uko_node"],
            content=row["content"],
            relevance_score=float(row["score"]),
            token_count=int(row["tokens"]),
            detail_depth=int(row["depth"]),
            metadata=meta,
        )
    ]


@given("a phase2 budget with max_tokens {max_t:d} and reserved_tokens {res_t:d}")
def step_given_phase2_budget(context: Context, max_t: int, res_t: int) -> None:
    context.phase2_budget = ContextBudget(max_tokens=max_t, reserved_tokens=res_t)


@given(
    "scorer weights: relevance={rel:g}, hierarchy={hier:g}, "
    "quality={qual:g}, recency={rec:g}"
)
def step_given_scorer_weights(
    context: Context, rel: float, hier: float, qual: float, rec: float
) -> None:
    context.scorer_weights = ScorerWeights(
        relevance=rel, hierarchy=hier, quality=qual, recency=rec
    )


# ---------------------------------------------------------------------------
# ContentHashDeduplicator — When / Then
# ---------------------------------------------------------------------------


@when("I deduplicate the fragments")
def step_deduplicate(context: Context) -> None:
    dedup = ContentHashDeduplicator()
    context.phase2_result = list(dedup.deduplicate(context.phase2_fragments))


@then("{count:d} fragment should remain after deduplication")
def step_dedup_count_singular(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count, (
        f"Expected {count}, got {len(context.phase2_result)}"
    )


@then("{count:d} fragments should remain after deduplication")
def step_dedup_count_plural(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count, (
        f"Expected {count}, got {len(context.phase2_result)}"
    )


@then("the surviving fragment should have relevance score {expected:g}")
def step_surviving_relevance(context: Context, expected: float) -> None:
    assert len(context.phase2_result) >= 1
    assert context.phase2_result[0].relevance_score == expected


# ---------------------------------------------------------------------------
# MaxDepthResolver — When / Then
# ---------------------------------------------------------------------------


@when("I resolve depth conflicts")
def step_resolve_depth(context: Context) -> None:
    resolver = MaxDepthResolver()
    context.phase2_result = list(resolver.resolve(context.phase2_fragments))


@then("{count:d} fragment should remain after depth resolution")
def step_depth_count_singular(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count, (
        f"Expected {count}, got {len(context.phase2_result)}"
    )


@then("{count:d} fragments should remain after depth resolution")
def step_depth_count_plural(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count, (
        f"Expected {count}, got {len(context.phase2_result)}"
    )


@then("the surviving fragment should have detail depth {expected:d}")
def step_surviving_depth(context: Context, expected: int) -> None:
    assert len(context.phase2_result) >= 1
    assert context.phase2_result[0].detail_depth == expected


@then('the resolved fragment order should be "{c1}", "{c2}", "{c3}"')
def step_resolved_order(context: Context, c1: str, c2: str, c3: str) -> None:
    contents = [f.content for f in context.phase2_result]
    assert contents == [c1, c2, c3], f"Expected [{c1}, {c2}, {c3}], got {contents}"


# ---------------------------------------------------------------------------
# WeightedCompositeScorer — When / Then
# ---------------------------------------------------------------------------


@when("I score the fragments with default weights")
def step_score_default(context: Context) -> None:
    scorer = WeightedCompositeScorer()
    context.phase2_result = list(scorer.score(context.phase2_fragments))


@when("I score the fragments with custom weights")
def step_score_custom(context: Context) -> None:
    scorer = WeightedCompositeScorer(weights=context.scorer_weights)
    context.phase2_result = list(scorer.score(context.phase2_fragments))


@when("I score the fragments using score_detailed")
def step_score_detailed(context: Context) -> None:
    scorer = WeightedCompositeScorer()
    context.scored_fragment_list = scorer.score_detailed(context.phase2_fragments)


@then("the scored fragment should have a composite relevance score")
def step_scored_has_composite(context: Context) -> None:
    assert len(context.phase2_result) >= 1
    score = context.phase2_result[0].relevance_score
    assert 0.0 <= score <= 1.0


@then('the scored fragment metadata should contain original relevance "{val}"')
def step_scored_original_relevance(context: Context, val: str) -> None:
    meta = context.phase2_result[0].metadata
    assert meta.get("_original_relevance") == val


@then("the scored fragment relevance score should be {expected:g}")
def step_scored_exact_relevance(context: Context, expected: float) -> None:
    assert len(context.phase2_result) >= 1
    actual = context.phase2_result[0].relevance_score
    assert abs(actual - expected) < 1e-4, f"Expected {expected}, got {actual}"


@then('the scored fragment metadata should contain component "{name}" near {val:g}')
def step_scored_component_near(context: Context, name: str, val: float) -> None:
    meta = context.phase2_result[0].metadata
    key = f"_score_{name}"
    actual = float(meta[key])
    assert abs(actual - val) < 0.01, f"Expected ~{val}, got {actual}"


@then("{count:d} scored fragments should be returned")
def step_scored_count(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count


@then("the scored fragment relevance score should be at most {max_val:g}")
def step_scored_at_most(context: Context, max_val: float) -> None:
    assert context.phase2_result[0].relevance_score <= max_val


@then("{count:d} ScoredFragment objects should be returned")
def step_scored_fragment_objects_count(context: Context, count: int) -> None:
    assert len(context.scored_fragment_list) == count


@then("each ScoredFragment should have a composite_score between 0.0 and 1.0")
def step_scored_fragment_range(context: Context) -> None:
    for sf in context.scored_fragment_list:
        assert 0.0 <= sf.composite_score <= 1.0


@then("each ScoredFragment should have score_components")
def step_scored_fragment_has_components(context: Context) -> None:
    for sf in context.scored_fragment_list:
        assert len(sf.score_components) > 0


# ---------------------------------------------------------------------------
# GreedyKnapsackPacker — When / Then
# ---------------------------------------------------------------------------


@when("I pack the fragments into budget")
def step_pack(context: Context) -> None:
    packer = GreedyKnapsackPacker()
    context.phase2_result = list(
        packer.pack(context.phase2_fragments, context.phase2_budget)
    )


@then("{count:d} fragment should be packed")
def step_packed_count_singular(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count, (
        f"Expected {count}, got {len(context.phase2_result)}"
    )


@then("{count:d} fragments should be packed")
def step_packed_count_plural(context: Context, count: int) -> None:
    assert len(context.phase2_result) == count, (
        f"Expected {count}, got {len(context.phase2_result)}"
    )


@then("the packed total tokens should be at most {max_tokens:d}")
def step_packed_total_at_most(context: Context, max_tokens: int) -> None:
    total = sum(f.token_count for f in context.phase2_result)
    assert total <= max_tokens, f"Expected <= {max_tokens}, got {total}"


@then('the packed fragment uko_node should be "{expected}"')
def step_packed_single_node(context: Context, expected: str) -> None:
    assert len(context.phase2_result) >= 1
    assert context.phase2_result[0].uko_node == expected


@then('the packed fragments should include uko_node "{expected}"')
def step_packed_includes_node(context: Context, expected: str) -> None:
    nodes = {f.uko_node for f in context.phase2_result}
    assert expected in nodes, f"Expected {expected} in {nodes}"


@then('the packed fragment for "{node}" should have depth {depth:d}')
def step_packed_fragment_depth(context: Context, node: str, depth: int) -> None:
    for frag in context.phase2_result:
        if frag.uko_node == node:
            assert frag.detail_depth == depth, (
                f"Expected depth {depth} for {node}, got {frag.detail_depth}"
            )
            return
    msg = f"No fragment found for node {node}"
    raise AssertionError(msg)


# ---------------------------------------------------------------------------
# Pipeline Integration — Given / When / Then
# ---------------------------------------------------------------------------

_TEST_PLAN_ID = "01JQTESTPN00000000000000BB"


@given("the ACMS pipeline with a ContentHashDeduplicator")
def step_pipeline_with_dedup(context: Context) -> None:
    context.phase2_pipeline = ACMSPipeline(
        deduplicator=ContentHashDeduplicator(),
    )


@given("the ACMS pipeline with a MaxDepthResolver")
def step_pipeline_with_resolver(context: Context) -> None:
    context.phase2_pipeline = ACMSPipeline(
        depth_resolver=MaxDepthResolver(),
    )


@given("the ACMS pipeline with a WeightedCompositeScorer")
def step_pipeline_with_scorer(context: Context) -> None:
    context.phase2_pipeline = ACMSPipeline(
        scorer=WeightedCompositeScorer(),
    )


@given("the ACMS pipeline with a GreedyKnapsackPacker")
def step_pipeline_with_packer(context: Context) -> None:
    context.phase2_pipeline = ACMSPipeline(
        packer=GreedyKnapsackPacker(),
    )


@when("I assemble context through the pipeline")
def step_assemble_pipeline(context: Context) -> None:
    budget = getattr(
        context, "phase2_budget", ContextBudget(max_tokens=100000, reserved_tokens=0)
    )
    payload = context.phase2_pipeline.assemble(
        plan_id=_TEST_PLAN_ID,
        fragments=context.phase2_fragments,
        budget=budget,
    )
    context.phase2_payload = payload


@when("I assemble context through the pipeline with budget")
def step_assemble_pipeline_with_budget(context: Context) -> None:
    payload = context.phase2_pipeline.assemble(
        plan_id=_TEST_PLAN_ID,
        fragments=context.phase2_fragments,
        budget=context.phase2_budget,
    )
    context.phase2_payload = payload


@then("the pipeline output should contain {count:d} fragment")
def step_pipeline_fragment_count_singular(context: Context, count: int) -> None:
    assert len(context.phase2_payload.fragments) == count, (
        f"Expected {count}, got {len(context.phase2_payload.fragments)}"
    )


@then("the pipeline output should contain {count:d} fragments")
def step_pipeline_fragment_count_plural(context: Context, count: int) -> None:
    assert len(context.phase2_payload.fragments) == count


@then("the pipeline output fragments should have updated relevance scores")
def step_pipeline_scorer_applied(context: Context) -> None:
    for frag in context.phase2_payload.fragments:
        # Scorer adds _original_relevance metadata
        assert "_original_relevance" in frag.metadata


@then("the pipeline output total tokens should be at most {max_tokens:d}")
def step_pipeline_total_at_most(context: Context, max_tokens: int) -> None:
    assert context.phase2_payload.total_tokens <= max_tokens