tests/test_brain.py

"""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