clustering.py

"""
Group failed runs by exact signature agreement and order the resulting patterns.

Two runs are clustered together only when they agree on what the verifier
rejected, where the error signal first appeared, how the agent behavior
contributed, and which reusable mechanism was involved. The goal is not latent
semantic similarity between traces; it is aggregating failures that plausibly
admit the same harness-level intervention.

Cluster-level evidence is computed from mechanical tree statistics rather than
by re-querying a model, so a bundle stays a deterministic function of its
records.
"""

from collections import Counter
from dataclasses import dataclass

from shrlm.optimization.taxonomy import (
    CAUSAL_WEIGHT,
    AgentMechanism,
    CausalStatus,
    FailingLevel,
)
from shrlm.optimization.types import (
    AttributionErrorKind,
    FailurePattern,
    FailureRecord,
    FailureSignature,
)

# Patterns with support below this are flagged, never dropped. A mechanism seen
# once may still be real, and silently discarding it would misrepresent the
# round as having found nothing there.
DEFAULT_MIN_SUPPORT = 2

DEFAULT_REPRESENTATIVES = 3

# The cutoff a run's ``collapse_ratio`` must clear to count as "collapsed"
# (mostly delegated to sub-calls rather than reasoned about directly). This is
# the one place the threshold is defined; other modules reporting the same
# signal (e.g. ``shrlm.experiment.collapse_and_attribution``) import it rather
# than hardcoding their own copy of the cutoff.
COLLAPSE_RATIO_THRESHOLD = 0.8

# Weights of the actionability terms. Recorded in the bundle config so a change
# is visible in the artifact rather than only in the code.
ACTIONABILITY_WEIGHTS: dict[str, float] = {
    "causal": 0.4,
    "grounded": 0.3,
    "surface": 0.2,
    "homogeneity": 0.1,
}


@dataclass(frozen=True)
class ClusteringConfig:
    min_support: int = DEFAULT_MIN_SUPPORT
    n_representatives: int = DEFAULT_REPRESENTATIVES


def attributed(records: list[FailureRecord]) -> list[FailureRecord]:
    """Records that carry a signature. Unattributed ones stay in the totals."""
    return [
        record
        for record in records
        if record.signature is not None and not record.attribution_failed
    ]


def grounded_fraction(records: list[FailureRecord]) -> float:
    if not records:
        return 0.0
    return sum(1 for record in records if record.level_grounded) / len(records)


def homogeneity(records: list[FailureRecord]) -> float:
    """
    How consistently the cluster's members describe the same thing.

    Members share a signature by construction, so divergent free-text detail is
    a signal that the enum member is being stretched across distinct behaviors.
    """
    if not records:
        return 0.0
    details = {
        (record.detail.agent_mechanism_detail.strip().lower() if record.detail else "")
        for record in records
    }
    return 1.0 - (len(details) - 1) / len(records)


def actionability(records: list[FailureRecord]) -> float:
    """
    How likely this cluster is to map onto a single harness surface.

    Orders the proposer's reading order only. It is not a prediction of edit
    success, and no promotion decision may consume it.
    """
    if not records:
        return 0.0

    signature = records[0].signature
    assert signature is not None  # attributed() filtered these

    causal = sum(CAUSAL_WEIGHT[record.signature.causal_status] for record in records) / len(records)
    surface = 0.0 if signature.agent_mechanism is AgentMechanism.OTHER else 1.0

    return (
        ACTIONABILITY_WEIGHTS["causal"] * causal
        + ACTIONABILITY_WEIGHTS["grounded"] * grounded_fraction(records)
        + ACTIONABILITY_WEIGHTS["surface"] * surface
        + ACTIONABILITY_WEIGHTS["homogeneity"] * homogeneity(records)
    )


def shared_symptoms(records: list[FailureRecord]) -> list[str]:
    """
    Mechanical statements about what the cluster's runs have in common.

    Every line here is arithmetic over tree statistics. Nothing consults a
    model, so the bundle is reproducible from the records alone.
    """
    n = len(records)
    stats = [record.stats for record in records]

    collapsed = sum(1 for s in stats if s.collapse_ratio > COLLAPSE_RATIO_THRESHOLD)
    fallback = sum(1 for s in stats if s.terminated_by_fallback)
    degraded = sum(1 for s in stats if s.suspected_lost_subcalls > 0)
    no_children = sum(1 for s in stats if s.n_nodes <= 1)
    errored = sum(1 for s in stats if s.n_errored > 0)

    median_children = sorted(s.n_nodes - 1 for s in stats)[n // 2]
    median_iterations = sorted(s.n_iterations for s in stats)[n // 2]
    median_depth = sorted(s.max_observed_depth for s in stats)[n // 2]

    lines = [
        f"median sub-calls per run: {median_children}",
        f"median iterations: {median_iterations}",
        f"median observed depth: {median_depth}",
    ]
    if collapsed:
        lines.append(f"largest sub-call took >80% of the root context in {collapsed}/{n} runs")
    if no_children:
        lines.append(f"no sub-calls issued at all in {no_children}/{n} runs")
    if fallback:
        lines.append(f"answer synthesized by the iteration fallback in {fallback}/{n} runs")
    if errored:
        lines.append(f"at least one sub-call errored in {errored}/{n} runs")
    if degraded:
        lines.append(
            f"trace known to be missing sub-calls in {degraded}/{n} runs "
            "(lower bound; child-level attribution is understated here)"
        )
    lines.extend(answer_shape_symptoms(records))
    return lines


def answer_shape_symptoms(records: list[FailureRecord]) -> list[str]:
    """
    What the produced answers have in common, as arithmetic over the verdicts.

    The tree statistics above were identical for every pattern of
    experiment_kimi ("no sub-calls issued at all in n/n runs") and told the
    proposer nothing; the discriminating facts were in the produced strings --
    quoted items, a missing answer line, an empty list -- which it could only
    see through truncated evidence quotes. Descriptive only: each line states a
    count, never a remedy (``bundle.assert_no_prescription``).
    """
    n = len(records)
    causes = Counter(
        record.verdict.cause.value if record.verdict.cause is not None else "pass"
        for record in records
    )
    lines = [
        "verifier causes: "
        + ", ".join(f"{cause} {count}/{n}" for cause, count in sorted(causes.items()))
    ]
    produced = [record.verdict.produced or "" for record in records]
    quoted = sum(1 for text in produced if _bracket_list_has_quoted_items(text))
    no_list = sum(1 for text in produced if "[" not in text.strip().split("\n")[-1])
    empty = sum(1 for text in produced if text.strip().endswith("[]"))
    if quoted:
        lines.append(f"produced answer list carried quoted items in {quoted}/{n} runs")
    if no_list:
        lines.append(f"no bracketed answer list on the last produced line in {no_list}/{n} runs")
    if empty:
        lines.append(f"produced answer was the empty list in {empty}/{n} runs")
    return lines


def _bracket_list_has_quoted_items(text: str) -> bool:
    last = text.strip().split("\n")[-1]
    start, end = last.find("["), last.rfind("]")
    if start < 0 or end <= start:
        return False
    body = last[start + 1 : end]
    return any(item.strip()[:1] in ("'", '"') for item in body.split(",") if item.strip())


def verifier_evidence(records: list[FailureRecord], limit: int = 3) -> list[str]:
    """Concrete produced-versus-gold pairs, so the cause is inspectable.

    ``produced`` comes first: the proposer's prompt bounds each entry at
    render time, and a 181-node gold list placed first hid the 20-character
    produced string that carried the whole diagnosis (POST_MORTEM.md
    section 6).
    """
    lines = []
    for record in sorted(records, key=lambda r: r.instance_id)[:limit]:
        lines.append(
            f"{record.instance_id}: produced {record.verdict.produced!r}, "
            f"expected {record.verdict.gold!r}"
        )
    return lines


def select_representatives(records: list[FailureRecord], count: int) -> list[str]:
    """
    A fixed selection rule, so ordering is stable under input permutation.

    Grounded records first, since their failing level is a checkable fact
    rather than a judgment.
    """
    ranked = sorted(
        records,
        key=lambda r: (not r.level_grounded, len(r.digest_sha256), r.instance_id),
    )
    return [record.instance_id for record in ranked[:count]]


def cluster_failures(
    records: list[FailureRecord], config: ClusteringConfig | None = None
) -> list[FailurePattern]:
    """Group by exact signature agreement, then order by distinct-instance
    support and actionability (see ``rank_patterns``)."""
    config = config or ClusteringConfig()

    groups: dict[tuple[str, str, str, str], list[FailureRecord]] = {}
    for record in attributed(records):
        assert record.signature is not None
        groups.setdefault(record.signature.key(), []).append(record)

    patterns = []
    for members in groups.values():
        signature = members[0].signature
        assert signature is not None
        patterns.append(
            FailurePattern(
                signature=signature,
                support=len(members),
                instance_support=len({record.instance_id for record in members}),
                instance_ids=sorted(record.instance_id for record in members),
                representatives=select_representatives(members, config.n_representatives),
                shared_symptoms=shared_symptoms(members),
                verifier_evidence=verifier_evidence(members),
                grounded_fraction=grounded_fraction(members),
                surface=signature.surface(),
                actionability=actionability(members),
                below_support_floor=len(members) < config.min_support,
            )
        )

    return rank_patterns(patterns) + environment_summary_patterns(records, config)


# Symptom prefix that marks a synthetic environment-summary pattern. The
# proposer prompt names this marker so the proposer can recognize summary
# entries; nothing else in the pipeline branches on it.
ENVIRONMENT_SUMMARY_MARKER = "environment summary (synthetic)"


def environment_summary_patterns(
    records: list[FailureRecord], config: ClusteringConfig | None = None
) -> list[FailurePattern]:
    """Low-actionability summary patterns over environment-skipped records.

    Environment-owned failures never reach the attributor (``mining.
    environment_caused``), so they cannot cluster into agent mechanisms.
    Discarding them entirely, though, forecloses a legitimate edit class:
    harness-side *handling* of environment failures (recovery-instruction
    treatment of sub-call errors, for example). This function restores the
    proposer-judgment layer as a ranking rather than a gate: one synthetic
    pattern per verifier cause, honest about its nature (mechanism OTHER,
    causal status UNATTRIBUTED, actionability 0.0, below the support floor,
    symptoms prefixed with ``ENVIRONMENT_SUMMARY_MARKER``), appended after
    every ranked real pattern so it is read last. The proposer may, at its
    own judgment, propose harness-side handling against it -- or skip it,
    which is the expected default.
    """
    config = config or ClusteringConfig()
    skipped = [
        record
        for record in records
        if record.attribution_error_kind is AttributionErrorKind.ENVIRONMENT
    ]
    if not skipped:
        return []

    by_cause: dict[str, list[FailureRecord]] = {}
    for record in skipped:
        assert record.verdict.cause is not None  # a failing verdict carries a cause
        by_cause.setdefault(record.verdict.cause.value, []).append(record)

    patterns = []
    for cause_value in sorted(by_cause):
        members = by_cause[cause_value]
        instance_ids = sorted({record.instance_id for record in members})
        reasons = sorted({record.attribution_error for record in members})
        patterns.append(
            FailurePattern(
                signature=FailureSignature(
                    verifier_cause=members[0].verdict.cause,
                    failing_level=FailingLevel.UNDETERMINED,
                    causal_status=CausalStatus.UNATTRIBUTED,
                    agent_mechanism=AgentMechanism.OTHER,
                ),
                support=len(members),
                instance_support=len(instance_ids),
                instance_ids=sorted(record.instance_id for record in members),
                representatives=instance_ids[: config.n_representatives],
                shared_symptoms=[
                    f"{ENVIRONMENT_SUMMARY_MARKER}: {len(members)} run(s) ended with "
                    f"verifier cause {cause_value}; no agent behavior was recorded for "
                    "these terminations and no per-run attribution was performed",
                    "the terminations themselves are outside harness reach; a candidate "
                    "against this pattern can only change harness-side handling around "
                    "such terminations",
                ],
                verifier_evidence=reasons,
                grounded_fraction=0.0,
                surface=None,
                actionability=0.0,
                below_support_floor=True,
            )
        )
    return patterns


def rank_patterns(patterns: list[FailurePattern]) -> list[FailurePattern]:
    """
    Order by distinct-instance support, then actionability, then signature.

    ``instance_support`` rather than run-level ``support`` (KTD7/R2): repeated
    failing attempts of one instance are corroboration, not breadth, and must
    not outrank a pattern seen across distinct instances. Lexicographic rather
    than a weighted blend: the paper orders by support *and* estimated
    actionability, and inventing an exchange rate between frequency and
    tractability would be a modeling choice with no grounding in the evidence.
    """
    return sorted(
        patterns,
        key=lambda p: (-p.instance_support, -p.actionability, p.signature.key()),
    )


def compute_marginals(records: list[FailureRecord]) -> dict[str, dict[str, int]]:
    """
    Per-component counts across all attributed records.

    The full four-tuple space is large relative to a round's failure count, so
    most clusters will be small. These marginals are the principled backoff:
    the by-surface view in particular has ten buckets and is directly
    consumable by a proposal stage that must target one surface anyway.
    """
    usable = attributed(records)

    causes: Counter[str] = Counter()
    levels: Counter[str] = Counter()
    statuses: Counter[str] = Counter()
    mechanisms: Counter[str] = Counter()
    surfaces: Counter[str] = Counter()

    for record in usable:
        signature = record.signature
        assert signature is not None
        causes[signature.verifier_cause.value] += 1
        levels[signature.failing_level.value] += 1
        statuses[signature.causal_status.value] += 1
        mechanisms[signature.agent_mechanism.value] += 1
        surface = signature.surface()
        surfaces[surface.value if surface else "unmapped"] += 1

    return {
        "by_cause": dict(sorted(causes.items())),
        "by_level": dict(sorted(levels.items())),
        "by_causal_status": dict(sorted(statuses.items())),
        "by_mechanism": dict(sorted(mechanisms.items())),
        "by_surface": dict(sorted(surfaces.items())),
    }