test_chat_json_repair.py

"""
chat_json's repair pass must not contaminate the result.

Bug (observed live on MiniCPM-V-4.6): the grader's first reply didn't parse, so
the old repair injected a user turn — "That was not valid JSON. Reply again..."
— and the grader then GRADED THAT MESSAGE, returning a real-looking but
nonsensical grade (score 0, explanation "incorrect JSON syntax"). The repair now
folds a format reminder into the existing user turn instead, and drops the bad
reply. These check that contract without any model/GPU.

    python3 -m pytest test_chat_json_repair.py
"""
import os

os.environ["RECALL_STUB"] = "1"

import llm


def _grading_messages():
    return [
        {"role": "system", "content": "You grade an answer. Return ONLY a JSON object."},
        {"role": "user", "content":
         "Question: How does friction generate heat?\n"
         "Reference answer: Friction converts motion into heat.\n"
         "Student answer: More friction means more heat\nGrade it."},
    ]


def test_repair_reasks_task_without_meta_turn(monkeypatch):
    calls: list[list[dict]] = []

    def fake_chat(messages, max_tokens=512):
        calls.append([dict(m) for m in messages])
        if len(calls) == 1:
            return "Sure, here is my assessment of the answer."  # no JSON
        return '{"score": 3, "explanation": "Mostly right.", "missed_concept": ""}'

    monkeypatch.setattr(llm, "chat", fake_chat)

    data = llm.chat_json(_grading_messages(), max_tokens=512)
    assert data == {"score": 3, "explanation": "Mostly right.", "missed_concept": ""}, data

    # The second (repair) call must re-ask the SAME grading task...
    second = calls[1]
    text = " ".join(m["content"] for m in second if isinstance(m.get("content"), str))
    assert "Student answer: More friction" in text, "lost the grading task"
    # ...with a format reminder folded in...
    assert "ONLY the raw JSON" in text
    # ...and WITHOUT the contaminating meta phrasing or the echoed bad reply.
    assert "not valid JSON" not in text
    assert all(m.get("role") != "assistant" for m in second), "bad reply re-fed"
    assert len(second) == len(_grading_messages()), "injected an extra turn"
    print("ok  repair re-asks the task, no contaminating meta turn")


def test_augment_handles_multimodal_user_content():
    # Image-PDF path passes list content (text + PIL images). The reminder must
    # ride along as a text part, not crash on the list.
    msgs = [
        {"role": "system", "content": "Generate quiz JSON."},
        {"role": "user", "content": ["<image>", "Generate the JSON array."]},
    ]
    out = llm._augment_last_user(msgs)
    assert isinstance(out[-1]["content"], list)
    assert any("ONLY the raw JSON" in p for p in out[-1]["content"] if isinstance(p, str))
    # original is untouched
    assert msgs[-1]["content"] == ["<image>", "Generate the JSON array."]
    print("ok  augment appends a text reminder to multimodal content")


def test_returns_none_when_every_attempt_fails(monkeypatch):
    # Both passes unparseable -> None, so callers fall back honestly (never a
    # garbage grade).
    monkeypatch.setattr(llm, "chat", lambda messages, max_tokens=512: "no json here")
    assert llm.chat_json(_grading_messages(), max_tokens=512) is None
    print("ok  None when no attempt yields JSON")


if __name__ == "__main__":
    import pytest

    raise SystemExit(pytest.main([__file__, "-q"]))