Zum Inhalt

tests/test_preflight.py

 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
from touchstone.config import ModelSpec
from touchstone.preflight import preflight_models
from touchstone.runner import StreamEvent


class _FakeClient:
    engine = "fake"
    engine_version = "0"

    def __init__(self, mode):
        self.mode = mode  # "content" | "reasoning" | "empty" | "boom"

    def stream(self, *, messages, model, max_tokens, temperature, seed, extra_body=None):
        if self.mode == "boom":
            raise RuntimeError("model not found")
        if self.mode == "content":
            yield StreamEvent(delta_text="4")
        if self.mode == "reasoning":
            yield StreamEvent(reasoning_text="denke…")
        yield StreamEvent(prompt_tokens=3, completion_tokens=1)


class _MultiModeClient:
    """Yields a different mode per stream() call, in order — to test that preflight_models
    processes every model independently (no short-circuit after the first)."""

    engine = "fake"
    engine_version = "0"

    def __init__(self, modes):
        self._modes = list(modes)
        self._i = 0

    def stream(self, *, messages, model, max_tokens, temperature, seed, extra_body=None):
        mode = self._modes[self._i]
        self._i += 1
        if mode == "content":
            yield StreamEvent(delta_text="4")
        if mode == "reasoning":
            yield StreamEvent(reasoning_text="denke…")
        yield StreamEvent(prompt_tokens=3, completion_tokens=1)


def _budget(_m):
    return 64


def test_preflight_classifies_content():
    [r] = preflight_models(_FakeClient("content"), [ModelSpec(id="m")], budget_for=_budget)
    assert (r.status, r.model) == ("ok", "m")


def test_preflight_classifies_reasoning_only():
    [r] = preflight_models(_FakeClient("reasoning"), [ModelSpec(id="m")], budget_for=_budget)
    assert r.status == "reasoning_only"
    assert r.reasoning_chars > 0


def test_preflight_classifies_empty():
    [r] = preflight_models(_FakeClient("empty"), [ModelSpec(id="m")], budget_for=_budget)
    assert r.status == "empty"


def test_preflight_never_raises_on_error():
    [r] = preflight_models(_FakeClient("boom"), [ModelSpec(id="m")], budget_for=_budget)
    assert r.status == "error"
    assert "not found" in r.detail


def test_preflight_processes_all_models_independently():
    results = preflight_models(
        _MultiModeClient(["content", "reasoning"]),
        [ModelSpec(id="a"), ModelSpec(id="b")],
        budget_for=_budget,
    )
    assert [r.model for r in results] == ["a", "b"]
    assert results[0].status == "ok"
    assert results[1].status == "reasoning_only"