Spaces:
Sleeping
Sleeping
File size: 2,728 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | 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
|