palimseste/reasoning.py

"""PALIMPSESTE — Fact chaining (multi-step reasoning).

The core limitation of an associative memory is that it can only retrieve
what it was explicitly taught. If the model knows:

    A → B   ("who won the world cup 2018?" → "france")
    B → C   ("what is the capital of france?" → "paris")

...it CANNOT answer:

    A → C   ("what is the capital of the country that won the world cup 2018?")

...because that question was never stored. This module bridges that gap.

How it works
------------
When :meth:`respond` hits a fallback (no associative match), the reasoner:

1. **Decomposes** the question: scans known questions for fragments that
   appear as substrings of the user's input. E.g., "what is the capital of
   france?" is a known question whose answer ("france") appears as a fragment.

2. **Sub-queries**: for each known question that overlaps with the input,
   retrieves its answer via normal :meth:`respond`.

3. **Composes**: replaces the overlapping fragment in the original question
   with the retrieved answer, creating a *resolved* question. E.g.:
   - Input: "what is the capital of the country that won the world cup 2018?"
   - Known: "who won the world cup 2018?" → "france"
   - Resolved: "what is the capital of france?" → "paris"

4. **Re-queries**: sends the resolved question through :meth:`respond` again.
   If it hits, the model has *reasoned* across two facts it was never
   explicitly taught to chain.

This is a two-hop reasoning chain. The same mechanism extends to N hops
recursively: each resolution can trigger another decomposition.

What this is NOT
----------------
This is not symbolic logic or a theorem prover. It's associative chaining —
the model finds overlapping concepts between its stored knowledge and the
novel question, resolves them one hop at a time, and re-queries. It's the
minimal mechanism that turns a memory into a reasoner.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import re
from .chat import Conversation, FALLBACK_RESPONSE

__all__ = ["Reasoner", "ChainStep", "ChainResult"]


@dataclass
class ChainStep:
    """One hop in a reasoning chain."""

    sub_question: str
    sub_answer: str
    resolved_question: str


@dataclass
class ChainResult:
    """Outcome of a reasoning attempt."""

    success: bool
    answer: str
    steps: list[ChainStep]
    original_question: str

    @property
    def n_hops(self) -> int:
        return len(self.steps)


@dataclass
class Reasoner:
    """Multi-step fact chaining for the :class:`Conversation`.

    Wraps a :class:`Conversation` and adds the ability to chain facts across
    hops when the direct associative retrieval fails.

    Parameters
    ----------
    conv : Conversation
        The conversation to enhance with reasoning.
    max_hops : int
        Maximum number of chaining hops before giving up (default 3).
    min_fragment_len : int
        Minimum length of a question fragment to be considered for overlap
        (default 8 chars). Shorter fragments cause false matches.
    """

    conv: Conversation
    max_hops: int = 3
    min_fragment_len: int = 8

    # ----------------------------------------------------------- reasoning
    def try_chain(self, question: str) -> ChainResult:
        """Attempt to answer a question by chaining known facts.

        Returns a :class:`ChainResult`. If ``success`` is True, ``answer``
        contains the chained answer. ``steps`` contains the reasoning trace.
        """
        steps: list[ChainStep] = []
        current_question = question.lower().strip()

        for hop in range(self.max_hops):
            # find a known question that overlaps with the current question
            best_match = self._find_subquestion(current_question)
            if best_match is None:
                break

            sub_q, fragment = best_match
            # get the answer to the sub-question
            sub_answer = self._query_known(sub_q)
            if not sub_answer or sub_answer == FALLBACK_RESPONSE:
                break

            # resolve: replace the fragment with the answer, then clean up
            resolved = self._resolve_question(current_question, fragment, sub_answer, sub_q)

            steps.append(ChainStep(
                sub_question=sub_q,
                sub_answer=sub_answer,
                resolved_question=resolved,
            ))

            # try to answer the resolved question
            final_answer = self._query_known(resolved)
            if final_answer and final_answer != FALLBACK_RESPONSE:
                return ChainResult(
                    success=True,
                    answer=final_answer,
                    steps=steps,
                    original_question=question,
                )

            # otherwise, try to chain further with the resolved question
            current_question = resolved

        return ChainResult(
            success=False,
            answer=FALLBACK_RESPONSE,
            steps=steps,
            original_question=question,
        )

    def _resolve_question(self, question: str, fragment: str,
                          answer: str, sub_question: str) -> str:
        """Replace the overlapping fragment with the answer, then clean up.

        The cleanup removes filler phrases like "the country that", "the city
        that", "the person who", etc. that connect the sub-question to the
        rest of the query. This produces a natural resolved question that
        is more likely to match a known question.
        """
        # Step 1: replace the fragment with the answer
        resolved = question.replace(fragment, answer)

        # Step 2: remove common filler phrases that connect sub-questions
        fillers = [
            r'\bthe country that\b',
            r'\bthe city that\b',
            r'\bthe person who\b',
            r'\bthe one who\b',
            r'\bthe place where\b',
            r'\bthe thing that\b',
            r'\bthat won\b',
            r'\bthat is\b',
            r'\bwho won\b',
        ]
        for filler in fillers:
            resolved = re.sub(filler, '', resolved, flags=re.IGNORECASE)

        # Step 3: clean up extra spaces
        resolved = re.sub(r'\s+', ' ', resolved).strip()

        # Step 4: try to match a known question pattern
        # If the resolved question looks like "what is the capital of  france"
        # (double space from filler removal), clean it
        resolved = re.sub(r'\s+of\s+', ' of ', resolved)
        resolved = re.sub(r'\s+', ' ', resolved).strip()

        return resolved

    # ----------------------------------------------------------- helpers
    def _find_subquestion(self, question: str) -> tuple[str, str] | None:
        """Find a known question whose text overlaps with ``question``.

        Returns ``(known_question, overlapping_fragment)`` or None.
        The overlap is a substring of the known question that appears in
        ``question`` and is at least ``min_fragment_len`` chars long.
        """
        known_qs = self.conv._known_questions
        if not known_qs:
            return None

        best = None
        best_len = 0
        for kq in known_qs:
            kq_lower = kq.lower().strip()
            if len(kq_lower) < self.min_fragment_len:
                continue
            # check if kq is a substring of question, or a significant part of it
            if kq_lower in question:
                # the known question itself appears in the input
                return (kq, kq_lower)
            # try progressively shorter fragments of the known question
            # (from the end, since answers usually follow questions)
            for frag_len in range(min(len(kq_lower), 40), self.min_fragment_len - 1, -2):
                fragment = kq_lower[-frag_len:]
                if fragment in question and len(fragment) > best_len:
                    best = (kq, fragment)
                    best_len = len(fragment)
        return best

    def _query_known(self, question: str) -> str:
        """Query the model for a question, return the answer or fallback."""
        # use a fresh conversation context to avoid history pollution
        saved_history = self.conv.history[:]
        self.conv.reset()
        answer = self.conv.respond(question, temperature=0.0, seed=0)
        self.conv.history = saved_history
        return answer

    # ----------------------------------------------------------- respond with chaining
    def respond(self, question: str, max_new_tokens: int = 200,
                temperature: float | None = None,
                seed: int | None = None) -> tuple[str, ChainResult | None]:
        """Respond, falling back to fact chaining if direct retrieval fails.

        Returns ``(answer, chain_result)``. If chaining was used,
        ``chain_result`` is non-None and contains the reasoning trace.
        """
        # try direct response first
        answer = self.conv.respond(question, max_new_tokens=max_new_tokens,
                                   temperature=temperature, seed=seed)
        if answer and answer != FALLBACK_RESPONSE and answer.strip():
            return answer, None

        # direct retrieval failed — try chaining
        self.conv.reset()  # clear the failed attempt from history
        chain_result = self.try_chain(question)
        if chain_result.success:
            # record the successful chain in history
            self.conv.history.append(_make_turn("user", question))
            self.conv.history.append(_make_turn("palimpseste", chain_result.answer))
            return chain_result.answer, chain_result

        # chaining also failed — return fallback
        self.conv.history.append(_make_turn("user", question))
        self.conv.history.append(_make_turn("palimpseste", FALLBACK_RESPONSE))
        return FALLBACK_RESPONSE, chain_result


def _make_turn(role: str, text: str):
    """Create a Turn without importing the dataclass at module level."""
    from .chat import Turn
    return Turn(role=role, text=text)