codebenders-datathon/tests/training/test_distill.py at 42c8f4ac7a0576181024706f396b396bf1b67027 · devcolor/codebenders-datathon · GitHub

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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
"""Tests for training.distill — teacher model distillation."""

import json
import pytest
from unittest.mock import patch, MagicMock

from training.distill import (
    validate_json,
    call_teacher,
    generate_pairs,
)


class TestValidateJson:
    def test_valid_json(self):
        result = validate_json('{"key": "value"}')
        assert result == {"key": "value"}

    def test_strips_markdown_fences(self):
        result = validate_json('```json\n{"key": "value"}\n```')
        assert result == {"key": "value"}

    def test_returns_none_for_invalid(self):
        assert validate_json("not json") is None

    def test_returns_none_for_empty(self):
        assert validate_json("") is None
        assert validate_json(None) is None

    def test_returns_none_for_non_dict(self):
        assert validate_json("[1, 2, 3]") is None


class TestCallTeacher:
    def test_calls_anthropic_backend(self):
        mock_client = MagicMock()
        mock_message = MagicMock()
        mock_message.content = [MagicMock(text='{"result": "ok"}')]
        mock_message.usage.input_tokens = 100
        mock_message.usage.output_tokens = 50
        mock_client.messages.create.return_value = mock_message

        with patch("training.distill._get_anthropic_client", return_value=mock_client):
            result = call_teacher(
                system="system prompt",
                user="user prompt",
                backend="anthropic",
                model="claude-sonnet-4-20250514",
            )

        assert result == '{"result": "ok"}'
        mock_client.messages.create.assert_called_once()

    def test_calls_ollama_backend(self):
        mock_response = {"message": {"content": '{"result": "ok"}'}}

        with patch("training.distill.ollama") as mock_ollama:
            mock_ollama.chat.return_value = mock_response
            result = call_teacher(
                system="system prompt",
                user="user prompt",
                backend="ollama",
                model="qwen3.5:27b",
            )

        assert result == '{"result": "ok"}'
        mock_ollama.chat.assert_called_once()


class TestGenerateExplainerPairs:
    def test_generates_pairs_from_seed_data(self, sample_school_config, sample_course_pairing_data):
        mock_response = json.dumps({
            "explanation": "Test explanation",
            "structural_factors": ["factor1"],
            "student_impact": "impact",
            "advisor_recommendation": "recommendation",
            "data_limitations": ["caveat"],
            "related_intervention": None,
        })

        with patch("training.distill.call_teacher", return_value=mock_response):
            pairs = generate_pairs(
                config=sample_school_config,
                seed_data=[sample_course_pairing_data],
                count=2,
                task="explainer",
            )

        assert len(pairs) == 2
        assert "messages" in pairs[0]
        assert len(pairs[0]["messages"]) == 3

    def test_skips_invalid_responses(self, sample_school_config, sample_course_pairing_data):
        with patch("training.distill.call_teacher", return_value="not json"):
            pairs = generate_pairs(
                config=sample_school_config,
                seed_data=[sample_course_pairing_data],
                count=3,
                task="explainer",
            )

        assert len(pairs) == 0


class TestGenerateSummarizerPairs:
    def test_generates_pairs_from_seed_data(self, sample_school_config, sample_query_result_data):
        mock_response = json.dumps({
            "summary": "Test summary",
            "key_insights": ["insight1"],
            "context": "context",
            "action_items": ["action"],
            "caveats": ["caveat"],
        })

        with patch("training.distill.call_teacher", return_value=mock_response):
            pairs = generate_pairs(
                config=sample_school_config,
                seed_data=[sample_query_result_data],
                count=2,
                task="summarizer",
            )

        assert len(pairs) == 2
        assert "messages" in pairs[0]