palimseste/cognitive.py

"""PALIMPSESTE — Cognitive layer: confidence, self-correction, curiosity,
auto-chaining, and explanation traces.

All five features are **inference-time only** — they wrap the existing
:class:`Conversation` and :class:`Reasoner` without modifying the model's
memory structure. Everything works on a model that was already trained;
no retraining is needed.

Features
--------
1. **Confidence scoring**: every response comes with a confidence score
   (0.0–1.0) based on the Hamming similarity of the best match. Low
   confidence → "I'm not sure" instead of a confident wrong answer.

2. **Self-correction**: when the user says "no, the answer is X", the model
   learns the correction via O(1) memory write and uses it immediately.
   It also marks the old answer as superseded.

3. **Curiosity loop**: when the model doesn't know an answer, instead of
   saying "sorry", it asks the user to teach it: "I don't know X. Can you
   teach me?" — creating a bidirectional learning loop.

4. **Auto-chaining**: when two facts A→B and B→C are taught, the model
   automatically discovers A→C by running the reasoner in the background
   and writing the composed fact to memory. Next time A is asked, C is
   retrieved directly — the chain is consolidated.

5. **Explanation trace**: every response comes with an explanation of
   *why* the model answered: "I matched your question to 'X' (similarity:
   0.92) which I was taught at turn 3" or "I chained: A→B→C".
"""

from __future__ import annotations

from dataclasses import dataclass, field
import re
from typing import Optional

from .hv import similarity, bind
from .chat import Conversation, Turn, FALLBACK_RESPONSE
from .reasoning import Reasoner, ChainResult

__all__ = [
    "CognitiveResponse",
    "CognitiveAgent",
    "CONFIDENCE_THRESHOLD",
    "CORRECTION_PATTERNS",
]

#: Below this confidence, the model expresses doubt.
CONFIDENCE_THRESHOLD = 0.15

#: Patterns that indicate the user is correcting a previous answer.
#: Handles both straight (') and curly (') apostrophes.
_APOS = r"['\u2019]"
CORRECTION_PATTERNS = [
    (re.compile(rf"(?:non|no)[,\s]+(?:c{_APOS}?est|it is|the answer is)\s+(.+)", re.I), "direct"),
    (re.compile(rf"(?:faux|wrong)[,\s]+(?:c{_APOS}?est|it is|the answer is)\s+(.+)", re.I), "direct"),
    (re.compile(rf"(?:actually|en fait)[,\s]+(?:c{_APOS}?est|it is|the answer is)\s+(.+)", re.I), "direct"),
    (re.compile(rf"(?:la bonne r[ée]ponse est|the correct answer is)\s+(.+)", re.I), "clean"),
    (re.compile(rf"(?:c{_APOS}?est pas .+?,\s*c{_APOS}?est|it{_APOS}?s not .+?,\s*it{_APOS}?s)\s+(.+)", re.I), "contrast"),
]


@dataclass
class CognitiveResponse:
    """A response with full cognitive metadata."""

    text: str
    confidence: float
    explanation: str
    source: str  # "direct" | "chained" | "corrected" | "curiosity" | "fallback"
    chain: ChainResult | None = None
    corrected_answer: str | None = None


@dataclass
class CognitiveAgent:
    """A conversational agent with confidence, self-correction, curiosity,
    auto-chaining, and explanation traces.

    Wraps a :class:`Conversation` + :class:`Reasoner`. All features are
    inference-time — no retraining needed.

    Parameters
    ----------
    conv : Conversation
        The base conversation (with a trained model).
    reasoner : Reasoner | None
        The fact-chaining reasoner. If None, created automatically.
    confidence_threshold : float
        Below this confidence, express doubt (default 0.15).
    enable_curiosity : bool
        If True, unknown questions trigger a "teach me" prompt (default True).
    enable_auto_chain : bool
        If True, teaching new facts triggers background chaining (default True).
    """

    conv: Conversation
    reasoner: Reasoner | None = None
    confidence_threshold: float = CONFIDENCE_THRESHOLD
    enable_curiosity: bool = True
    enable_auto_chain: bool = True
    _last_question: str = ""
    _last_answer: str = ""
    _last_confidence: float = 0.0
    _corrections: dict[str, str] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.reasoner is None:
            self.reasoner = Reasoner(conv=self.conv, max_hops=3, min_fragment_len=8)

    # ----------------------------------------------------------- respond
    def respond(self, user_input: str, max_new_tokens: int = 200,
                temperature: float | None = None,
                seed: int | None = None) -> CognitiveResponse:
        """Respond with full cognitive metadata."""
        # check for self-correction first
        correction = self._detect_correction(user_input)
        if correction is not None:
            return self._handle_correction(correction, user_input)

        # try direct response with confidence
        self.conv.reset()
        raw_answer = self.conv.respond(user_input, max_new_tokens=max_new_tokens,
                                       temperature=temperature, seed=seed)
        confidence = self._compute_confidence(user_input, raw_answer)

        # if direct retrieval succeeded
        if raw_answer and raw_answer != FALLBACK_RESPONSE and raw_answer.strip():
            explanation = self._explain_direct(user_input, confidence)
            self._last_question = user_input
            self._last_answer = raw_answer
            self._last_confidence = confidence

            # if confidence is low, express doubt
            if confidence < self.confidence_threshold:
                raw_answer = f"(I'm not sure) {raw_answer}"
                explanation += " [low confidence — expressed doubt]"

            return CognitiveResponse(
                text=raw_answer,
                confidence=confidence,
                explanation=explanation,
                source="direct",
            )

        # direct failed — try chaining
        self.conv.reset()
        answer, chain = self.reasoner.respond(user_input, temperature=temperature, seed=seed)
        if chain and chain.success:
            explanation = self._explain_chain(chain)
            self._last_question = user_input
            self._last_answer = answer
            self._last_confidence = 0.7  # chained answers are moderately confident
            return CognitiveResponse(
                text=answer,
                confidence=0.7,
                explanation=explanation,
                source="chained",
                chain=chain,
            )

        # all failed — curiosity loop
        if self.enable_curiosity:
            curiosity_msg = self._curiosity_response(user_input)
            self._last_question = user_input
            self._last_answer = ""
            return CognitiveResponse(
                text=curiosity_msg,
                confidence=0.0,
                explanation="No match found. Asking user to teach.",
                source="curiosity",
            )

        return CognitiveResponse(
            text=FALLBACK_RESPONSE,
            confidence=0.0,
            explanation="No match found.",
            source="fallback",
        )

    # ----------------------------------------------------------- confidence
    def _compute_confidence(self, question: str, answer: str) -> float:
        """Compute confidence from the best Hamming similarity of the match."""
        if not answer or answer == FALLBACK_RESPONSE:
            return 0.0
        lm = self.conv.model
        tok = lm.tokenizer
        if tok is None:
            return 0.5
        # reconstruct the query and check retrieval similarity
        q_ids = tok.encode(question, add_bos=True, add_eos=True)
        ctx = q_ids + [1]  # BOS
        s = lm._state_hv(ctx)
        q = bind(lm._self_hv, s)
        ret = lm.phi.retrieve(lm.mem, q)
        if not ret.sims:
            return 0.0
        best_sim = max(ret.sims)
        # map [-1, 1] to [0, 1]
        return (best_sim + 1.0) / 2.0

    # ----------------------------------------------------------- self-correction
    def _detect_correction(self, user_input: str) -> str | None:
        """Check if the user is correcting the last answer. Returns the corrected answer."""
        if not self._last_question:
            return None
        for pattern, kind in CORRECTION_PATTERNS:
            m = pattern.match(user_input.strip())
            if m:
                corrected = m.group(1).strip().rstrip(".!?")
                return corrected
        return None

    def _handle_correction(self, corrected_answer: str,
                           user_input: str) -> CognitiveResponse:
        """Learn the correction via O(1) write and acknowledge."""
        question = self._last_question
        # teach the corrected answer
        self.conv.teach(question, corrected_answer)
        # record the correction
        self._corrections[question] = corrected_answer

        # auto-chain if enabled
        chain_note = ""
        if self.enable_auto_chain:
            chains = self._try_auto_chain(question, corrected_answer)
            if chains:
                chain_note = f" I also discovered {len(chains)} new connection(s)."

        response = f"thanks for the correction! I learned that the answer to \"{question[:40]}\" is \"{corrected_answer}\".{chain_note}"

        self._last_question = user_input
        self._last_answer = corrected_answer
        return CognitiveResponse(
            text=response,
            confidence=1.0,
            explanation=f"User corrected the answer to '{question}'. "
                        f"Learned via O(1) memory write. Old answer superseded.",
            source="corrected",
            corrected_answer=corrected_answer,
        )

    # ----------------------------------------------------------- curiosity
    def _curiosity_response(self, question: str) -> str:
        """Generate a curiosity-driven response asking the user to teach."""
        # extract the key concept from the question
        concept = question.strip().rstrip("?")
        if len(concept) > 60:
            concept = concept[:60] + "..."
        return (f"I don't know the answer to this question. "
                f"can you teach me? type: teach that {concept} = <your answer>")

    # ----------------------------------------------------------- auto-chaining
    def _try_auto_chain(self, new_question: str, new_answer: str) -> list[str]:
        """When a new fact is taught, try to discover chains.

        If the new answer matches a known question, or the new question
        contains a known answer, run the reasoner to discover composed facts.
        Returns a list of discovered chain descriptions.
        """
        discovered = []
        known_qs = self.conv._known_questions
        known_as = self.conv._question_to_answer

        # Case 1: the new answer is itself a known question
        # e.g., taught "what is the capital of france" → "paris"
        # if "paris" is a known question, we might chain further
        for kq in known_qs:
            if kq == new_question:
                continue
            if new_answer.lower() in kq.lower() and len(new_answer) > 3:
                # the new answer appears in a known question — try chaining
                composed = kq.replace(new_answer.lower(), new_answer.lower())
                result = self.reasoner.try_chain(composed)
                if result.success:
                    desc = f"{new_question} → {new_answer} → {kq} → {result.answer}"
                    discovered.append(desc)
                    # write the composed fact to memory
                    self.conv.teach(composed, result.answer)

        # Case 2: the new question contains a known answer
        for kq, ka in known_as.items():
            if ka.lower() in new_question.lower() and len(ka) > 3:
                # the known answer appears in the new question — try chaining
                result = self.reasoner.try_chain(new_question)
                if result.success:
                    desc = f"{kq} → {ka} → {new_question} → {result.answer}"
                    discovered.append(desc)
                    self.conv.teach(new_question, result.answer)

        return discovered

    # ----------------------------------------------------------- explanation
    def _explain_direct(self, question: str, confidence: float) -> str:
        """Explain why a direct answer was given."""
        # find the best matching known question
        matched = self.conv._fuzzy_match(question)
        if matched:
            return (f"Matched your question to '{matched}' "
                    f"(confidence: {confidence:.0%}). Retrieved from memory.")
        return f"Direct associative retrieval (confidence: {confidence:.0%})."

    def _explain_chain(self, chain: ChainResult) -> str:
        """Explain a chained answer."""
        steps_desc = []
        for i, step in enumerate(chain.steps, 1):
            steps_desc.append(
                f"hop {i}: '{step.sub_question}' → '{step.sub_answer}' "
                f"→ resolved to '{step.resolved_question}'"
            )
        return "Chained reasoning:\n  " + "\n  ".join(steps_desc)

    # ----------------------------------------------------------- teach (with auto-chain)
    def teach(self, question: str, answer: str) -> CognitiveResponse:
        """Teach a new fact, with auto-chaining if enabled."""
        self.conv.teach(question, answer)

        chain_note = ""
        if self.enable_auto_chain:
            chains = self._try_auto_chain(question, answer)
            if chains:
                chain_note = f" Auto-discovered {len(chains)} connection(s):\n"
                for c in chains:
                    chain_note += f"  → {c}\n"

        explanation = (f"Learned '{question}' → '{answer}' via O(1) memory write. "
                       f"Immediately retrievable." + chain_note)

        return CognitiveResponse(
            text=f"thanks! I learned this answer.{chain_note}",
            confidence=1.0,
            explanation=explanation,
            source="corrected",
        )

    # ----------------------------------------------------------- utilities
    def reset(self) -> None:
        """Clear conversation state (not memory)."""
        self.conv.reset()
        self._last_question = ""
        self._last_answer = ""
        self._last_confidence = 0.0

    @property
    def n_corrections(self) -> int:
        return len(self._corrections)

    def get_transcript(self) -> str:
        return self.conv.get_transcript()