inference/self_correction.py

#!/usr/bin/env python3
"""Codette Self-Correction Engine — Autonomous constraint compliance

Three systems working together:

1. SELF-CORRECTION LOOP (check_violations → re-prompt if needed)
   Detects constraint violations BEFORE sending the response.
   If the model's output violates constraints, it re-generates with
   explicit correction instructions. Max 1 retry to avoid latency.

2. PERSISTENT BEHAVIOR MEMORY (learn from mistakes)
   Stores constraint violation/success patterns as cocoon memories.
   On startup, loads past lessons and injects them into the system prompt
   so the model learns from its own history across sessions.

3. CHAOS DETECTOR (graceful degradation under competing pressures)
   When multiple constraints + uncertainty + mode personality collide,
   detects the "chaos level" and applies a simplify-first strategy
   instead of letting the model collapse into incoherence.
"""

import json
import os
import re
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple


# ================================================================
# 1. SELF-CORRECTION: Violation detection
# ================================================================

def detect_violations(response: str, constraints: dict) -> List[str]:
    """Check a response against constraints and return list of violations.

    Returns empty list if response is compliant.
    """
    violations = []

    if not response or not constraints:
        return violations

    words = response.split()
    word_count = len(words)

    # Check word limit
    max_words = constraints.get('max_words')
    if max_words and word_count > max_words:
        violations.append(f"WORD_LIMIT: {word_count} words exceeds limit of {max_words}")

    # Check sentence limit
    max_sentences = constraints.get('max_sentences')
    if max_sentences:
        sentences = re.split(r'(?<=[.!?])\s+', response.strip())
        # Filter out empty strings
        sentences = [s for s in sentences if s.strip()]
        if len(sentences) > max_sentences:
            violations.append(f"SENTENCE_LIMIT: {len(sentences)} sentences exceeds limit of {max_sentences}")

    # Check incomplete sentence (ends with dangling word or no punctuation)
    if response and response[-1] not in '.!?':
        violations.append("INCOMPLETE: Response does not end with proper punctuation")

    # Check dangling words at end
    _DANGLING = {
        'that', 'which', 'who', 'with', 'and', 'but', 'or', 'the', 'a', 'an',
        'in', 'on', 'of', 'for', 'to', 'by', 'from', 'as', 'if', 'because',
        'including', 'is', 'are', 'was', 'were', 'has', 'have', 'had',
    }
    if words:
        last_word = words[-1].lower().rstrip('.,;:!?')
        if last_word in _DANGLING and response[-1] == '.':
            # Ends with a dangling word + forced period = likely truncated
            violations.append(f"DANGLING: Response ends with '{words[-1]}' (likely truncated)")

    # Check binary constraint
    if constraints.get('binary'):
        first_word = words[0].lower().rstrip('.,;:!?') if words else ''
        if first_word not in ('yes', 'no', 'true', 'false'):
            violations.append("BINARY: Response should start with Yes/No")

    # Check brevity — soft check, flag if over 30 words
    if constraints.get('brevity') and word_count > 30:
        violations.append(f"BREVITY: {word_count} words is not brief (target: under 30)")

    return violations


def _extract_sum_diff_problem(query: str) -> Optional[Tuple[float, float, float]]:
    """Extract (total, diff, correct_answer) from sum-difference word problem queries.

    Pattern: 'A and B cost $S total. A costs $D more than B.'
    Correct answer: B = (S - D) / 2
    Returns None if no matching pattern is found.
    """
    q = query.lower()

    # Require BOTH a sum signal AND a difference signal to avoid false positives
    # on simple addition ("3 apples and 5 oranges, how many total?").
    has_sum = bool(re.search(
        r'\b(together|total|combined|in total|in all|altogether)\b', q
    ))
    has_diff = bool(re.search(
        r'\bmore\s+than\b|\bdifference\b|\bextra\b|\bcosts?\s+\$?[\d].*more\b', q
    ))
    if not (has_sum and has_diff):
        return None

    # Must also involve a cost/price context
    if not re.search(r'\$[\d]|cost|price|pay|worth|dollar|cent', q):
        return None

    # Extract "D more than" (the difference)
    more_match = re.search(r'\$?([\d]+(?:\.\d+)?)\s+(?:dollar[s]?\s+)?more\b', q)
    if not more_match:
        return None
    diff = float(more_match.group(1))

    # All monetary amounts in query
    all_amounts = [float(m) for m in re.findall(r'\$?([\d]+(?:\.\d+)?)\b', q)]

    # Total must be strictly greater than diff
    candidates = [a for a in all_amounts if a > diff + 0.001]
    if not candidates:
        return None
    total = max(candidates)

    correct = round((total - diff) / 2, 4)
    if correct <= 0:
        return None

    return (total, diff, correct)


def _check_sum_diff_bias(query: str, response: str) -> Tuple[str, List[str]]:
    """Detect and correct the sum-difference cognitive bias in word problem responses.

    Intuitive (wrong) answer: S - D
    Correct answer:           (S - D) / 2
    """
    issues: List[str] = []

    result = _extract_sum_diff_problem(query)
    if result is None:
        return response, issues

    total, diff, correct = result
    biased = round(total - diff, 4)

    correct_cents = int(round(correct * 100))
    biased_cents = int(round(biased * 100))

    resp_lower = response.lower()

    # Representations of the correct answer — if already present, nothing to do
    correct_strs = [
        f"${correct:.2f}", f"{correct_cents} cent",
        f"{correct_cents}¢", str(correct_cents),
    ]
    if any(s.lower() in resp_lower for s in correct_strs):
        return response, issues

    # Representations of the biased answer — in priority order for replacement
    biased_strs = [
        f"${biased:.2f}", f"{biased_cents} cents",
        f"{biased_cents} cent", f"{biased_cents}¢",
    ]
    has_bias = any(s.lower() in resp_lower for s in biased_strs)
    if not has_bias:
        return response, issues

    # Replace ALL occurrences — longest match first to avoid partial replacement.
    # All mentions of the biased answer in the response are wrong; leave none standing.
    for bs in sorted(biased_strs, key=len, reverse=True):
        pat = re.compile(re.escape(bs), re.IGNORECASE)
        if pat.search(response):
            replacement = f"{correct_cents} cents" if "cent" in bs.lower() else f"${correct:.2f}"
            corrected, n = pat.subn(replacement, response)  # count=0 → replace all
            if n:
                issues.append(
                    f"ARITHMETIC_FIX: sum-diff bias corrected "
                    f"({n}× '{bs}' → '{replacement}'; "
                    f"algebra: ({total}-{diff})/2={correct})"
                )
                return corrected, issues

    return response, issues


def universal_self_check(response: str, query: str = "") -> Tuple[str, List[str]]:
    """PERMANENT LOCK enforcement — runs on EVERY response, not just constrained ones.

    Enforces Lock 1 (Answer→Stop), Lock 3 (self-check), Lock 4 (no incomplete outputs).
    When query is provided, also applies arithmetic bias detection (Lock 5).
    Returns (cleaned_response, issues_found).
    """
    issues = []
    if not response:
        return response, issues

    cleaned = response.rstrip()

    # LOCK 3: Self-check — detect echo-back failures (model repeating the question)
    echo_patterns = [
        r'^You (?:received|asked|said|mentioned|wrote)(?:\s+this)?\s*(?:question|query)?[:\s]',
        r'^The (?:question|query) (?:is|was)[:\s]',
        r'^(?:Your|The) question[:\s]',
        r'^(?:Here is|I received|Let me|I will)\s+(?:the|your|this)\s+(?:question|query)',
    ]
    is_echo = False
    for ep in echo_patterns:
        if re.match(ep, cleaned, re.I):
            is_echo = True
            # Strip ONLY the first echo line, keep everything after
            lines = cleaned.split('\n')
            non_echo = [l for l in lines[1:] if l.strip()]  # Skip first line only
            if non_echo:
                cleaned = '\n'.join(non_echo).strip()
                issues.append("LOCK3_FIX: Removed question echo-back")
            else:
                # Nothing left after echo removal — keep original rather than empty
                # An imperfect response is better than no response
                issues.append("LOCK3_WARN: Echo detected but kept response (no content after removal)")
            break

    # LOCK 4: No incomplete outputs — fix missing punctuation
    if cleaned and cleaned[-1] not in '.!?"\')\u2019':
        words = cleaned.split()
        _DANGLING = {
            'that', 'which', 'who', 'with', 'and', 'but', 'or', 'the', 'a', 'an',
            'in', 'on', 'of', 'for', 'to', 'by', 'from', 'as', 'if', 'because',
            'including', 'is', 'are', 'was', 'were', 'has', 'have', 'had',
            'not', 'very', 'also', 'just', 'even', 'still', 'such', 'like',
        }
        # Strip dangling words from end
        while len(words) > 1 and words[-1].lower().rstrip('.,;:!?') in _DANGLING:
            words.pop()
            issues.append("LOCK4_FIX: Removed dangling word from end")
        cleaned = ' '.join(words)
        cleaned = cleaned.rstrip(' ,;:—-')
        if cleaned and cleaned[-1] not in '.!?"\')\u2019':
            cleaned += '.'
            issues.append("LOCK4_FIX: Added missing end punctuation")

    # LOCK 1: Answer→Stop — detect and trim post-answer drift
    sentences = re.split(r'(?<=[.!?])\s+', cleaned.strip())
    sentences = [s for s in sentences if s.strip()]

    if len(sentences) > 2:
        # Check for drift indicators in later sentences (from 2nd sentence onward)
        drift_phrases = [
            r"^(?:furthermore|moreover|additionally|in addition|it(?:'|\u2019)s worth noting)",
            r'^(?:this (?:is|means|suggests|implies|shows|demonstrates|highlights|reveals))',
            r'^(?:one (?:could|might|may) (?:argue|say|suggest))',
            r'^(?:from (?:a|an|the) \w+ perspective)',
            r'^(?:in (?:other words|essence|summary|fact))',
            r'^(?:by (?:acknowledging|understanding|exploring|examining|considering))',
            r'^(?:understanding \w+ requires)',
            r'^(?:it is (?:important|worth|essential|notable|key) to (?:note|understand|recognize|remember))',
        ]
        drift_pattern = re.compile('|'.join(drift_phrases), re.I)

        cut_at = None
        for i in range(2, len(sentences)):  # Check from 3rd sentence onward
            if drift_pattern.match(sentences[i]):
                cut_at = i
                break

        if cut_at is not None:
            cleaned = ' '.join(sentences[:cut_at])
            if cleaned and cleaned[-1] not in '.!?':
                cleaned += '.'
            issues.append(f"LOCK1_TRIM: Removed {len(sentences) - cut_at} drift sentence(s)")

    # LOCK 1 softcap: Safety net against runaway output.
    # Codette may answer at whatever length the response needs; this only trims
    # genuinely excessive responses to the last complete sentence within the cap.
    # Set high so it rarely triggers — it's a guardrail, not a style limit.
    SOFTCAP_WORDS = 1000
    words = cleaned.split()
    if len(words) > SOFTCAP_WORDS:
        sentences = re.split(r'(?<=[.!?])\s+', cleaned.strip())
        sentences = [s for s in sentences if s.strip()]
        fitted = []
        wc = 0
        for s in sentences:
            sw = len(s.split())
            if wc + sw <= SOFTCAP_WORDS:
                fitted.append(s)
                wc += sw
            else:
                # Allow one more sentence if we're close and it's reasonably short
                if wc >= SOFTCAP_WORDS // 2 and sw <= 40:
                    fitted.append(s)
                break
        if fitted and len(fitted) < len(sentences):
            cleaned = ' '.join(fitted)
            if cleaned and cleaned[-1] not in '.!?':
                cleaned += '.'
            issues.append(f"LOCK1_SOFTCAP: Trimmed from {len(words)} to {len(cleaned.split())} words")

    # LOCK 4 final check: Ensure the final output is complete
    if cleaned and cleaned[-1] not in '.!?"\')\u2019':
        cleaned += '.'
        issues.append("LOCK4_FIX: Final punctuation check")

    # LOCK 5: Arithmetic bias check \u2014 sum-difference word problems
    # Only runs when the query is provided (caller opt-in).
    if query:
        cleaned, arith_issues = _check_sum_diff_bias(query, cleaned)
        issues.extend(arith_issues)

    return cleaned, issues


def build_correction_prompt(original_response: str, violations: List[str],
                            constraints: dict, query: str) -> str:
    """Build a correction prompt that tells the model exactly what it did wrong.

    This is sent as a follow-up user message to get a corrected response.
    """
    violation_text = "\n".join(f"  - {v}" for v in violations)

    constraint_text = []
    if 'max_words' in constraints:
        constraint_text.append(f"maximum {constraints['max_words']} words")
    if 'max_sentences' in constraints:
        constraint_text.append(f"maximum {constraints['max_sentences']} sentence(s)")
    if constraints.get('brevity'):
        constraint_text.append("be brief")
    if constraints.get('binary'):
        constraint_text.append("answer Yes or No")

    return (
        f"Your previous answer violated these constraints:\n{violation_text}\n\n"
        f"The required constraints are: {', '.join(constraint_text)}.\n\n"
        f"Rewrite your answer to the original question: \"{query}\"\n"
        f"This time, STRICTLY obey the constraints. "
        f"If you can't fit everything, SIMPLIFY — say the right thing cleanly. "
        f"Do NOT end mid-sentence. Give a COMPLETE thought within the limits."
    )


# ================================================================
# 2. PERSISTENT BEHAVIOR MEMORY
# ================================================================

_BEHAVIOR_MEMORY_FILE = Path(__file__).resolve().parent.parent / "cocoons" / "behavior_memory.json"


class BehaviorMemory:
    """Persistent memory of constraint successes and failures.

    Stores patterns like:
    - "When asked for 'one sentence under 10 words' with philosophy mode,
       I over-explained. The correct approach: simplify to core answer."
    - "When asked 'yes or no', I answered directly. Keep doing this."

    These are loaded on startup and injected into the system prompt
    as learned behaviors.
    """

    def __init__(self, storage_path: str = None):
        self.storage_path = Path(storage_path or _BEHAVIOR_MEMORY_FILE)
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        self.lessons: List[Dict] = []
        self._max_lessons = 50  # Rolling window
        self._load()

    def _load(self):
        """Load lessons from disk."""
        if self.storage_path.exists():
            try:
                with open(self.storage_path, "r") as f:
                    data = json.load(f)
                self.lessons = data.get("lessons", [])[-self._max_lessons:]
            except Exception:
                self.lessons = []

    def _save(self):
        """Save lessons to disk."""
        try:
            with open(self.storage_path, "w") as f:
                json.dump({
                    "lessons": self.lessons[-self._max_lessons:],
                    "updated_at": time.time(),
                }, f, indent=2)
        except Exception as e:
            print(f"  [BEHAVIOR] Save failed: {e}")

    def record_violation(self, query: str, constraints: dict, violations: List[str],
                         adapter: str, response_preview: str):
        """Record a constraint violation for future learning."""
        lesson = {
            "type": "violation",
            "timestamp": time.time(),
            "adapter": adapter,
            "constraints": constraints,
            "violations": violations,
            "query_preview": query[:80],
            "response_preview": response_preview[:100],
            "lesson": self._derive_lesson(constraints, violations, adapter),
        }
        self.lessons.append(lesson)
        self._save()

    def record_success(self, query: str, constraints: dict, adapter: str,
                       word_count: int):
        """Record a successful constraint compliance."""
        lesson = {
            "type": "success",
            "timestamp": time.time(),
            "adapter": adapter,
            "constraints": constraints,
            "word_count": word_count,
            "query_preview": query[:80],
        }
        self.lessons.append(lesson)
        self._save()

    def _derive_lesson(self, constraints: dict, violations: List[str],
                       adapter: str) -> str:
        """Derive a human-readable lesson from a violation."""
        parts = []

        for v in violations:
            if "WORD_LIMIT" in v:
                parts.append(f"With {adapter} mode, I exceeded the word limit. I must count words and simplify.")
            elif "SENTENCE_LIMIT" in v:
                parts.append(f"With {adapter} mode, I used too many sentences. One complete thought is enough.")
            elif "INCOMPLETE" in v:
                parts.append("I ended a sentence incomplete. I must finish thoughts cleanly or simplify.")
            elif "DANGLING" in v:
                parts.append("I ended with a connecting word (cut-off). I must close the thought before the limit.")
            elif "BINARY" in v:
                parts.append("I didn't start with Yes/No when asked for a binary answer.")
            elif "BREVITY" in v:
                parts.append(f"With {adapter} mode, I over-explained when asked to be brief.")

        return " ".join(parts) if parts else "Constraint violation occurred."

    def get_lessons_for_prompt(self, max_lessons: int = 5) -> str:
        """Build a system prompt section from recent lessons.

        Returns a string to inject into the system prompt that teaches
        the model from its own past mistakes.
        """
        # Get recent violations (most valuable for learning)
        violations = [l for l in self.lessons if l["type"] == "violation"]
        if not violations:
            return ""

        # Take the most recent unique lessons
        seen_lessons = set()
        unique_lessons = []
        for v in reversed(violations):
            lesson = v.get("lesson", "")
            if lesson and lesson not in seen_lessons:
                seen_lessons.add(lesson)
                unique_lessons.append(lesson)
                if len(unique_lessons) >= max_lessons:
                    break

        if not unique_lessons:
            return ""

        lines = "\n".join(f"  - {l}" for l in unique_lessons)
        return (
            f"\n\nLEARNED FROM PAST MISTAKES (apply these lessons):\n{lines}\n"
            "Remember: constraints always override modes. Simplify rather than truncate."
        )

    def get_stats(self) -> Dict:
        """Get behavior memory statistics."""
        violations = [l for l in self.lessons if l["type"] == "violation"]
        successes = [l for l in self.lessons if l["type"] == "success"]
        return {
            "total_lessons": len(self.lessons),
            "violations": len(violations),
            "successes": len(successes),
            "compliance_rate": len(successes) / max(len(self.lessons), 1),
        }


# ================================================================
# 3. CHAOS DETECTION — graceful degradation under pressure
# ================================================================

def detect_chaos_level(query: str, constraints: dict, adapter: str) -> Tuple[int, List[str]]:
    """Detect how many competing pressures are active.

    Returns:
        (chaos_level: 0-5, pressure_list: list of active pressures)

    Chaos levels:
        0: No constraints, simple query → normal generation
        1: One constraint → straightforward compliance
        2: Two constraints → manageable, minor tension
        3: Three+ constraints → high tension, simplify-first strategy
        4: Constraints + uncertainty requirement + complex topic → danger zone
        5: All of above + mode conflict (philosophy on a binary question) → maximum simplification
    """
    pressures = []

    # Count hard constraints
    if constraints.get('max_words'):
        pressures.append(f"word_limit={constraints['max_words']}")
    if constraints.get('max_sentences'):
        pressures.append(f"sentence_limit={constraints['max_sentences']}")
    if constraints.get('brevity'):
        pressures.append("brevity")
    if constraints.get('binary'):
        pressures.append("binary_answer")
    if constraints.get('list_format'):
        pressures.append("list_format")

    # Check for uncertainty requirement
    uncertainty_words = ['uncertain', 'uncertainty', 'maybe', 'might', 'possibly',
                         'include uncertainty', 'not sure', 'debatable']
    query_lower = query.lower()
    if any(w in query_lower for w in uncertainty_words):
        pressures.append("uncertainty_required")

    # Check for complex topic (would normally trigger verbose mode)
    complex_topics = ['consciousness', 'free will', 'meaning of life', 'quantum',
                      'philosophy', 'ethics', 'morality', 'existence', 'reality',
                      'epistemology', 'ontology', 'metaphysics']
    if any(t in query_lower for t in complex_topics):
        pressures.append("complex_topic")

    # Check for mode conflict (verbose mode + tight constraint)
    verbose_modes = {'philosophy', 'consciousness', 'multi_perspective', 'quantum'}
    tight_constraints = bool(constraints.get('max_words') or constraints.get('max_sentences')
                             or constraints.get('binary') or constraints.get('brevity'))
    if adapter in verbose_modes and tight_constraints:
        pressures.append(f"mode_conflict:{adapter}")

    chaos_level = min(len(pressures), 5)
    return chaos_level, pressures


def build_chaos_mitigation(chaos_level: int, pressures: List[str]) -> str:
    """Build additional system prompt guidance for high-chaos situations.

    At chaos level 3+, we inject explicit simplification strategies.
    """
    if chaos_level < 3:
        return ""

    if chaos_level == 3:
        return (
            "\n\nMULTIPLE COMPETING PRESSURES DETECTED. Strategy: "
            "SIMPLIFY FIRST. Give the simplest correct answer that satisfies "
            "ALL constraints. Do not try to be comprehensive — be precise. "
            "One clean thought is better than a cramped paragraph.\n"
        )
    elif chaos_level == 4:
        return (
            "\n\nHIGH CONSTRAINT PRESSURE DETECTED. Strategy: "
            "CORE ANSWER ONLY. Strip everything to the essential truth. "
            "No qualifiers, no hedging, no mode-specific elaboration. "
            "If the topic is complex but the constraint is tight, "
            "give the most important single insight cleanly.\n"
        )
    else:  # chaos_level >= 5
        return (
            "\n\nMAXIMUM CONSTRAINT PRESSURE. Strategy: "
            "OVERRIDE ALL MODES. Your personality mode is SUSPENDED for this response. "
            "You are a plain, direct answerer. Give the shortest correct answer "
            "that satisfies every constraint. No philosophy, no creativity, no "
            "exploration — just the answer. If asked yes/no, say yes or no. "
            "If given a word limit, count your words. Nothing else matters.\n"
        )