File size: 3,431 Bytes
1f71c7d | 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 | """Tests for the Brain: multi-expert ensemble."""
from __future__ import annotations
import pytest
from ensemble import Brain, BrainConfig, Expert
MATH_QA = [
("what is two plus two", "two plus two equals four"),
("what is pi", "pi is approximately three point one four"),
] * 3
GEO_QA = [
("what is the capital of france", "the capital of france is paris"),
("what is the capital of japan", "the capital of japan is tokyo"),
] * 3
@pytest.fixture
def brain():
b = Brain()
b.add_expert(Expert.from_qa_pairs(MATH_QA, domain="math", D=3000))
b.add_expert(Expert.from_qa_pairs(GEO_QA, domain="geography", D=3000))
return b
class TestBrainRoster:
def test_add_and_list(self):
b = Brain()
assert b.n_experts == 0
b.add_expert(Expert.from_qa_pairs(MATH_QA, domain="m", D=2000))
assert b.n_experts == 1
assert b.list_experts() == ["m"]
def test_remove(self):
b = Brain()
b.add_expert(Expert.from_qa_pairs(MATH_QA, domain="m", D=2000))
assert b.remove_expert("m") is True
assert b.n_experts == 0
assert b.remove_expert("nope") is False
def test_add_with_custom_name(self):
b = Brain()
b.add_expert(Expert.from_qa_pairs(MATH_QA, domain="m", D=2000),
name="custom")
assert "custom" in b.list_experts()
class TestBrainQuery:
def test_query_returns_result(self, brain):
res = brain.query("what is pi", max_new_tokens=30)
assert isinstance(res.answer, str)
assert res.n_experts == 2
assert res.dominant_expert in ("math", "geography")
def test_query_routes_to_right_expert(self, brain):
# math question -> math should dominate or contribute
res = brain.query("what is pi", max_new_tokens=30)
assert res.dominant_expert == "math"
# geography question -> geography
res2 = brain.query("what is the capital of france", max_new_tokens=30)
assert res2.dominant_expert == "geography"
def test_per_expert_answers_populated(self, brain):
res = brain.query("what is pi", max_new_tokens=30)
assert set(res.per_expert_answers.keys()) == {"math", "geography"}
def test_empty_brain_query(self):
b = Brain()
res = b.query("anything")
assert res.answer == ""
assert res.n_experts == 0
class TestBrainThink:
def test_think_returns_result(self, brain):
result = brain.think(n_cycles=3)
assert result.n_cycles == 3
assert isinstance(result.concepts, list)
def test_think_concepts_accumulate(self, brain):
brain.think(n_cycles=4)
n1 = brain.n_concepts
brain.think(n_cycles=4)
n2 = brain.n_concepts
assert n2 >= n1 # never decreases
class TestBrainSelfModify:
def test_self_modify_returns_result_or_none(self, brain):
r = brain.self_modify()
# may accept or reject, but shouldn't crash
assert r is None or hasattr(r, "accepted")
def test_self_modify_history(self, brain):
for _ in range(3):
brain.self_modify()
# at least some attempts were recorded
assert brain._modifier is not None
class TestBrainConfig:
def test_custom_config(self):
cfg = BrainConfig(coupling_strength=3.0, n_iterations=10)
b = Brain(config=cfg)
assert b.config.coupling_strength == 3.0
|