tests/test_integration.py

"""End-to-end integration test: the full ENSEMBLE pipeline."""
from __future__ import annotations

import pytest

from ensemble import Expert, Brain


MATH_QA = [
    ("what is two plus two", "two plus two equals four"),
    ("what is three times three", "three times three equals nine"),
    ("what is pi", "pi is approximately three point one four"),
    ("what is ten minus four", "ten minus four equals six"),
    ("what is five times five", "five times five equals twenty five"),
] * 5

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"),
    ("what is the capital of italy", "the capital of italy is rome"),
    ("what is the capital of egypt", "the capital of egypt is cairo"),
    ("what is the capital of spain", "the capital of spain is madrid"),
] * 5


def test_full_pipeline(tmp_path):
    """Build -> save (compressed) -> reload -> brain -> query -> think."""
    # 1. build experts
    math = Expert.from_qa_pairs(MATH_QA, domain="math", D=3000, seed=1)
    geo = Expert.from_qa_pairs(GEO_QA, domain="geography", D=3000, seed=2)

    # 2. save and verify compression
    rm = math.save(tmp_path / "math.exp")
    rg = geo.save(tmp_path / "geo.exp")
    assert rm.expert_size_bytes < rm.source_size_bytes
    assert rg.expert_size_bytes < rg.source_size_bytes

    # 3. reload — memory and answer capability round-trip
    math2 = Expert.load(rm.path)
    geo2 = Expert.load(rg.path)
    assert math2.n_traces == math.n_traces
    assert geo2.n_traces == geo.n_traces
    # the reloaded expert can still answer questions
    assert isinstance(math2.answer("what is pi"), str)

    # 4. assemble a brain and query
    brain = Brain()
    brain.add_expert(math2)
    brain.add_expert(geo2)
    assert brain.n_experts == 2

    # 5. query routes correctly
    r1 = brain.query("what is pi", max_new_tokens=30)
    assert r1.dominant_expert == "math"
    r2 = brain.query("what is the capital of france", max_new_tokens=30)
    assert r2.dominant_expert == "geography"

    # 6. think produces concepts without crashing
    think = brain.think(n_cycles=3)
    assert think.n_cycles == 3

    # 7. self-modify doesn't crash
    brain.self_modify()

    # 8. remove an expert (Lego)
    assert brain.remove_expert("math") is True
    assert brain.n_experts == 1


def test_compression_at_realistic_size(tmp_path):
    """At realistic dataset sizes, compression should be substantial."""
    text = (
        "The mitochondria is the powerhouse of the cell and produces energy. "
        "Photosynthesis converts sunlight into chemical energy in plants. "
        "DNA contains the genetic instructions for living organisms. "
    ) * 30
    e = Expert.from_text(text, domain="bio", D=2000)
    result = e.save(tmp_path / "bio.exp")
    # realistic text should compress well (>> 3x)
    assert result.compression_ratio > 3.0


def test_multiple_experts_couple():
    """Adding more experts should not break the attractor."""
    experts = [
        Expert.from_qa_pairs([("what is x", f"x is thing {i}")] * 3,
                             domain=f"d{i}", D=2000)
        for i in range(4)
    ]
    brain = Brain()
    for e in experts:
        brain.add_expert(e)
    res = brain.query("what is x", max_new_tokens=15)
    assert res.n_experts == 4