test_cohort_comparison_block.py

import json

import pytest

from oddish.blocks.block import BlockParseError

from api.services.blocks.analyzer.cohort.cohort_comparison_block import (
    SCHEMA_VERSION,
    CohortComparisonBlock,
    CohortInput,
)
from api.services.blocks.analyzer.cohort import cohort_prompts as cp

TRIAL = {
    "trial_id": "t1",
    "components": [
        {
            "trajectory_component": "testing_public",
            "step_ids": [34, 35],
            "summary": "Ran mvn test for a baseline.",
        }
    ],
    "covered_steps": 2,
    "span": 35,
    "coverage": 1.0,
}


def _block(**overrides):
    return CohortComparisonBlock(
        CohortInput(
            task_name="demo-task",
            successful=overrides.get("successful", [TRIAL]),
            failing=overrides.get("failing", [{**TRIAL, "trial_id": "t2"}]),
        ),
        instructions_template=cp.load_cohort_prompt_template(),
    )


def _raw(evidence, schema_version=99):
    return json.dumps(
        {
            "schema_version": schema_version,
            "cohort_success": ["t1"],
            "cohort_failure": ["t2"],
            "summary": "Agents took a test baseline before editing.",
            "categories": [
                {
                    "category": "testing_verification",
                    "label": None,
                    "successful": [
                        {
                            "behavior_description": "ran the verifier first",
                            "evidence": evidence,
                        }
                    ],
                    "failing": [],
                }
            ],
        }
    )


GOOD_EVIDENCE = {
    "trial_id": "t1",
    "trajectory_component": "testing_public",
    "step_ids": [34, 35],
    "quote": "Ran mvn test for a baseline.",
}


def test_prompt_contains_both_cohorts_and_definitions():
    prompt = _block().build_prompt()
    assert "SUCCESSFUL" in prompt and "FAILING" in prompt
    assert "t1" in prompt and "t2" in prompt
    assert "behavior_discovery:" in prompt  # definition, not a bare label


def test_to_output_parses_and_stamps_schema_version():
    out = _block().to_output(_raw([GOOD_EVIDENCE]))
    # The block owns schema_version; a model-supplied value is overwritten.
    # Asserted against the constant: pinning the literal made a deliberate
    # version bump look like a regression.
    assert out["schema_version"] == SCHEMA_VERSION
    assert SCHEMA_VERSION != 99
    assert out["categories"][0]["category"] == "testing_verification"


def test_short_model_name_keeps_dots_that_belong_to_the_name():
    """A generic split on "." would turn gpt-5.4 into "4"."""
    from api.services.blocks.analyzer.cohort.cohort_prompts import (
        short_model_name,
    )

    assert short_model_name("global.anthropic.claude-opus-4-8") == "claude-opus-4-8"
    assert short_model_name("anthropic/claude-fable-5") == "claude-fable-5"
    assert short_model_name("gpt-5.4") == "gpt-5.4"
    assert short_model_name("gemini-3.5-flash") == "gemini-3.5-flash"


def test_short_model_name_strips_every_region_prefix():
    """Opus 4.1 / Opus 4 have no "global." inference profile: they are stored
    as "us.anthropic...", and a global-only strip left them long."""
    from api.services.blocks.analyzer.cohort.cohort_prompts import (
        short_model_name,
    )

    assert (
        short_model_name("us.anthropic.claude-opus-4-1-20250805-v1:0")
        == "claude-opus-4-1-20250805-v1:0"
    )
    assert short_model_name("eu.anthropic.claude-sonnet-4-5") == "claude-sonnet-4-5"
    assert short_model_name("apac.anthropic.claude-haiku-4-5") == "claude-haiku-4-5"
    assert short_model_name("bedrock/apn.amazon.nova-pro-v1:0") == "nova-pro-v1:0"


def test_model_counts_are_ordered_and_stripped():
    out = _block(
        successful=[
            {**TRIAL, "trial_id": "t1", "model": "global.anthropic.claude-opus-4-8"},
            {**TRIAL, "trial_id": "t3", "model": "global.anthropic.claude-opus-4-8"},
            {**TRIAL, "trial_id": "t4", "model": "gemini-3.5-flash"},
        ],
    ).to_output(_raw([GOOD_EVIDENCE]))
    assert out["models"]["successful"] == [
        {"model": "claude-opus-4-8", "trials": 2},
        {"model": "gemini-3.5-flash", "trials": 1},
    ]


def test_mode_is_single_when_one_cohort_is_empty():
    """All-failed and all-succeeded tasks are the cases a reader most wants
    explained; the payload has to say which so the UI drops a column rather
    than drawing an empty one."""
    out = _block(failing=[]).to_output(_raw([GOOD_EVIDENCE]))
    assert out["mode"] == "single"
    assert out["models"]["failing"] == []
    assert _block().to_output(_raw([GOOD_EVIDENCE]))["mode"] == "comparison"


def test_summary_survives_a_clean_comparison():
    out = _block().to_output(_raw([GOOD_EVIDENCE]))
    assert out["dropped"]["categories"] == 0
    assert out["summary"]


def test_summary_is_dropped_when_a_category_is():
    """A headline written against categories that validation then removed is
    an unsourced claim above sourced rows -- exactly what the citation check
    exists to prevent, so it must not outlive them."""
    fabricated = {**GOOD_EVIDENCE, "trial_id": "does-not-exist"}
    out = _block().to_output(_raw([fabricated]))
    assert out["dropped"]["categories"] == 1
    assert out["categories"] == []
    assert "summary" not in out


def test_to_output_validates_citations_before_the_block_persists():
    """Validation must happen in the transform, not after block.run().

    The AnalyzerBlock persists whatever the transform returns, so a citation
    filtered downstream would still be served on every later cache hit.
    """
    fabricated = {**GOOD_EVIDENCE, "trial_id": "does-not-exist"}
    out = _block().to_output(_raw([fabricated]))
    assert out["categories"] == []
    assert out["dropped"]["evidence"] == 1
    assert out["dropped"]["observations"] == 1


def test_to_output_reports_thin_coverage():
    thin = {**TRIAL, "trial_id": "t9", "coverage": 0.1}
    out = _block(successful=[TRIAL, thin]).to_output(_raw([GOOD_EVIDENCE]))
    assert out["thin_coverage"] == ["t9"]


def test_to_output_accepts_a_retired_taxonomy_value():
    # Stored summaries still carry retired labels; the citation resolves
    # because the component really is on the trial, not because an enum
    # happened to list it.
    retired = {
        "trajectory_component": "thinking_diagnose",
        "step_ids": [7, 9],
        "summary": "Worked out why the assertion tripped.",
    }
    trial = {**TRIAL, "components": [retired]}
    cited = {
        "trial_id": "t1",
        "trajectory_component": "thinking_diagnose",
        "step_ids": [7, 9],
        "quote": "Worked out why the assertion tripped.",
    }
    out = _block(successful=[trial]).to_output(_raw([cited]))
    assert out["dropped"]["evidence"] == 0
    assert len(out["categories"]) == 1


def test_to_output_rejects_malformed_json():
    # BlockParseError subclasses ValueError, so AnalyzerBlock's transform
    # contract still holds now that parsing goes through Block.parse.
    with pytest.raises(BlockParseError):
        _block().to_output("not json")


def test_to_output_rejects_a_non_object_reply():
    # Previously reached CohortComparisonOutput(**data) and raised TypeError,
    # which is not a ValueError and so escaped the transform contract.
    with pytest.raises(BlockParseError):
        _block().to_output("[]")