chen/tests/conftest.py at main · twomathematicians-code/chen · 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
"""Shared pytest fixtures for the CHEN test suite."""

from __future__ import annotations

import pytest

from chen.backends.mock import MockBackend
from chen.core.expert import Expert, ExpertRole
from chen.core.memory import InMemoryMemory


@pytest.fixture
def mock_backend_small() -> MockBackend:
    """A small 3B mock backend."""
    return MockBackend(params_m=3_000, role_hint="analyst")


@pytest.fixture
def mock_backend_large() -> MockBackend:
    """A larger 8B mock backend."""
    return MockBackend(params_m=8_000, role_hint="reasoner")


@pytest.fixture
def three_experts() -> list[Expert]:
    """The canonical 3-expert garage: Analyst, Reasoner, Synthesizer."""
    return [
        Expert(
            name="analyst",
            role=ExpertRole.ANALYST,
            backend=MockBackend(params_m=3_000, role_hint="analyst"),
        ),
        Expert(
            name="reasoner",
            role=ExpertRole.REASONER,
            backend=MockBackend(params_m=8_000, role_hint="reasoner"),
        ),
        Expert(
            name="synthesizer",
            role=ExpertRole.SYNTHESIZER,
            backend=MockBackend(params_m=3_000, role_hint="synthesizer"),
        ),
    ]


@pytest.fixture
def four_experts() -> list[Expert]:
    """A 4-expert garage: Analyst, Reasoner, Coder, Synthesizer."""
    return [
        Expert(
            name="analyst",
            role=ExpertRole.ANALYST,
            backend=MockBackend(params_m=3_000, role_hint="analyst"),
        ),
        Expert(
            name="reasoner",
            role=ExpertRole.REASONER,
            backend=MockBackend(params_m=8_000, role_hint="reasoner"),
        ),
        Expert(
            name="coder",
            role=ExpertRole.CODER,
            backend=MockBackend(params_m=7_000, role_hint="coder"),
        ),
        Expert(
            name="synthesizer",
            role=ExpertRole.SYNTHESIZER,
            backend=MockBackend(params_m=3_000, role_hint="synthesizer"),
        ),
    ]


@pytest.fixture
def fresh_memory() -> InMemoryMemory:
    """A fresh in-memory store (no entries)."""
    return InMemoryMemory()