tests/test_conversation.py

"""Tests for the Conversation layer (``palimseste.chat``).

Verifies:
  - multi-turn memory: the conversation history is folded into context
  - live learning (teach): a new Q/A pair is instantly retrievable
  - fuzzy matching: typos / near-misses match to known questions
  - fallback: unseen questions return a graceful response, not empty string
  - streaming: respond_streaming yields tokens
  - teach command parsing (teach that ... = ...)
  - conversation reset clears history but not M
  - transcript generation
"""

from __future__ import annotations

import numpy as np
import pytest

from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.chat import Conversation, Turn, FALLBACK_RESPONSE


def _model(D=5000, radius=0, ctx_window=64):
    cfg = PalimpsesteConfig(D=D, context_window=ctx_window, kernel_radius=radius,
                            temperature=0.3)
    lm = PalimpsesteForCausalLM(config=cfg)
    pairs = [
        ("hello", "hi i am palimpseste"),
        ("who are you", "i am palimpseste a hypervectorial cortex"),
        ("do you use a gpu", "no i do not use a gpu"),
    ]
    full = "".join(q + a for q, a in pairs)
    lm.build_tokenizer(full)
    lm.train_on_qa_pairs(pairs)
    return lm, pairs


# ----------------------------------------------------------------- fuzzy matching
class TestFuzzyMatching:
    def test_exact_match(self):
        lm, pairs = _model()
        conv = Conversation(model=lm)
        conv.register_questions(pairs)
        assert conv._fuzzy_match("hello") == "hello"

    def test_typo_matches(self):
        lm, pairs = _model()
        conv = Conversation(model=lm, fuzzy_threshold=0.7)
        conv.register_questions(pairs)
        # "do you use gpu" (missing "a") should match "do you use a gpu"
        matched = conv._fuzzy_match("do you use gpu")
        assert matched == "do you use a gpu"

    def test_no_match_returns_none(self):
        lm, pairs = _model()
        conv = Conversation(model=lm, fuzzy_threshold=0.9)
        conv.register_questions(pairs)
        assert conv._fuzzy_match("xyz123abc") is None

    def test_empty_known_questions(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        assert conv._fuzzy_match("anything") is None


# ----------------------------------------------------------------- fallback
class TestFallback:
    def test_unseen_question_returns_fallback(self):
        lm, _ = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        resp = conv.respond("xyz123abc", temperature=0.0, seed=0)
        assert resp == FALLBACK_RESPONSE

    def test_seen_question_returns_answer(self):
        lm, pairs = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        conv.register_questions(pairs)
        resp = conv.respond("hello", temperature=0.0, seed=0)
        assert resp == "hi i am palimpseste"
        assert resp != FALLBACK_RESPONSE


# ----------------------------------------------------------------- live learning
class TestLiveLearning:
    def test_teach_then_respond(self):
        lm, _ = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        # before teaching: unknown question -> fallback
        resp_before = conv.respond("what is the capital of france",
                                   temperature=0.0, seed=0)
        assert resp_before == FALLBACK_RESPONSE
        # teach the new Q/A
        msg = conv.teach("what is the capital of france", "paris")
        assert "learned" in msg.lower()
        # after teaching: should retrieve the answer
        # reset history so the fallback context doesn't interfere
        conv.reset()
        resp_after = conv.respond("what is the capital of france",
                                  temperature=0.0, seed=0)
        assert resp_after == "paris"

    def test_teach_grows_memory(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        n_before = len(lm.mem)
        conv.teach("new question", "new answer")
        assert len(lm.mem) > n_before

    def test_teach_registers_for_fuzzy(self):
        lm, _ = _model()
        conv = Conversation(model=lm, fuzzy_threshold=0.8)
        conv.teach("what time is it", "i do not know")
        # a near-miss should now fuzzy-match
        matched = conv._fuzzy_match("what time is it now")
        assert matched is not None


# ----------------------------------------------------------------- teach command
class TestTeachCommand:
    def test_parse_teach_that(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        assert conv.try_parse_teach("teach that color = blue") is True
        assert conv._known_questions[-1] == "color"

    def test_parse_teach_english(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        assert conv.try_parse_teach("teach color = blue") is True

    def test_parse_teach_no_equals(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        assert conv.try_parse_teach("teach that color blue") is False

    def test_parse_teach_no_prefix(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        assert conv.try_parse_teach("color = blue") is False

    def test_parse_teach_empty_answer(self):
        lm, _ = _model()
        conv = Conversation(model=lm)
        assert conv.try_parse_teach("teach that color = ") is False


# ----------------------------------------------------------------- multi-turn
class TestMultiTurn:
    def test_history_recorded(self):
        lm, pairs = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        conv.register_questions(pairs)
        conv.respond("hello", temperature=0.0, seed=0)
        assert len(conv.history) == 2  # one Q + one A
        assert conv.history[0].role == "user"
        assert conv.history[0].text == "hello"
        assert conv.history[1].role == "palimpseste"

    def test_reset_clears_history(self):
        lm, pairs = _model()
        conv = Conversation(model=lm)
        conv.register_questions(pairs)
        conv.respond("hello", temperature=0.0, seed=0)
        assert len(conv.history) > 0
        conv.reset()
        assert len(conv.history) == 0

    def test_reset_does_not_clear_memory(self):
        lm, pairs = _model()
        conv = Conversation(model=lm)
        n_before = len(lm.mem)
        conv.respond("hello", temperature=0.0, seed=0)
        conv.reset()
        assert len(lm.mem) == n_before  # M untouched

    def test_turn_count(self):
        lm, pairs = _model()
        conv = Conversation(model=lm)
        conv.register_questions(pairs)
        assert conv.turn_count == 0
        conv.respond("hello", temperature=0.0, seed=0)
        assert conv.turn_count == 1
        conv.respond("who are you", temperature=0.0, seed=0)
        assert conv.turn_count == 2


# ----------------------------------------------------------------- streaming
class TestStreaming:
    def test_streaming_yields_tokens(self):
        lm, pairs = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        conv.register_questions(pairs)
        tokens = list(conv.respond_streaming("hello", temperature=0.0, seed=0))
        assert len(tokens) > 0
        # reconstruct the text
        text = "".join(t for t, _ in tokens)
        assert "hi" in text

    def test_streaming_cold_start(self):
        lm, _ = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        tokens = list(conv.respond_streaming("xyz123", temperature=0.0, seed=0))
        # should yield the fallback response
        text = "".join(t for t, _ in tokens)
        assert "sorry" in text.lower() or len(text) > 0


# ----------------------------------------------------------------- transcript
class TestTranscript:
    def test_transcript_format(self):
        lm, pairs = _model(radius=0)
        conv = Conversation(model=lm, fuzzy_threshold=0.99)
        conv.register_questions(pairs)
        conv.respond("hello", temperature=0.0, seed=0)
        transcript = conv.get_transcript()
        assert "You:" in transcript
        assert "Palimpseste:" in transcript
        assert "hello" in transcript


# ----------------------------------------------------------------- context building
class TestContextBuilding:
    def test_context_includes_history(self):
        lm, pairs = _model(radius=0, ctx_window=128)
        conv = Conversation(model=lm, max_history_turns=2, fuzzy_threshold=0.99)
        conv.register_questions(pairs)
        # first turn
        conv.respond("hello", temperature=0.0, seed=0)
        # build context for second turn — should include history
        ctx = conv._build_context_ids("who are you")
        # context should be longer than just the question
        q_only = lm.tokenizer.encode("who are you", add_bos=True, add_eos=True) + [1]  # +BOS
        assert len(ctx) > len(q_only)