tests/test_killer.py

"""Tests for the killer generation layer (7 fixes)."""

import pytest
import numpy as np

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.killer import (
    NgramPredictor, MultiScaleEncoder, NativeRAG,
    IterativeRefiner, TemplateExtractor, massive_ingestion_plan,
    NgramPrediction, MultiScaleResult, RAGResult, RefinementResult, Template,
)


def _build_model(D=5000, ctx=128, radius=200):
    cfg = PalimpsesteConfig(D=D, context_window=ctx, kernel_radius=radius, temperature=0.0)
    lm = PalimpsesteForCausalLM(config=cfg)
    pairs = [
        ("hello", "hi i am palimpseste"),
        ("who are you", "i am palimpseste a hypervectorial cortex"),
        ("what is python", "python is a programming language"),
        ("what is the capital of france", "the capital of france is paris"),
        ("who won the world cup 2018", "france"),
        ("what is recursion", "recursion is when a function calls itself"),
    ]
    lm.build_tokenizer("".join(q + a for q, a in pairs))
    lm.train_on_qa_pairs(pairs)
    return lm, pairs


# ================================================================ BPE
class TestBPETokenizer:
    def test_bpe_encode_decode(self):
        from palimseste.bpe import BPETokenizer
        lm, pairs = _build_model()
        bpe = BPETokenizer(encoder=lm.encoder, vocab_size=300)
        corpus = " ".join(q + a for q, a in pairs)
        bpe.train(corpus)
        ids = bpe.encode("hello world", add_bos=True, add_eos=True)
        assert len(ids) > 0
        text = bpe.decode(ids)
        assert "hello" in text

    def test_bpe_reduces_tokens(self):
        from palimseste.bpe import BPETokenizer
        lm, pairs = _build_model()
        bpe = BPETokenizer(encoder=lm.encoder, vocab_size=500)
        bpe.train(" ".join(q + a for q, a in pairs))
        char_ids = lm.tokenizer.encode("hello world", add_bos=True, add_eos=True)
        bpe_ids = bpe.encode("hello world", add_bos=True, add_eos=True)
        # BPE should produce fewer or equal tokens
        assert len(bpe_ids) <= len(char_ids)

    def test_bpe_vocab_size(self):
        from palimseste.bpe import BPETokenizer
        lm, _ = _build_model()
        bpe = BPETokenizer(encoder=lm.encoder, vocab_size=300)
        bpe.train("hello world test")
        assert bpe.vocab_size_actual > 4  # at least special tokens + some bytes

    def test_lm_accepts_bpe(self):
        """LM should accept a BPE tokenizer via attach_tokenizer."""
        from palimseste.bpe import BPETokenizer
        lm, pairs = _build_model()
        bpe = BPETokenizer(encoder=lm.encoder, vocab_size=300)
        bpe.train(" ".join(q + a for q, a in pairs))
        lm.attach_tokenizer(bpe)
        # vocab_size should be actual, not target
        assert lm.config.vocab_size == bpe.vocab_size_actual


# ================================================================ N-GRAM
class TestNgramPredictor:
    def test_learn_ngrams(self):
        lm, pairs = _build_model()
        ngp = NgramPredictor(lm=lm)
        ngp.learn_ngrams([q + a for q, a in pairs])
        assert len(ngp._bigram_stats) > 0
        assert len(ngp._trigram_stats) > 0

    def test_predict_next(self):
        lm, pairs = _build_model()
        ngp = NgramPredictor(lm=lm)
        ngp.learn_ngrams([q + a for q, a in pairs])
        q_ids = lm.tokenizer.encode("who are you", add_bos=True, add_eos=True) + [1]
        pred = ngp.predict_next(q_ids)
        assert isinstance(pred, NgramPrediction)
        assert len(pred.token_ids) >= 0

    def test_predict_bigram(self):
        lm, pairs = _build_model()
        ngp = NgramPredictor(lm=lm)
        ngp.learn_ngrams([q + a for q, a in pairs])
        q_ids = lm.tokenizer.encode("who are you", add_bos=True, add_eos=True) + [1]
        pred = ngp.predict_bigram(q_ids)
        assert isinstance(pred, NgramPrediction)

    def test_cold_start(self):
        lm, _ = _build_model()
        ngp = NgramPredictor(lm=lm)
        pred = ngp.predict_next([0] * 100)  # unknown context
        assert pred.confidence >= 0.0


# ================================================================ MULTI-SCALE
class TestMultiScaleEncoder:
    def test_push_returns_result(self):
        lm, _ = _build_model()
        mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
        result = mse.push(5, lm.tokenizer)
        assert isinstance(result, MultiScaleResult)
        assert result.state_hv.D == lm.config.D

    def test_peek_empty(self):
        lm, _ = _build_model()
        mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
        result = mse.peek(lm.tokenizer)
        assert result is None  # empty at start

    def test_peek_after_push(self):
        lm, _ = _build_model()
        mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
        mse.push(5, lm.tokenizer)
        result = mse.peek(lm.tokenizer)
        assert result is not None

    def test_reset(self):
        lm, _ = _build_model()
        mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
        mse.push(5, lm.tokenizer)
        mse.reset()
        result = mse.peek(lm.tokenizer)
        assert result is None  # reset clears state


# ================================================================ RAG
class TestNativeRAG:
    def test_retrieve_context(self):
        lm, _ = _build_model()
        rag = NativeRAG(lm=lm, top_k=3)
        from palimseste.hv import bind, random_hv
        q = random_hv(lm.config.D)
        boosted, sims = rag.retrieve_context(q)
        assert boosted.D == lm.config.D

    def test_generate_with_rag(self):
        lm, _ = _build_model()
        rag = NativeRAG(lm=lm, top_k=3)
        result = rag.generate_with_rag("who are you", max_new_tokens=20)
        assert isinstance(result, RAGResult)
        assert len(result.text) >= 0


# ================================================================ REFINER
class TestIterativeRefiner:
    def test_refine_response(self):
        lm, _ = _build_model()
        refiner = IterativeRefiner(lm=lm, max_passes=2)
        result = refiner.refine_response("who are you", max_new_tokens=20)
        assert isinstance(result, RefinementResult)
        assert result.n_passes == 2

    def test_empty_question(self):
        lm, _ = _build_model()
        refiner = IterativeRefiner(lm=lm, max_passes=1)
        result = refiner.refine_response("xyz123unknown", max_new_tokens=10)
        assert isinstance(result, RefinementResult)


# ================================================================ TEMPLATES
class TestTemplateExtractor:
    def test_extract(self):
        extractor = TemplateExtractor()
        templates = extractor.extract([
            "python is a language",
            "gravity is a force",
            "the cat is a mammal",
        ])
        assert len(templates) > 0
        assert any("is a" in t.pattern for t in templates)

    def test_match_template(self):
        extractor = TemplateExtractor()
        extractor.extract(["python is a language"])
        match = extractor.match_template("java is a language")
        assert match is not None

    def test_no_match(self):
        extractor = TemplateExtractor()
        match = extractor.match_template("xyz123 random text")
        assert match is None

    def test_fill_template(self):
        extractor = TemplateExtractor()
        result = extractor.fill_template("{0} is a {1}", ["python", "language"])
        assert result == "python is a language"

    def test_capital_pattern(self):
        extractor = TemplateExtractor()
        extractor.extract([
            "the capital of france is paris",
            "the capital of japan is tokyo",
        ])
        match = extractor.match_template("the capital of italy is rome")
        assert match is not None
        assert "capital" in match[0]


# ================================================================ INGESTION PLAN
class TestMassiveIngestion:
    def test_plan_structure(self):
        plan = massive_ingestion_plan()
        assert "priority_1" in plan
        assert "total_estimated" in plan
        assert "tokens" in plan["total_estimated"]

    def test_plan_has_sources(self):
        plan = massive_ingestion_plan()
        for key in ["priority_1", "priority_2", "priority_3", "priority_4"]:
            assert "source" in plan[key]
            assert "estimated_tokens" in plan[key]