"""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"