palimseste/evolution.py

"""PALIMPSESTE — Evolution layer: closing the gap with transformer LLMs.

5 cognitive upgrades that make PALIMPSESTE competitive on reasoning quality:

1. **ResponseSynthesizer**: when multiple facts are retrieved, combine them
   into a novel, coherent answer — not just echo the closest match.

2. **EntityTracker**: tracks entities mentioned in conversation and resolves
   pronouns ("it", "that", "this thing") across turns.

3. **QueryRouter**: classifies intent (factual, comparison, how-to, opinion,
   code) and routes to the best retrieval strategy.

4. **CodePatternBank**: stores and retrieves code snippets indexed by
   language + task.

5. **ConfidenceCalibrator**: estimates uncertainty from retrieval quality
   and hedges appropriately ("I think...", "I'm not sure but...").

All operate within the HV paradigm. No transformer. No gradient.
"""

from __future__ import annotations

import re
import time
from dataclasses import dataclass, field
from typing import Optional
import numpy as np

__all__ = [
    "ResponseSynthesizer",
    "EntityTracker",
    "QueryRouter",
    "CodePatternBank",
    "ConfidenceCalibrator",
    "SynthesisResult",
    "Entity",
    "QueryIntent",
    "CodePattern",
    "CalibrationResult",
]


# ================================================================
# 1. RESPONSE SYNTHESIZER
# ================================================================

@dataclass
class SynthesisResult:
    """Result of multi-fact synthesis."""
    answer: str
    n_facts_used: int
    facts_used: list[str]
    synthesis_type: str  # "single" | "merge" | "compare" | "compose"
    confidence: float


class ResponseSynthesizer:
    """Combine multiple retrieved facts into novel answers.

    When the user asks a question that touches multiple topics (e.g., "tell
    me about Python and how it compares to Java"), a simple retrieval returns
    one fact. This synthesizer:

    1. Decomposes the query into key topics
    2. Retrieves facts for each topic independently
    3. Combines them using template-based synthesis

    This is what GPT-4 does naturally via attention over its context.
    PALIMPSESTE does it via explicit multi-query retrieval + template merge.
    """

    # Synthesis templates for combining facts
    TEMPLATES = {
        'compare': "regarding {topic_a}: {fact_a}. regarding {topic_b}: {fact_b}.",
        'merge': "{fact_a} additionally, {fact_b}.",
        'compose': "{fact_a} this means that {fact_b}.",
        'single': "{fact_a}",
    }

    def __init__(self, conv) -> None:
        """conv: a Conversation object with respond() and _fuzzy_match()."""
        self.conv = conv

    def synthesize(self, question: str, max_new_tokens: int = 100) -> SynthesisResult:
        """Synthesize a response from potentially multiple facts.

        1. Try direct retrieval first.
        2. If the question mentions multiple topics, retrieve each separately.
        3. Combine using the appropriate template.
        """
        # Step 1: try direct retrieval
        self.conv.reset()
        direct = self.conv.respond(question, max_new_tokens=max_new_tokens,
                                    temperature=0.0, seed=0)

        from .chat import FALLBACK_RESPONSE
        if direct and direct != FALLBACK_RESPONSE and direct.strip():
            return SynthesisResult(
                answer=direct,
                n_facts_used=1,
                facts_used=[direct],
                synthesis_type='single',
                confidence=1.0,
            )

        # Step 2: decompose into topics
        topics = self._extract_topics(question)
        if len(topics) <= 1:
            return SynthesisResult(
                answer=direct if direct != FALLBACK_RESPONSE else "",
                n_facts_used=0,
                facts_used=[],
                synthesis_type='single',
                confidence=0.0,
            )

        # Step 3: retrieve facts for each topic
        facts: dict[str, str] = {}
        for topic in topics:
            self.conv.reset()
            topic_q = self._make_question(topic, question)
            fact = self.conv.respond(topic_q, max_new_tokens=60,
                                      temperature=0.0, seed=0)
            if fact and fact != FALLBACK_RESPONSE and fact.strip():
                facts[topic] = fact.strip()

        if not facts:
            return SynthesisResult(
                answer="",
                n_facts_used=0,
                facts_used=[],
                synthesis_type='single',
                confidence=0.0,
            )

        # Step 4: synthesize
        if len(facts) == 1:
            fact = list(facts.values())[0]
            return SynthesisResult(
                answer=fact, n_facts_used=1, facts_used=[fact],
                synthesis_type='single', confidence=0.8,
            )

        # Multiple facts — choose template
        question_lower = question.lower()
        if any(w in question_lower for w in ['compare', 'difference', 'vs', 'versus', 'or']):
            synth_type = 'compare'
        elif any(w in question_lower for w in ['and', 'also', 'both', 'plus']):
            synth_type = 'merge'
        else:
            synth_type = 'compose'

        fact_list = list(facts.values())
        topic_list = list(facts.keys())

        template = self.TEMPLATES[synth_type]
        if synth_type == 'compare' and len(topic_list) >= 2:
            answer = template.format(
                topic_a=topic_list[0], fact_a=fact_list[0],
                topic_b=topic_list[1], fact_b=fact_list[1],
            )
        else:
            # Merge facts sequentially
            parts = []
            for i, (topic, fact) in enumerate(facts.items()):
                if i == 0:
                    parts.append(fact)
                else:
                    parts.append(fact.lower() if fact[0:1].isupper() else fact)
            answer = '. '.join(parts) + '.'

        return SynthesisResult(
            answer=answer,
            n_facts_used=len(facts),
            facts_used=list(facts.values()),
            synthesis_type=synth_type,
            confidence=0.7,
        )

    def _extract_topics(self, question: str) -> list[str]:
        """Extract key topics from a multi-topic question."""
        q = question.lower().strip().rstrip('?')

        # Pattern: "X and Y", "X vs Y", "X or Y", "compare X and Y"
        for sep in [' vs ', ' versus ', ' or ']:
            if sep in q:
                parts = q.split(sep)
                if len(parts) >= 2:
                    topics = []
                    for p in parts:
                        p = p.strip()
                        p = re.sub(r'^(what is|what are|tell me about|explain)\s+', '', p)
                        p = re.sub(r'^(compare|difference between)\s+', '', p)
                        topics.append(p.strip())
                    return [t for t in topics if t]

        if ' and ' in q:
            parts = q.split(' and ')
            if len(parts) >= 2:
                topics = []
                for p in parts:
                    p = re.sub(r'^(what is|what are|tell me about|explain)\s+', '', p.strip())
                    topics.append(p.strip())
                return [t for t in topics if t]

        # Pattern: "X's Y" → topic = X
        m = re.match(r"what\s+is\s+(.+?)'s\s+(.+)", q)
        if m:
            return [m.group(1).strip()]

        return []

    def _make_question(self, topic: str, original: str) -> str:
        """Make a retrieval question for a single topic."""
        topic = topic.strip().rstrip('.?,!')
        # Simple "what is X" format works best for associative retrieval
        return f"what is {topic}"


# ================================================================
# 2. ENTITY TRACKER
# ================================================================

@dataclass
class Entity:
    """A tracked entity in conversation."""
    name: str
    first_mentioned_turn: int
    last_mentioned_turn: int
    mentions: int = 1
    attributes: dict[str, str] = field(default_factory=dict)


class EntityTracker:
    """Track entities and resolve pronouns across conversation turns.

    When the user says "what about its speed?" after talking about Python,
    the tracker knows "its" = Python and rewrites the query.

    GPT-4 does this implicitly via attention. PALIMPSESTE does it
    explicitly via entity tracking — more transparent and debuggable.
    """

    PRONOUNS = {'it', 'its', 'this', 'that', 'these', 'those', 'he', 'she',
                'his', 'her', 'they', 'their', 'them', 'one'}

    def __init__(self) -> None:
        self._entities: dict[str, Entity] = {}
        self._turn: int = 0
        self._last_entity: str | None = None

    def update(self, user_text: str) -> str:
        """Process user input: extract entities and resolve pronouns.

        Returns the (possibly rewritten) user text with pronouns replaced.
        """
        self._turn += 1
        words = user_text.split()

        # Check for pronouns BEFORE extracting new entities
        has_pronoun = any(p in user_text.lower().split() for p in self.PRONOUNS)
        pronoun_to_resolve = self._last_entity if (has_pronoun and self._last_entity) else None

        # Extract entities: meaningful words that aren't pronouns or stop words
        stop_words = self.PRONOUNS | {'the', 'a', 'an', 'is', 'are',
                'was', 'were', 'what', 'who', 'how', 'why', 'when', 'where',
                'do', 'does', 'did', 'can', 'could', 'would', 'should',
                'will', 'about', 'tell', 'me', 'and', 'or', 'but', 'to',
                'of', 'in', 'on', 'at', 'for', 'with', 'from', 'by', 'its'}
        new_entities_this_turn: list[str] = []
        for word in words:
            clean = word.strip('.,!?;:"\'()[]').lower()
            if len(clean) < 3 or clean in stop_words:
                continue
            if clean in self._entities:
                ent = self._entities[clean]
                ent.last_mentioned_turn = self._turn
                ent.mentions += 1
                new_entities_this_turn.append(clean)
            else:
                self._entities[clean] = Entity(
                    name=clean,
                    first_mentioned_turn=self._turn,
                    last_mentioned_turn=self._turn,
                )
                new_entities_this_turn.append(clean)

        # Update last entity only if we found new ones this turn
        if new_entities_this_turn:
            self._last_entity = new_entities_this_turn[-1]

        # Resolve pronouns using the PREVIOUS entity (before this turn's extraction)
        resolved = user_text
        if pronoun_to_resolve:
            for pronoun in ['its', 'it', 'this', 'that', 'these', 'those']:
                resolved = re.sub(
                    rf'\b{pronoun}\b',
                    pronoun_to_resolve,
                    resolved,
                    count=1,
                    flags=re.IGNORECASE,
                )

        return resolved

    def get_entity(self, name: str) -> Entity | None:
        return self._entities.get(name.lower())

    def get_last_entity(self) -> str | None:
        return self._last_entity

    @property
    def n_entities(self) -> int:
        return len(self._entities)

    def all_entities(self) -> list[Entity]:
        return sorted(self._entities.values(),
                       key=lambda e: e.last_mentioned_turn, reverse=True)

    def reset(self) -> None:
        self._entities.clear()
        self._turn = 0
        self._last_entity = None


# ================================================================
# 3. QUERY ROUTER
# ================================================================

@dataclass
class QueryIntent:
    """Classified intent of a user query."""
    intent: str  # "factual" | "comparison" | "howto" | "opinion" | "code" | "definition" | "greeting"
    topics: list[str]
    confidence: float
    strategy: str  # recommended retrieval strategy


class QueryRouter:
    """Classify intent and route to the best retrieval strategy.

    Different question types need different approaches:
    - Factual: direct retrieval
    - Comparison: multi-topic synthesis
    - How-to: step-by-step retrieval
    - Code: code pattern bank
    - Opinion: template-based with hedging
    - Definition: direct retrieval with template

    GPT-4 handles all these via the same attention mechanism.
    PALIMPSESTE routes explicitly — more efficient and transparent.
    """

    INTENT_PATTERNS = [
        ('greeting', re.compile(r'\b(hello|hi|hey|good morning|good evening|howdy|sup)\b', re.I), 'direct'),
        ('comparison', re.compile(r'\b(compare|difference|vs|versus|better than|worse than|or)\b', re.I), 'synthesize'),
        ('howto', re.compile(r'\b(how (do|to|can)|steps? to|way to|guide)\b', re.I), 'stepwise'),
        ('code', re.compile(r'\b(code|function|program|script|write|implement|algorithm|debug)\b', re.I), 'code'),
        ('opinion', re.compile(r'\b(think|opinion|believe|feel|should i)\b', re.I), 'hedge'),
        ('definition', re.compile(r'\b(what is|what are|define|definition|explain|meaning)\b', re.I), 'direct'),
        ('factual', re.compile(r'\b(who|where|when|why|which|whose)\b', re.I), 'direct'),
    ]

    def __init__(self) -> None:
        pass

    def classify(self, question: str) -> QueryIntent:
        """Classify the intent of a question."""
        q_lower = question.lower().strip()

        # Try each pattern in order
        for intent, pattern, strategy in self.INTENT_PATTERNS:
            if pattern.search(q_lower):
                topics = self._extract_topics(question)
                return QueryIntent(
                    intent=intent,
                    topics=topics,
                    confidence=0.85,
                    strategy=strategy,
                )

        # Default: factual
        return QueryIntent(
            intent='factual',
            topics=[],
            confidence=0.5,
            strategy='direct',
        )

    def _extract_topics(self, question: str) -> list[str]:
        """Extract key topic words from the question."""
        q = question.lower()
        # Remove question words and stop words
        stop = {'what', 'is', 'are', 'was', 'were', 'the', 'a', 'an',
                'who', 'where', 'when', 'why', 'how', 'do', 'does', 'did',
                'can', 'could', 'would', 'should', 'will', 'tell', 'me',
                'about', 'of', 'in', 'on', 'at', 'for', 'with', 'and', 'or',
                'to', 'from', 'by', 'this', 'that', 'it', 'its'}
        words = re.findall(r'[a-z]+', q)
        topics = [w for w in words if w not in stop and len(w) > 2]
        return topics[:5]


# ================================================================
# 4. CODE PATTERN BANK
# ================================================================

@dataclass
class CodePattern:
    """A stored code pattern."""
    id: int
    language: str
    task: str
    code: str
    description: str
    keywords: list[str]


class CodePatternBank:
    """Store and retrieve code snippets.

    A dedicated memory for code patterns, indexed by language and task.
    When the user asks "how to read a file in Python", the bank retrieves
    the most relevant code pattern.

    GPT-4 has code in its training data mixed with everything else.
    PALIMPSESTE has a dedicated code bank — faster and more precise.
    """

    def __init__(self) -> None:
        self._patterns: list[CodePattern] = []
        self._next_id: int = 0

    def store(self, language: str, task: str, code: str,
              description: str = "", keywords: list[str] | None = None) -> CodePattern:
        """Store a code pattern."""
        pattern = CodePattern(
            id=self._next_id,
            language=language.lower(),
            task=task.lower(),
            code=code,
            description=description,
            keywords=keywords or [],
        )
        self._next_id += 1
        self._patterns.append(pattern)
        return pattern

    def retrieve(self, query: str, language: str | None = None,
                 top_k: int = 3) -> list[tuple[CodePattern, float]]:
        """Retrieve code patterns matching a query.

        Uses keyword overlap scoring (simple but effective for code).
        """
        query_lower = query.lower()
        query_words = set(re.findall(r'[a-z_]+', query_lower))

        scored: list[tuple[CodePattern, float]] = []
        for pattern in self._patterns:
            if language and pattern.language != language.lower():
                continue

            # Score by keyword overlap
            pattern_words = set(pattern.keywords) | {pattern.task} | set(pattern.language.split())
            code_words = set(re.findall(r'[a-z_]+', pattern.code.lower()))
            pattern_words |= code_words

            overlap = len(query_words & pattern_words)
            total = len(query_words)
            score = overlap / max(total, 1)

            # Boost for task match
            if pattern.task in query_lower:
                score += 0.3

            if score > 0:
                scored.append((pattern, min(score, 1.0)))

        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

    def load_defaults(self) -> None:
        """Load a set of common code patterns."""
        defaults = [
            ('python', 'read file', 'with open("file.txt") as f:\n    content = f.read()',
             'Read a text file', ['open', 'read', 'file']),
            ('python', 'write file', 'with open("file.txt", "w") as f:\n    f.write("hello")',
             'Write to a text file', ['write', 'file']),
            ('python', 'function', 'def square(x):\n    return x * x',
             'Define a function', ['def', 'function', 'return']),
            ('python', 'class', 'class Dog:\n    def __init__(self, name):\n        self.name = name',
             'Define a class', ['class', 'init', 'self']),
            ('python', 'loop', 'for i in range(10):\n    print(i)',
             'For loop', ['for', 'range', 'loop']),
            ('python', 'list comprehension', 'squares = [x*x for x in range(10)]',
             'List comprehension', ['list', 'comprehension']),
            ('python', 'dictionary', 'd = {"key": "value"}\nprint(d["key"])',
             'Dictionary operations', ['dict', 'dictionary', 'key']),
            ('python', 'try except', 'try:\n    result = 1 / 0\nexcept ZeroDivisionError:\n    print("error")',
             'Error handling', ['try', 'except', 'error']),
            ('python', 'sort list', 'sorted_list = sorted(my_list, key=lambda x: x.name)',
             'Sort a list', ['sort', 'sorted']),
            ('python', 'json parse', 'import json\ndata = json.loads(\'{"key": "value"}\')',
             'Parse JSON', ['json', 'parse', 'loads']),
        ]
        for lang, task, code, desc, kw in defaults:
            self.store(lang, task, code, desc, kw)

    @property
    def n_patterns(self) -> int:
        return len(self._patterns)


# ================================================================
# 5. CONFIDENCE CALIBRATOR
# ================================================================

@dataclass
class CalibrationResult:
    """Result of confidence calibration."""
    text: str
    raw_confidence: float
    calibrated_confidence: float
    hedged: bool
    hedge_phrase: str | None


class ConfidenceCalibrator:
    """Calibrate confidence and hedge appropriately.

    The raw confidence from Hamming similarity is often overconfident
    (sim=1.0 for exact matches, even if the answer is wrong). This
    calibrator adjusts:

    1. Checks if the retrieved answer is actually relevant to the question
    2. Applies hedging for uncertain answers
    3. Adds "I think" / "I'm not sure" when appropriate

    GPT-4 calibrates via RLHF. PALIMPSESTE calibrates via retrieval quality.
    """

    HEDGE_PHRASES = {
        'low': "i'm not sure, but ",
        'medium': "i think ",
        'high': "",  # no hedge for high confidence
    }

    def __init__(self, low_threshold: float = 0.3, high_threshold: float = 0.7) -> None:
        self.low_threshold = low_threshold
        self.high_threshold = high_threshold

    def calibrate(self, text: str, raw_confidence: float,
                  question: str = "", answer: str = "") -> CalibrationResult:
        """Calibrate confidence and apply hedging.

        Parameters
        ----------
        text : str
            The response text to potentially hedge.
        raw_confidence : float
            Raw confidence from retrieval (0-1).
        question : str
            The original question (for relevance check).
        answer : str
            The retrieved answer (for length check).

        Returns
        -------
        CalibrationResult
        """
        calibrated = raw_confidence

        # Factor 1: answer length — very short answers are slightly less reliable
        if answer and len(answer.strip()) < 5:
            calibrated *= 0.8

        # Factor 2: question-answer relevance (word overlap)
        if question and answer:
            q_words = set(re.findall(r'[a-z]+', question.lower()))
            a_words = set(re.findall(r'[a-z]+', answer.lower()))
            stop = {'what', 'is', 'the', 'a', 'an', 'of', 'to', 'in', 'for',
                    'and', 'or', 'how', 'who', 'why', 'when', 'where'}
            q_meaningful = q_words - stop
            a_meaningful = a_words - stop
            if q_meaningful:
                overlap = len(q_meaningful & a_meaningful) / len(q_meaningful)
                # Blend: 70% original confidence + 30% relevance
                calibrated = 0.7 * raw_confidence + 0.3 * overlap

        # Determine hedge level
        if calibrated < self.low_threshold:
            hedge_level = 'low'
        elif calibrated < self.high_threshold:
            hedge_level = 'medium'
        else:
            hedge_level = 'high'

        hedge = self.HEDGE_PHRASES[hedge_level]
        hedged_text = text
        hedged = False

        if hedge and text.strip():
            # Don't double-hedge
            if not text.lower().startswith(("i'm not sure", "i think", "i believe")):
                hedged_text = hedge + text
                hedged = True

        return CalibrationResult(
            text=hedged_text,
            raw_confidence=raw_confidence,
            calibrated_confidence=calibrated,
            hedged=hedged,
            hedge_phrase=hedge if hedge else None,
        )