tests/test_tutor_flashcards.py

from unittest.mock import MagicMock

from app.agents.tutor_agent import TutorAgent, _Flashcard, _FlashcardsPayload


def card(front: str, back: str = "A grounded answer.", indexes: list[int] | None = None) -> _Flashcard:
    return _Flashcard(
        front=front,
        back=back,
        source_chunk_indexes=[0] if indexes is None else indexes,
    )


def payload(cards: list[_Flashcard]) -> _FlashcardsPayload:
    return _FlashcardsPayload(cards=cards)


def agent_with(*responses: _FlashcardsPayload) -> tuple[TutorAgent, MagicMock]:
    client = MagicMock()
    client.structured_complete.side_effect = responses
    return TutorAgent(client=client), client


CHUNKS = [{"chunk_index": 0, "text": "Attention uses query-key similarity."}]


def test_ten_valid_cards_use_one_logical_cerebras_call():
    expected = [card(f"Question {index}?") for index in range(10)]
    agent, client = agent_with(payload(expected))

    result = agent.generate_flashcards("attention", CHUNKS, "graduate")

    assert result.cards == expected
    assert client.structured_complete.call_count == 1


def test_invalid_first_pass_cards_trigger_one_bounded_repair_call():
    first_pass = payload([
        card("Valid question?"),
        card(""),
        card("  VALID   question?  "),
        card("What is $x?"),
        card("Vector notation?", r"Use \v for the vector."),
        card("According to the chunk, what happens?"),
        card("Bad source?", indexes=[1]),
        card("Negative source?", indexes=[-1]),
        card("Empty answer?", back="   "),
    ])
    repaired = [card(f"Repair question {index}?") for index in range(9)]
    agent, client = agent_with(first_pass, payload(repaired))

    result = agent.generate_flashcards("attention", CHUNKS, "graduate")

    assert [item.front for item in result.cards] == [
        "Valid question?",
        *(f"Repair question {index}?" for index in range(9)),
    ]
    assert client.structured_complete.call_count == 2
    repair_messages = client.structured_complete.call_args_list[1].args[0]
    repair_prompt = repair_messages[-1]["content"]
    assert "Generate exactly 9 replacement flashcards" in repair_prompt
    assert "Valid question?" in repair_prompt


def test_invalid_repair_results_do_not_trigger_a_third_call():
    agent, client = agent_with(
        payload([card("Only valid card?")]),
        payload([
            card("Only valid card?"),
            card("Broken $math?"),
            card("According to the source, what is it?"),
        ]),
    )

    result = agent.generate_flashcards("attention", CHUNKS, "graduate")

    assert [item.front for item in result.cards] == ["Only valid card?"]
    assert client.structured_complete.call_count == 2