composer_replication/replaysim/tests/test_replaysim.py

"""Replaysim normalization tests — the skip_dj passthrough path.

The full data-juicer path requires `pip install -e .[replaysim]` which we
defer to the user's environment. These tests verify:

1. The package imports cleanly without data-juicer installed.
2. `DJNormalizer(skip_dj=True)` is a working passthrough.
3. The DPOPair → DJ-record → NormalizedDPOPair shape transforms are
   lossless modulo the metadata field.
4. The DPOPair dict shape (TypedDict) is what we expect.
"""
from __future__ import annotations

import pytest

from composer_replication.replaysim import (
    DJNormalizer,
    NormalizedDPOPair,
    replay_and_normalize_trace,
)
from composer_replication.replaysim.normalize import (
    _dj_record_to_normalized,
    _dpo_pair_to_dj_record,
)


def _make_pair(
    state_id: str,
    state_messages: list[dict] | None = None,
    chosen: str = "Four.",
    rejected: str = "Five.",
    n_teachers_agreeing: int = 2,
) -> dict:
    """Helper — DPOPair is a TypedDict, so dicts work directly."""
    return {
        "state_id": state_id,
        "state_messages": state_messages or [{"role": "user", "content": "What is 2+2?"}],
        "chosen": chosen,
        "rejected": rejected,
        "n_teachers_agreeing": n_teachers_agreeing,
    }


def test_dpo_pair_to_dj_record_shape():
    """Records carry BOTH flat-string and chat-messages shapes for chosen/rejected.

    See default.yaml header for why: data-juicer's text_length_filter et al
    consume the flat strings; chat-aware consumers and the round-trip use
    the *_messages fields.
    """
    p = _make_pair("s1")
    rec = _dpo_pair_to_dj_record(p)
    assert rec["state_id"] == "s1"
    assert rec["messages"] == [{"role": "user", "content": "What is 2+2?"}]
    # Flat-string shape (drives text_length_filter, words_num_filter, ...)
    assert rec["chosen"] == "Four."
    assert rec["rejected"] == "Five."
    assert isinstance(rec["chosen"], str)
    assert isinstance(rec["rejected"], str)
    # Chat-messages shape (preserved for chat-aware ops)
    assert rec["chosen_messages"] == [{"role": "assistant", "content": "Four."}]
    assert rec["rejected_messages"] == [{"role": "assistant", "content": "Five."}]
    assert rec["n_teachers_agreeing"] == 2


def test_dj_record_to_normalized_roundtrip():
    p = _make_pair("s2", chosen="C", rejected="R", n_teachers_agreeing=3)
    rec = _dpo_pair_to_dj_record(p)
    rec["__dj_meta__"] = {"ops_applied": ["text_length_filter"]}
    norm = _dj_record_to_normalized(rec)
    assert isinstance(norm, NormalizedDPOPair)
    assert norm.state_id == "s2"
    assert norm.chosen_messages == [{"role": "assistant", "content": "C"}]
    assert norm.rejected_messages == [{"role": "assistant", "content": "R"}]
    assert norm.n_teachers_agreeing == 3
    assert norm.metadata == {"ops_applied": ["text_length_filter"]}


def test_dj_record_to_normalized_preserves_state_messages():
    """The conversation context (state_messages) must round-trip."""
    multi_turn = [
        {"role": "user", "content": "What is 2+2?"},
        {"role": "assistant", "content": "Let me think."},
        {"role": "user", "content": "Just give me a number."},
    ]
    p = _make_pair("s3", state_messages=multi_turn)
    rec = _dpo_pair_to_dj_record(p)
    norm = _dj_record_to_normalized(rec)
    assert norm.state_messages == multi_turn


def test_dj_normalizer_skip_dj_passthrough():
    """skip_dj=True: bypasses data-juicer entirely, just does shape conversion."""
    pairs = [
        _make_pair("s1", chosen="c1", rejected="r1"),
        _make_pair("s2", chosen="c2", rejected="r2"),
    ]
    normalizer = DJNormalizer(skip_dj=True)
    out = normalizer.normalize(pairs)
    assert len(out) == 2
    assert all(isinstance(o, NormalizedDPOPair) for o in out)
    assert out[0].state_id == "s1"
    assert out[1].state_id == "s2"
    assert out[0].metadata == {"skipped": True}
    assert out[1].metadata == {"skipped": True}


def test_dj_normalizer_skip_dj_preserves_count():
    """Passthrough must not drop records — only filter ops do that."""
    pairs = [_make_pair(f"s{i}") for i in range(10)]
    normalizer = DJNormalizer(skip_dj=True)
    out = normalizer.normalize(pairs)
    assert len(out) == 10


def test_dj_normalizer_default_recipe_path_exists():
    """The default recipe ships with the package."""
    assert DJNormalizer.DEFAULT_RECIPE.exists(), \
        f"Default recipe missing at {DJNormalizer.DEFAULT_RECIPE}"


def test_dj_normalizer_real_path_requires_data_juicer():
    """Without skip_dj, instantiation requires data-juicer or fails clearly."""
    try:
        import data_juicer  # type: ignore[import-not-found]  # noqa: F401
    except ImportError:
        with pytest.raises(RuntimeError, match="data-juicer"):
            DJNormalizer(skip_dj=False)
    else:
        # data-juicer IS installed; verify init succeeds with default recipe
        normalizer = DJNormalizer(skip_dj=False)
        assert normalizer.recipe_path == DJNormalizer.DEFAULT_RECIPE


def test_replay_and_normalize_trace_signature():
    """Convenience function is callable and importable. Smoke-only — we
    don't run it against OpenRouter from CI."""
    assert callable(replay_and_normalize_trace)
    # It's an async function
    import inspect
    assert inspect.iscoroutinefunction(replay_and_normalize_trace)


def test_record_handles_missing_optional_fields():
    """A DPOPair dict missing some optional fields shouldn't crash the converter."""
    minimal = {"state_id": "x", "chosen": "a", "rejected": "b"}
    rec = _dpo_pair_to_dj_record(minimal)
    assert rec["state_id"] == "x"
    assert rec["messages"] == []        # missing state_messages → empty list
    assert rec["n_teachers_agreeing"] == 0  # missing → default 0


# ---------------------------------------------------------------------
# Dual-shape contract (Wave 13 review Suggestion 3)
# ---------------------------------------------------------------------
#
# data-juicer's text_length_filter / words_num_filter /
# special_characters_filter / document_deduplicator all expect string-typed
# fields under `text_keys`. Earlier the converter wrapped chosen/rejected
# into list-of-dicts (chat-messages), which would have caused those ops to
# crash or no-op silently. The fix carries BOTH shapes:
#   - chosen / rejected            → flat strings  (filter ops)
#   - chosen_messages / rejected_messages → list-of-dicts (chat-aware ops + round-trip)
#
# The tests below pin that contract.


def test_record_chosen_rejected_are_flat_strings_for_dj_text_ops():
    """text_length_filter & friends expect text_keys to point at strings.

    If we ever regress to wrapping `chosen`/`rejected` into list-of-dicts,
    data-juicer's text-key ops break. Keep this red-line explicit.
    """
    p = _make_pair(
        "s_strings",
        chosen="A long-enough chosen response.",
        rejected="A long-enough rejected response.",
    )
    rec = _dpo_pair_to_dj_record(p)

    assert isinstance(rec["chosen"], str)
    assert isinstance(rec["rejected"], str)
    assert rec["chosen"] == "A long-enough chosen response."
    assert rec["rejected"] == "A long-enough rejected response."
    # Sanity: text_length_filter style usage works without crashing.
    assert len(rec["chosen"]) >= 8
    assert len(rec["rejected"]) >= 8


def test_record_chosen_rejected_messages_carry_chat_shape():
    """The *_messages variants preserve the chat-template-aware shape."""
    p = _make_pair("s_msgs", chosen="hello world", rejected="goodbye world")
    rec = _dpo_pair_to_dj_record(p)

    assert isinstance(rec["chosen_messages"], list)
    assert isinstance(rec["rejected_messages"], list)
    assert rec["chosen_messages"] == [
        {"role": "assistant", "content": "hello world"}
    ]
    assert rec["rejected_messages"] == [
        {"role": "assistant", "content": "goodbye world"}
    ]
    # Both shapes must agree on content.
    assert rec["chosen_messages"][0]["content"] == rec["chosen"]
    assert rec["rejected_messages"][0]["content"] == rec["rejected"]


def test_dj_record_to_normalized_uses_chat_messages_when_present():
    """When *_messages fields are present, the round-trip uses them directly
    (does not re-wrap the flat string)."""
    rec = {
        "state_id": "s_present",
        "messages": [{"role": "user", "content": "q"}],
        "chosen": "some flat str — should be ignored when _messages present",
        "rejected": "another flat str",
        "chosen_messages": [
            {"role": "assistant", "content": "real chosen"},
        ],
        "rejected_messages": [
            {"role": "assistant", "content": "real rejected"},
        ],
        "n_teachers_agreeing": 4,
    }
    norm = _dj_record_to_normalized(rec)
    assert norm.chosen_messages == [{"role": "assistant", "content": "real chosen"}]
    assert norm.rejected_messages == [{"role": "assistant", "content": "real rejected"}]
    assert norm.n_teachers_agreeing == 4


def test_dj_record_to_normalized_falls_back_to_flat_strings():
    """When *_messages fields are absent (e.g. an op only rewrote the flat
    string), the round-trip wraps the flat string into a single assistant
    turn so downstream consumers always see the chat-messages shape."""
    rec = {
        "state_id": "s_fallback",
        "messages": [{"role": "user", "content": "q"}],
        "chosen": "rewritten chosen",
        "rejected": "rewritten rejected",
        # NOTE: no chosen_messages / rejected_messages
        "n_teachers_agreeing": 1,
    }
    norm = _dj_record_to_normalized(rec)
    assert norm.chosen_messages == [
        {"role": "assistant", "content": "rewritten chosen"}
    ]
    assert norm.rejected_messages == [
        {"role": "assistant", "content": "rewritten rejected"}
    ]


def test_round_trip_preserves_strings_through_skip_dj():
    """End-to-end shape sanity: pair → normalize(skip_dj=True) → assert
    chat-messages content matches original strings."""
    pairs = [
        _make_pair("rt1", chosen="alpha", rejected="beta", n_teachers_agreeing=2),
        _make_pair("rt2", chosen="gamma", rejected="delta", n_teachers_agreeing=3),
    ]
    out = DJNormalizer(skip_dj=True).normalize(pairs)
    assert len(out) == 2
    assert out[0].chosen_messages[0]["content"] == "alpha"
    assert out[0].rejected_messages[0]["content"] == "beta"
    assert out[1].chosen_messages[0]["content"] == "gamma"
    assert out[1].rejected_messages[0]["content"] == "delta"


# ---------------------------------------------------------------------
# End-to-end test against the real data-juicer engine.
# ---------------------------------------------------------------------
#
# Install path tried during Wave 13 fix: `pip install py-data-juicer`
# (the canonical PyPI distribution name; `data-juicer` redirects there).
# If that succeeded in the runtime environment, the e2e test runs the
# actual op-graph from default.yaml against a tiny fixture and verifies
# the dual-shape contract holds at the JSONL boundary. If data-juicer
# is NOT importable, the test is skipped.

try:
    import data_juicer  # type: ignore[import-not-found]  # noqa: F401
    _HAS_DJ = True
except ImportError:
    _HAS_DJ = False


@pytest.mark.skipif(not _HAS_DJ, reason="data-juicer not installed")
def test_dj_normalizer_e2e_default_recipe(tmp_path):
    """E2E: real data-juicer engine + default.yaml on a 3-record fixture.

    Verifies:
      1. The engine runs without a type-mismatch crash on the flat-string
         text_keys (this is the bug Wave 13 Suggestion 3 flagged).
      2. Output records survive the round-trip with both shapes intact.
    """
    pairs = [
        _make_pair(
            "e2e1",
            chosen="A reasonably long chosen response with several words.",
            rejected="A reasonably long rejected response with several words.",
            n_teachers_agreeing=2,
        ),
        _make_pair(
            "e2e2",
            chosen="Another solid chosen completion that has enough text.",
            rejected="Another solid rejected completion that has enough text.",
            n_teachers_agreeing=3,
        ),
        _make_pair(
            "e2e3",
            chosen="Third chosen example with sufficient length to pass.",
            rejected="Third rejected example with sufficient length to pass.",
            n_teachers_agreeing=2,
        ),
    ]

    normalizer = DJNormalizer(skip_dj=False)
    out = normalizer.normalize(pairs)

    # Length filter, etc., should NOT drop any of these — all are
    # comfortably within bounds. If we get back fewer than 1, the op-graph
    # is misconfigured.
    assert len(out) >= 1
    for n in out:
        assert isinstance(n, NormalizedDPOPair)
        # Round-trip should always give us chat-messages shape on the way out.
        assert isinstance(n.chosen_messages, list)
        assert isinstance(n.rejected_messages, list)
        assert n.chosen_messages and n.chosen_messages[0]["role"] == "assistant"
        assert n.rejected_messages and n.rejected_messages[0]["role"] == "assistant"