palimseste/chat.py

"""PALIMPSESTE — Conversational layer: multi-turn memory, live learning, fuzzy matching.

The :class:`PalimpsesteForCausalLM` is a pure associative memory — it maps a
context hypervector to a next-token hypervector. This module wraps it with the
state needed for real conversation:

  - **Multi-turn memory**: the conversation history (previous Q/A turns) is
    folded into the context so the model can reference earlier exchanges.
    "comment as-tu dit?" works because the prior answer is in the context.
  - **Live learning** (``teach``): the user can teach a new Q/A pair at
    runtime. It's written to ``M`` in ``O(1)`` per token — no retraining, no
    gradient. The very next message can use it. This is PALIMPSESTE's defining
    feature: learning = writing.
  - **Fuzzy matching**: before generating, we scan the known questions in the
    corpus and find the closest by character similarity. If the user's input
    is a near-miss of a known question ("tu utilise" vs "tu utilises"), we
    substitute the canonical form so the Hamming-radius retrieval hits.
  - **Fallback**: if no known question is close enough and the associative
    retrieval returns nothing, we return a graceful "je ne sais pas" response
    instead of an empty string.

The conversation state is serializable alongside the model, so a dialogue
can be saved and resumed.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import math
import numpy as np

from .hv import HV, bind, similarity
from .lm import PalimpsesteForCausalLM, PalimpsesteConfig
from .tokenizer import BOS, EOS, PAD

__all__ = ["Conversation", "Turn", "FALLBACK_RESPONSE"]

#: Response when the model has no matching knowledge.
FALLBACK_RESPONSE = "sorry, I don't know the answer to that question. you can teach me with: teach that <question> = <answer>"


@dataclass
class Turn:
    """One conversational turn."""

    role: str  # "user" or "palimpseste"
    text: str


@dataclass
class Conversation:
    """A multi-turn conversation with live learning and fuzzy matching.

    Wraps a :class:`PalimpsesteForCausalLM` with conversational state.

    Parameters
    ----------
    model : PalimpsesteForCausalLM
        The trained (or fresh) language model.
    max_history_turns : int
        How many previous Q/A turns to fold into the context (default 3).
        Older turns are dropped to stay within the context window.
    fuzzy_threshold : float
        Character-level similarity (0..1) above which a user input is
        treated as a fuzzy match to a known question (default 0.75).
    learn_live : bool
        If True, unknown questions trigger a fallback that invites the user
        to teach the answer (default True). If False, just returns the
        fallback string.
    """

    model: PalimpsesteForCausalLM
    max_history_turns: int = 3
    fuzzy_threshold: float = 0.75
    learn_live: bool = True
    history: list[Turn] = field(default_factory=list)
    _known_questions: list[str] = field(default_factory=list)
    _question_to_answer: dict[str, str] = field(default_factory=dict)
    _turn_index: int = 0
    _entity_tracker: object | None = field(default=None)

    def __post_init__(self) -> None:
        self._index_known_questions()
        # Enable entity tracking for pronoun resolution
        from .evolution import EntityTracker
        self._entity_tracker = EntityTracker()

    # ----------------------------------------------------------- known Q indexing
    def _index_known_questions(self) -> None:
        """Extract known questions from the model's memory tags.

        Tags formatted ``qa{n}@{i}`` mark Q/A training traces. We reconstruct
        the question→answer map by scanning the tokenizer + memory.
        """
        self._known_questions = []
        self._question_to_answer = {}
        # We can't easily decode stored HVs back to text (the memory stores
        # HVs, not text), so we rely on the model's respond() to recover
        # answers for questions we know. The fuzzy matcher needs a list of
        # known questions — these must be registered via register_questions()
        # or teach().
        pass

    def register_questions(self, pairs: list[tuple[str, str]]) -> None:
        """Register a list of (question, answer) pairs for fuzzy matching.

        This does NOT train the model (use :meth:`teach` or
        ``model.train_on_qa_pairs`` for that). It just populates the fuzzy
        index so that near-miss user inputs can be matched to canonical forms.
        """
        for q, a in pairs:
            self._known_questions.append(q)
            self._question_to_answer[q] = a

    # ----------------------------------------------------------- fuzzy matching
    def _fuzzy_match(self, user_input: str) -> str | None:
        """Find the closest known question by character-level similarity.

        Uses a simple normalized edit-distance ratio (difflib). Returns the
        canonical question if similarity >= ``fuzzy_threshold``, else None.
        """
        if not self._known_questions:
            return None
        import difflib
        best_q = None
        best_score = 0.0
        for q in self._known_questions:
            score = difflib.SequenceMatcher(None, user_input.lower().strip(),
                                            q.lower().strip()).ratio()
            if score > best_score:
                best_score = score
                best_q = q
        if best_score >= self.fuzzy_threshold:
            return best_q
        return None

    # ----------------------------------------------------------- multi-turn context
    def _build_context_ids(self, user_input: str) -> list[int]:
        """Build the token-id context including conversation history.

        The context is::

            [BOS hist_q1 EOS BOS hist_a1 EOS] ...
            [BOS user_input EOS BOS]

        Previous turns are included so the model can reference them. The
        current question is **always** fully included — history is truncated
        to leave room, never the question itself.
        """
        tok = self.model.tokenizer
        assert tok is not None
        window = self.model.config.context_window

        # current question + answer-start BOS (must always fit)
        q_ids = tok.encode(user_input, add_bos=True, add_eos=True)
        current = q_ids + [BOS]

        # budget for history = window - len(current)
        history_budget = window - len(current)
        if history_budget <= 0:
            # question alone fills the window — no history
            return current[-window:]

        # fold in recent history, oldest first, until budget exhausted
        ctx: list[int] = []
        recent = self.history[-self.max_history_turns * 2:]  # Q/A pairs
        for turn in recent:
            ids = tok.encode(turn.text, add_bos=True, add_eos=True)
            if len(ctx) + len(ids) > history_budget:
                break
            ctx.extend(ids)

        return ctx + current

    # ----------------------------------------------------------- respond (chat)
    def respond(self, user_input: str, max_new_tokens: int = 200,
                temperature: float | None = None, seed: int | None = None) -> str:
        """Respond to a user message, with multi-turn memory and fuzzy matching.

        Strategy (associative memory is not a transformer — context pollution
        from history can break retrieval, so we try clean first):

        1. Entity tracking: resolve pronouns ("it", "that") using context.
        2. Fuzzy-match the input to a known question (handles typos).
        3. Try generating with a **clean** context (question only, no history).
        4. If that fails, retry **with** history.
        5. If both fail, return a fallback that invites the user to teach.
        """
        assert self.model.tokenizer is not None

        # Entity tracking: resolve pronouns from conversation history
        resolved_input = user_input
        if self.history and self._entity_tracker:
            resolved_input = self._entity_tracker.update(user_input)

        # fuzzy match: if the user's input is close to a known question,
        # use the canonical form for better retrieval
        matched = self._fuzzy_match(resolved_input)
        effective_input = matched if matched else resolved_input

        tok = self.model.tokenizer
        # clean context: just the question + BOS (no history)
        q_ids = tok.encode(effective_input, add_bos=True, add_eos=True)
        clean_ctx = q_ids + [BOS]

        # try clean retrieval first
        logits = self.model._logits(clean_ctx)
        is_cold = (logits.max() == 0.0 and logits.min() == 0.0)

        response = ""
        if not is_cold:
            response = self.model.generate_from_ids(
                clean_ctx, max_new_tokens=max_new_tokens,
                temperature=temperature, seed=seed, stop_on_eos=True,
            )

        # if clean retrieval failed, try with history (multi-turn context)
        if not response.strip() and self.history:
            ctx_ids = self._build_context_ids(effective_input)
            logits = self.model._logits(ctx_ids)
            is_cold = (logits.max() == 0.0 and logits.min() == 0.0)
            if not is_cold:
                response = self.model.generate_from_ids(
                    ctx_ids, max_new_tokens=max_new_tokens,
                    temperature=temperature, seed=seed, stop_on_eos=True,
                )

        if not response.strip():
            response = FALLBACK_RESPONSE

        # record the turn
        self.history.append(Turn(role="user", text=user_input))
        self.history.append(Turn(role="palimpseste", text=response))
        self._turn_index += 1
        return response

    # ----------------------------------------------------------- live learning
    def teach(self, question: str, answer: str) -> str:
        """Teach a new Q/A pair at runtime. O(1) per token. Immediately usable.

        This is PALIMPSESTE's defining feature: learning is writing. The new
        pair is appended to ``M`` and becomes retrievable on the very next
        message. No retraining, no gradient, no epoch.

        Returns a confirmation message.
        """
        assert self.model.tokenizer is not None
        # expand vocab if needed — ONLY for CharTokenizer.
        # BPETokenizer.build_vocab() would re-train and destroy merges.
        tok_type = type(self.model.tokenizer).__name__
        if tok_type == "CharTokenizer":
            self.model.tokenizer.build_vocab(question)
            self.model.tokenizer.build_vocab(answer)
        # write the episode to M
        n = self._train_one_qa(question, answer)
        # register for fuzzy matching
        self._known_questions.append(question)
        self._question_to_answer[question] = answer
        return f"thanks! I learned this answer ({n} tokens written). I can now respond to this question."

    def _train_one_qa(self, question: str, answer: str) -> int:
        """Write a single Q/A episode to M. O(1) per token."""
        tok = self.model.tokenizer
        assert tok is not None
        q_ids = tok.encode(question, add_bos=True, add_eos=True)
        a_ids = tok.encode(answer, add_bos=True, add_eos=True)
        episode = q_ids + a_ids
        n = 0
        for i in range(1, len(episode)):
            ctx = episode[:i]
            target = episode[i]
            s = self.model._state_hv(ctx)
            q_addr = bind(self.model._self_hv, s)
            target_hv = tok.token_hv(target)
            self.model.mem.write(q_addr, target_hv, weight=1.0,
                                 tag=f"live@{self._turn_index}:{i}")
            n += 1
        return n

    # ----------------------------------------------------------- parse teach command
    def try_parse_teach(self, user_input: str) -> bool:
        """Check if the user input is a teach command.

        Supported formats:
          - ``teach that <question> = <answer>``
          - ``teach <question> = <answer>``

        If matched, teaches the pair and returns True. Otherwise returns false.
        """
        text = user_input.strip()
        prefixes = ["teach that ", "teach ", "learn that "]
        for prefix in prefixes:
            if text.lower().startswith(prefix):
                rest = text[len(prefix):]
                if "=" in rest:
                    q, a = rest.split("=", 1)
                    q = q.strip()
                    a = a.strip()
                    if q and a:
                        msg = self.teach(q, a)
                        # record the interaction in history
                        self.history.append(Turn(role="user", text=user_input))
                        self.history.append(Turn(role="palimpseste", text=msg))
                        return True
        return False

    # ----------------------------------------------------------- streaming
    def respond_streaming(self, user_input: str, max_new_tokens: int = 200,
                          temperature: float | None = None,
                          seed: int | None = None):
        """Generator that yields tokens one at a time as they are generated.

        Uses the same clean-first strategy as :meth:`respond`.

        Yields (token_text, token_id) tuples. Useful for a responsive chat UI.
        """
        assert self.model.tokenizer is not None
        tok = self.model.tokenizer
        matched = self._fuzzy_match(user_input)
        effective_input = matched if matched else user_input

        T = self.model.config.temperature if temperature is None else temperature
        rng = np.random.default_rng(seed) if seed is not None else self.model.rng

        # clean context first
        q_ids = tok.encode(effective_input, add_bos=True, add_eos=True)
        ctx_ids = q_ids + [BOS]

        # check if clean retrieval will work
        logits = self.model._logits(ctx_ids)
        is_cold = (logits.max() == 0.0 and logits.min() == 0.0)
        if is_cold and self.history:
            # fallback to multi-turn context
            ctx_ids = self._build_context_ids(effective_input)
            logits = self.model._logits(ctx_ids)
            is_cold = (logits.max() == 0.0 and logits.min() == 0.0)

        if is_cold:
            # no match at all — yield fallback
            self.history.append(Turn(role="user", text=user_input))
            self.history.append(Turn(role="palimpseste", text=FALLBACK_RESPONSE))
            yield (FALLBACK_RESPONSE, -1)
            return

        ids = list(ctx_ids)
        response_parts: list[str] = []
        for _ in range(max_new_tokens):
            logits = self.model._logits(ids)
            if logits.max() == 0.0 and logits.min() == 0.0:
                break
            if T <= 1e-6:
                next_id = int(np.argmax(logits))
            else:
                z = logits / T
                z -= z.max()
                e = np.exp(z)
                probs = e / e.sum()
                probs[PAD] = 0.0
                s = probs.sum()
                if s <= 0:
                    break
                probs /= s
                next_id = int(rng.choice(len(probs), p=probs))
            ids.append(next_id)
            if next_id == EOS:
                break
            token_text = tok.id2char[next_id] if 0 <= next_id < len(tok.id2char) else ""
            if token_text in ("<pad>", "<bos>", "<eos>", "<unk>"):
                continue
            response_parts.append(token_text)
            yield (token_text, next_id)

        response = "".join(response_parts)
        if not response.strip():
            response = FALLBACK_RESPONSE
        self.history.append(Turn(role="user", text=user_input))
        self.history.append(Turn(role="palimpseste", text=response))

    # ----------------------------------------------------------- utilities
    def reset(self) -> None:
        """Clear conversation history (start a new dialogue). Does NOT clear M."""
        self.history.clear()

    @property
    def turn_count(self) -> int:
        return self._turn_index

    def get_transcript(self) -> str:
        """Return the full conversation as readable text."""
        lines = []
        for turn in self.history:
            speaker = "You" if turn.role == "user" else "Palimpseste"
            lines.append(f"{speaker}: {turn.text}")
        return "\n".join(lines)