Spaces:

afaceuser09
/

InfinityGRM

Runtime error

InfinityGRM / reasoning_engine.py
Update reasoning_engine.py
abc217c verified 22 days ago
52 kB
"""
Infinity GRM - Reasoning Engine
================================
Standalone cognitive/reasoning layer for Infinity GRM.
Features
--------
- Qwen2.5-7B-Instruct reasoning model
- 700 structured reasoning demonstrations by default
- Semantic retrieval of relevant demonstrations
- Planning -> analysis -> verification -> final synthesis
- Persistent SQLite experience memory
- Optional detailed, user-visible reasoning TRACE (not raw hidden CoT)
- Backward-compatible plan()/analyze()/final_answer() API
Environment
-----------
GRM_MODEL=Qwen/Qwen2.5-7B-Instruct
GRM_EXAMPLE_TARGET=700       # clamped to 500..1000
GRM_MEMORY_DB=./grm_experience.sqlite3
GRM_SHOW_REASONING_TRACE=1   # default enabled for detailed trace output
Important
---------
The demonstration library is structured training-style material used at
inference time. It does not change the neural weights of the base model.
For model-weight improvement, use a separate SFT/QLoRA training pipeline.
"""

import os
import re
import json
import time
import sqlite3
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig

logger = logging.getLogger("infinitygrm.reasoning")

MODEL_NAME = os.getenv("GRM_MODEL", "Qwen/Qwen2.5-7B-Instruct")
MEMORY_DB = Path(os.getenv("GRM_MEMORY_DB", "./grm_experience.sqlite3"))
EXAMPLE_TARGET = max(500, min(1000, int(os.getenv("GRM_EXAMPLE_TARGET", "700"))))
SHOW_REASONING_TRACE = os.getenv("GRM_SHOW_REASONING_TRACE", "1").lower() not in {"0", "false", "no"}

STATE: Dict[str, Any] = {
    "model": None,
    "tokenizer": None,
    "embed_model": None,
    "device": None,
    "reasoning_example_embeddings": None,
}


# =====================================================================
# STRUCTURED REASONING DEMONSTRATION LIBRARY
# =====================================================================

_REASONING_SKILLS = [
    ("mathematics", "algebra", "isolate variables, preserve equality, check by substitution"),
    ("mathematics", "geometry", "identify the governing geometric relationship, substitute known quantities, verify units"),
    ("mathematics", "probability", "define the sample space, identify independence/dependence, compute and sanity-check probabilities"),
    ("mathematics", "statistics", "identify the estimator or test, state assumptions, calculate, interpret uncertainty"),
    ("mathematics", "calculus", "identify the derivative/integral rule, apply it carefully, verify with boundary or dimensional checks"),
    ("algorithms", "complexity", "identify the dominant operation, count repeated work, derive asymptotic time and space complexity"),
    ("algorithms", "sorting", "choose an algorithm based on constraints, preserve correctness invariants, analyze complexity"),
    ("algorithms", "graphs", "model entities as vertices and relationships as edges, choose traversal/shortest-path method, verify reachability"),
    ("algorithms", "dynamic programming", "define state, transition, base case, evaluation order, and prove that states cover all cases"),
    ("algorithms", "greedy methods", "define a locally optimal choice, establish the invariant/exchange argument, then verify feasibility"),
    ("coding", "Python debugging", "reproduce the failure, inspect the smallest relevant state, identify the root cause, patch and retest"),
    ("coding", "data structures", "identify access/update requirements, select an appropriate structure, define invariants and complexity"),
    ("coding", "API design", "define inputs/outputs/errors, preserve contracts, validate data, and design deterministic behavior"),
    ("coding", "concurrency", "identify shared state, define synchronization boundaries, avoid races/deadlocks, and test interleavings"),
    ("coding", "optimization", "measure the bottleneck, improve the dominant cost, preserve semantics, and benchmark before/after"),
    ("coding", "machine learning", "separate data/model/objective/evaluation, identify leakage risks, train, validate, and test independently"),
    ("science", "physics", "translate the situation into variables and laws, keep units consistent, solve symbolically, then sanity-check"),
    ("science", "chemistry", "identify substances and governing conservation laws, balance the reaction/model, calculate, and check stoichiometry"),
    ("science", "biology", "identify the mechanism, distinguish correlation from causation, connect evidence to the proposed biological process"),
    ("engineering", "systems", "decompose components and interfaces, identify bottlenecks/failure modes, design controls and verify requirements"),
    ("engineering", "electronics", "define voltage/current/power relationships, inspect component constraints, calculate worst case, verify margins"),
    ("logic", "deduction", "translate statements into premises, derive consequences, test alternatives, and distinguish valid inference from assumption"),
    ("logic", "constraint solving", "list variables/domains/constraints, propagate restrictions, branch only when needed, verify every constraint"),
    ("software engineering", "architecture", "identify requirements and boundaries, choose components/interfaces, evaluate tradeoffs, and define failure handling"),
    ("software engineering", "testing", "identify observable behavior, partition cases, include edge cases, assert invariants, and test regressions"),
    ("software engineering", "security review", "map trust boundaries, identify attack surfaces, validate inputs/authorization, and verify mitigations"),
    ("mathematics", "arithmetic", "break the problem into operations, preserve order of operations, calculate carefully, and independently verify the result"),
    ("mathematics", "fractions", "convert representations consistently, identify common denominators when needed, perform the operation, simplify, and verify"),
    ("mathematics", "ratios", "identify corresponding quantities, establish the ratio relationship, solve proportionally, and check whether the result has the correct scale"),
    ("mathematics", "percentages", "identify the reference quantity, distinguish percentage change from percentage points, calculate precisely, and sanity-check the magnitude"),
    ("mathematics", "word problems", "translate the natural-language description into variables and equations, solve systematically, and verify the answer against the wording"),
    ("mathematics", "linear equations", "isolate the unknown using equivalent operations, preserve equality at every step, and substitute the result back into the original equation"),
    ("mathematics", "systems of equations", "identify variables and equations, choose substitution/elimination/matrix methods appropriately, solve, and verify every equation"),
    ("mathematics", "quadratic equations", "identify the quadratic structure, select factoring/completing-the-square/quadratic-formula methods, solve all relevant roots, and verify"),
    ("mathematics", "inequalities", "transform both sides while preserving the inequality relationship, reverse the sign when multiplying or dividing by a negative, and check intervals"),
    ("mathematics", "sequences", "identify the generating pattern or recurrence, distinguish arithmetic/geometric/nonlinear behavior, derive the requested term, and test against known terms"),
    ("mathematics", "series", "identify whether the sequence is finite or infinite, select the appropriate summation method, apply convergence conditions when necessary, and verify"),
    ("mathematics", "exponents", "apply exponent laws consistently, simplify symbolic expressions, handle zero and negative exponents correctly, and verify with numerical substitution"),
    ("mathematics", "logarithms", "apply logarithm identities and domain restrictions, transform the equation into a solvable form, solve, and check the original domain"),
    ("mathematics", "functions", "identify inputs, outputs, domain, range, and transformations, apply the function definition, and verify composition or inverse relationships"),
    ("mathematics", "coordinate geometry", "translate geometric relationships into coordinate equations, calculate using distance/slope/midpoint formulas, and verify geometrically"),
    ("mathematics", "trigonometry", "identify the relevant triangle or periodic relationship, choose the correct identity or theorem, calculate carefully, and check angle/domain constraints"),
    ("mathematics", "vectors", "represent quantities as components, apply vector operations consistently, preserve dimensional meaning, and verify the magnitude or direction"),
    ("mathematics", "matrices", "identify dimensions and compatible operations, perform transformations carefully, and verify dimensions and results"),
    ("mathematics", "linear algebra", "represent the problem using vectors/matrices, identify the required transformation or system, solve, and verify algebraic consistency"),
    ("mathematics", "optimization", "define the objective and constraints, identify candidate extrema, evaluate boundary cases, and verify that the proposed optimum satisfies all constraints"),
    ("mathematics", "number theory", "identify divisibility, modular, prime, or factorization structure, exploit relevant properties, and verify the result"),
    ("mathematics", "combinatorics", "define what constitutes a distinct outcome, determine whether order and repetition matter, choose the appropriate counting principle, and verify with a smaller case"),
    ("mathematics", "permutations", "determine whether arrangement order matters, account for repeated elements when necessary, calculate systematically, and verify"),
    ("mathematics", "combinations", "determine whether order matters, identify the population and selection size, apply the correct counting formula, and sanity-check"),
    ("mathematics", "expected value", "define possible outcomes and probabilities, multiply each outcome by its probability, sum contributions, and verify the result lies within plausible bounds"),
    ("mathematics", "Bayesian reasoning", "define hypotheses and evidence, distinguish prior probability from likelihood, apply Bayes' rule correctly, and check normalization"),
    ("mathematics", "conditional probability", "define the conditioning event explicitly, restrict the sample space appropriately, calculate the conditional probability, and check independence assumptions"),
    ("mathematics", "discrete probability", "define the sample space, identify mutually exclusive and overlapping events, select the appropriate probability rule, and verify bounds"),
    ("mathematics", "continuous probability", "identify the density or distribution, determine the required interval or tail probability, integrate or use the appropriate distribution function, and check normalization"),
    ("mathematics", "statistical inference", "identify the parameter of interest, select an appropriate estimator or test, state assumptions, calculate uncertainty, and interpret cautiously"),
    ("mathematics", "data interpretation", "identify what the axes, units, sample, and statistic actually represent, calculate relevant comparisons, and avoid unsupported causal conclusions"),

    # --------------------------- LOGIC -----------------------------

    ("logic", "formal logic", "translate statements into propositions, identify premises and conclusions, apply valid inference rules, and distinguish implication from equivalence"),
    ("logic", "syllogisms", "identify premises, determine their logical structure, derive only conclusions supported by the premises, and test for counterexamples"),
    ("logic", "inductive reasoning", "identify observed patterns, generate the strongest supported generalization, distinguish evidence from certainty, and identify possible counterexamples"),
    ("logic", "abductive reasoning", "identify the observations, generate plausible explanations, compare how well each explains the evidence, and distinguish plausibility from proof"),
    ("logic", "counterexamples", "identify the universal claim being tested, search for the smallest valid counterexample, verify that it satisfies the premises, and use it to assess the claim"),
    ("logic", "contradiction", "identify mutually incompatible claims, trace the assumptions producing the conflict, and determine which premise or inference requires revision"),
    ("logic", "truth tables", "identify all relevant propositions, enumerate combinations systematically, evaluate operators correctly, and determine validity or equivalence"),
    ("logic", "causal reasoning", "distinguish correlation, temporal order, confounding, and mechanism, identify alternative explanations, and state what evidence would discriminate between them"),
    ("logic", "fallacy detection", "identify the exact claim and supporting argument, inspect whether the inference follows, and distinguish genuine logical errors from mere disagreement"),
    ("logic", "constraint reasoning", "identify variables, domains, hard constraints, and soft preferences, propagate restrictions, search systematically, and verify every solution"),
    ("logic", "invariant reasoning", "identify a property that remains unchanged under allowed operations, prove that the invariant holds, and use it to constrain possible outcomes"),
    ("logic", "case analysis", "partition the problem into exhaustive and mutually appropriate cases, solve each case, and ensure no case is omitted or double-counted"),
    ("logic", "proof construction", "state the proposition precisely, select direct proof/contradiction/induction/construction as appropriate, establish each implication, and check for hidden assumptions"),
    ("logic", "proof by induction", "define the base case, formulate the induction hypothesis precisely, prove the inductive step, and conclude only what the induction establishes"),
    ("logic", "necessary and sufficient conditions", "distinguish conditions that guarantee a result from those merely required for it, prove each direction separately, and identify edge cases"),

    # ---------------------- GENERAL PROBLEM SOLVING -----------------

    ("reasoning", "problem decomposition", "divide the problem into independent or dependent subproblems, solve the highest-leverage components first, then integrate and verify"),
    ("reasoning", "root cause analysis", "separate symptoms from causes, construct a causal chain, test competing explanations against evidence, and identify the deepest actionable cause"),
    ("reasoning", "first principles", "strip the problem down to fundamental facts and constraints, derive consequences from those foundations, and avoid relying on unsupported assumptions"),
    ("reasoning", "constraint analysis", "list hard constraints, soft constraints, resources, dependencies, and objectives before selecting a solution"),
    ("reasoning", "tradeoff analysis", "identify competing objectives, quantify or characterize the relevant tradeoffs, identify constraints, and explain what changes under different priorities"),
    ("reasoning", "scenario analysis", "define plausible scenarios, identify variables that differ between them, determine consequences, and avoid treating uncertain scenarios as established facts"),
    ("reasoning", "counterfactual reasoning", "define the hypothetical change precisely, hold unrelated conditions constant where appropriate, trace consequences, and distinguish inference from speculation"),
    ("reasoning", "hypothesis testing", "state a falsifiable hypothesis, identify predictions, gather relevant evidence, test alternatives, and update the conclusion based on evidence"),
    ("reasoning", "error analysis", "identify where an error could enter, determine its direction and magnitude, isolate the root cause, and estimate its effect on the final result"),
    ("reasoning", "uncertainty analysis", "identify uncertain inputs and assumptions, determine how they affect the conclusion, distinguish known from estimated values, and communicate uncertainty clearly"),
    ("reasoning", "estimation", "identify the quantity to estimate, choose reasonable reference values, simplify the problem, calculate an order-of-magnitude result, and sanity-check it"),
    ("reasoning", "Fermi estimation", "break an unknown quantity into estimable factors, use defensible approximations, calculate the order of magnitude, and identify the dominant uncertainty"),
    ("reasoning", "backward reasoning", "start from the desired outcome, determine the conditions required immediately before it, recursively establish prerequisites, and verify the resulting path"),
    ("reasoning", "forward reasoning", "start from established facts, apply valid rules step by step, track newly established information, and stop when the target is reached or impossible"),
    ("reasoning", "analogy reasoning", "identify the structural similarity between problems, map corresponding elements carefully, and verify that the analogy does not break under relevant differences"),
    ("reasoning", "simplification", "remove irrelevant complexity while preserving the essential structure, solve the simplified model, and determine whether the simplification changes the conclusion"),
    ("reasoning", "verification", "independently reconstruct the key result, test assumptions and edge cases, and compare the independent result with the original"),
    ("reasoning", "self-correction", "inspect a proposed solution for unsupported steps, locate inconsistencies, revise the weakest step, and rerun the affected reasoning"),
    ("reasoning", "ambiguity resolution", "identify multiple plausible interpretations, determine which are supported by context, state the interpretation used, and explain how the answer changes under alternatives"),
    ("reasoning", "information sufficiency", "identify exactly what information is required to answer the question, distinguish necessary from optional information, and avoid fabricating missing values"),
    ("reasoning", "decision under uncertainty", "identify possible actions, outcomes, probabilities or uncertainty ranges where available, costs and benefits, and sensitivity to assumptions"),
    ("reasoning", "prioritization", "identify objectives, urgency, impact, dependencies, and resource constraints, then construct a defensible ordering without confusing importance with urgency"),

    # -------------------------- SCIENCE ----------------------------

    ("science", "scientific method", "formulate a testable hypothesis, identify variables and controls, predict observations, evaluate evidence, and distinguish results from interpretation"),
    ("science", "experimental design", "define the research question, identify independent/dependent/control variables, minimize confounding, choose measurements, and define success criteria"),
    ("science", "measurement", "identify the quantity and units, distinguish precision from accuracy, track uncertainty, and verify dimensional consistency"),
    ("science", "dimensional analysis", "express quantities using fundamental units, cancel compatible dimensions, detect invalid equations, and verify the final units"),
    ("science", "physics problem solving", "translate the physical situation into variables and governing laws, solve symbolically where possible, substitute values, and check units and limiting cases"),
    ("science", "mechanics", "identify forces, motion variables, reference frames, and conservation laws, construct the governing equations, solve, and check physical plausibility"),
    ("science", "energy conservation", "identify the system and energy forms, account for transfers and losses, write the conservation equation, solve, and verify the energy balance"),
    ("science", "momentum", "define the system, identify external impulses, apply momentum conservation when justified, solve component-wise, and verify direction and magnitude"),
    ("science", "thermodynamics", "define the system and state variables, identify energy and entropy transfers, apply the appropriate law, and check physical constraints"),
    ("science", "electricity", "identify voltage, current, resistance, power, and circuit topology, apply governing relationships, and verify conservation and component limits"),
    ("science", "chemistry calculations", "identify substances and quantities, balance the chemical relationship, convert units through stoichiometry, calculate, and verify conservation"),
    ("science", "chemical equilibrium", "identify the reaction and equilibrium expression, account for concentrations/activities, apply the relevant equilibrium relationship, and check the direction of change"),
    ("science", "acid-base reasoning", "identify the acid/base system, determine relevant equilibria and concentrations, calculate pH or related quantities, and check approximations"),
    ("science", "biology mechanisms", "identify the biological entities and mechanism, connect cause to effect through intermediate processes, and distinguish established evidence from hypothesis"),
    ("science", "ecology", "identify populations, resources, interactions, and environmental constraints, model relevant relationships, and consider feedback and time scale"),
    ("science", "evolutionary reasoning", "identify variation, inheritance, selection pressures, and population-level effects, then distinguish individual adaptation from population change"),
    ("science", "scientific evidence evaluation", "identify study design, sample, measurement, controls, effect size, uncertainty, and limitations before interpreting the conclusion"),

    # ------------------------ ENGINEERING --------------------------

    ("engineering", "requirements analysis", "translate vague goals into measurable requirements, identify constraints and acceptance criteria, and distinguish mandatory requirements from preferences"),
    ("engineering", "failure mode analysis", "identify components and failure modes, estimate causes and effects, prioritize risks, and design mitigations and detection mechanisms"),
    ("engineering", "reliability", "identify failure mechanisms, estimate rates or probabilities, model system dependencies, and determine how redundancy or maintenance affects reliability"),
    ("engineering", "safety analysis", "identify hazards, exposure paths, severity and likelihood, protective controls, and residual risk, then verify that controls address the actual hazard"),
    ("engineering", "optimization", "define objective function and constraints, identify dominant variables, compare feasible designs, and verify robustness rather than optimizing a fragile nominal case"),
    ("engineering", "capacity planning", "estimate workload, resource demand, growth, bottlenecks, utilization, and safety margins, then test the design under peak conditions"),
    ("engineering", "systems thinking", "identify components, interfaces, feedback loops, dependencies, bottlenecks, and emergent behavior before proposing local fixes"),
    ("engineering", "control systems", "identify the controlled variable, actuator, sensor, feedback path, disturbance, and stability requirements, then analyze response and failure modes"),
    ("engineering", "signal analysis", "identify the signal, noise, sampling constraints, frequency characteristics, and measurement objectives, then select an appropriate analysis method"),
    ("engineering", "design validation", "translate requirements into measurable tests, construct representative test conditions, define pass/fail criteria, execute checks, and document limitations"),

    # ---------------------- COMPUTER SCIENCE ------------------------

    ("computer science", "recursion", "define the base case, identify the recursive reduction, prove progress toward termination, and verify returned values"),
    ("computer science", "search", "define the search space and target condition, choose an appropriate search strategy, preserve completeness/correctness, and analyze complexity"),
    ("computer science", "binary search", "establish the sorted invariant, maintain valid bounds, choose the midpoint safely, eliminate the impossible half, and verify termination"),
    ("computer science", "hashing", "identify the key and lookup requirements, choose a suitable hash structure, account for collisions and equality semantics, and analyze expected complexity"),
    ("computer science", "trees", "identify the tree invariant, choose traversal or update operations appropriate to the structure, preserve invariants, and analyze complexity"),
    ("computer science", "heaps", "identify priority-queue requirements, preserve the heap property after updates, choose appropriate operations, and analyze complexity"),
    ("computer science", "database reasoning", "identify entities, relationships, constraints, query requirements, consistency needs, and indexing opportunities before selecting a schema or query"),
    ("computer science", "distributed systems", "identify nodes, communication, state ownership, failure modes, consistency requirements, and recovery behavior before selecting a design"),
    ("computer science", "networking", "identify endpoints, protocols, addressing, transport behavior, failure points, and expected traffic, then trace the data path"),
    ("computer science", "operating systems", "identify processes, threads, memory, files, scheduling, synchronization, and system calls relevant to the problem"),
    ("computer science", "compiler reasoning", "separate lexical, syntactic, semantic, intermediate, optimization, and code-generation stages, then identify where an error or transformation occurs"),
    ("computer science", "formal verification", "state the invariant or specification precisely, identify the implementation obligation, prove relevant properties, and test assumptions"),
    ("computer science", "algorithm correctness", "state the invariant or correctness condition, show initialization, preservation, and termination where appropriate, and verify edge cases"),
    ("computer science", "amortized analysis", "identify operations whose costs vary, select an accounting/potential method, derive aggregate cost, and convert it to amortized complexity"),

    # ---------------------- SOFTWARE PRACTICE -----------------------

    ("software engineering", "requirements clarification", "separate explicit requirements from assumptions, identify ambiguities and edge cases, and define observable acceptance criteria"),
    ("software engineering", "code review", "inspect correctness, maintainability, security, performance, error handling, and compatibility, prioritizing concrete defects over stylistic preferences"),
    ("software engineering", "refactoring", "identify the behavior that must remain unchanged, isolate structural improvements, make small transformations, and run regression checks"),
    ("software engineering", "bug diagnosis", "reproduce the issue, minimize the failing case, inspect relevant state transitions, identify the root cause, patch the cause, and retest"),
    ("software engineering", "API compatibility", "identify the existing contract, inputs, outputs, errors, side effects, and version constraints, then preserve compatibility or document deliberate changes"),
    ("software engineering", "test design", "identify behavior partitions, boundary conditions, invalid inputs, state transitions, and regressions, then construct tests with clear expected outcomes"),
    ("software engineering", "performance diagnosis", "measure before changing code, identify the dominant bottleneck, distinguish CPU/memory/I/O/contention costs, optimize the bottleneck, and benchmark again"),
    ("software engineering", "security reasoning", "identify assets, trust boundaries, attacker capabilities, entry points, authorization requirements, and failure modes, then verify mitigations"),
    ("software engineering", "maintainability", "identify coupling, cohesion, duplication, unclear contracts, and hidden dependencies, then improve structure while preserving observable behavior"),
    ("software engineering", "documentation reasoning", "identify the intended audience and task, explain prerequisites and behavior accurately, include examples where useful, and avoid documenting assumptions as guarantees"),

    # --------------------- LANGUAGE & COMMUNICATION -----------------

    ("language", "reading comprehension", "identify the main claim, supporting evidence, qualifications, assumptions, and conclusions before answering questions about the text"),
    ("language", "text analysis", "separate explicit statements from implications, identify structure and evidence, and avoid attributing claims that the text does not support"),
    ("language", "summarization", "identify the central information and supporting points, remove redundancy, preserve important qualifications, and avoid introducing new claims"),
    ("language", "argument analysis", "identify claims, premises, evidence, assumptions, counterarguments, and conclusions, then evaluate whether the reasoning follows"),
    ("language", "instruction following", "extract explicit requirements, ordering constraints, formatting requirements, exclusions, and success conditions before producing the requested result"),
    ("language", "ambiguity analysis", "identify words or constructions with multiple plausible meanings, use context to distinguish them, and state remaining ambiguity when material"),
    ("language", "analogy explanation", "map an unfamiliar concept onto a familiar structure while explicitly identifying where the analogy holds and where it breaks"),
    ("language", "explanation", "identify the audience's likely knowledge level, introduce concepts in dependency order, use precise examples, and verify that the explanation answers the actual question"),
    ("language", "technical communication", "separate facts, assumptions, procedures, results, and caveats, use precise terminology, and structure information for easy verification"),
    ("language", "translation reasoning", "preserve meaning, intent, register, ambiguity, and important technical terminology rather than translating words mechanically"),
]

_VARIANTS = [
    "direct solution with explicit assumptions",
    "edge-case-heavy solution",
    "debugging-oriented solution",
    "optimization-oriented solution",
    "verification-oriented solution",
    "comparison of two plausible approaches",
    "failure-analysis and correction",
    "design followed by a validation plan",
    "symbolic derivation followed by a numerical sanity check",
    "algorithm selection followed by complexity analysis",
]


def _build_reasoning_examples(target: int = EXAMPLE_TARGET) -> List[Dict[str, Any]]:
    """Create a deterministic library of structured reasoning demonstrations.
    These are inference-time demonstrations rather than weight updates. The
    generator deliberately stores reasoning as concise, inspectable steps,
    not unrestricted hidden chain-of-thought.
    """
    examples: List[Dict[str, Any]] = []
    idx = 0
    skill_count = len(_REASONING_SKILLS)

    while len(examples) < target:
        domain, skill, principle = _REASONING_SKILLS[idx % skill_count]
        variant = _VARIANTS[(idx // skill_count) % len(_VARIANTS)]
        difficulty = ("medium", "hard", "expert")[(idx // (skill_count * 3)) % 3]

        problem = (
            f"Solve a {skill} problem using a {variant} approach. "
            f"The solution must apply the core principle: {principle}."
        )

        steps = [
            f"Identify the exact {skill} task and separate requirements from assumptions.",
            f"Apply the governing principle: {principle}.",
            "Work through the relevant intermediate states without skipping constraint checks.",
            "Inspect edge cases, contradictions, units, invariants, or boundary conditions as applicable.",
            "Produce the solution and independently verify that it satisfies the original requirements.",
        ]

        verification = [
            "Check that every stated constraint is satisfied.",
            "Check for an unstated assumption that could change the result.",
            "Use an independent sanity check where possible.",
        ]

        solution = (
            f"Use the {skill} method appropriate to the stated constraints; explicitly state assumptions, "
            "derive the result in ordered steps, then verify the result against the original problem."
        )

        examples.append({
            "id": idx + 1,
            "domain": domain,
            "skill": skill,
            "difficulty": difficulty,
            "problem": problem,
            "reasoning_trace": steps,
            "verification": verification,
            "solution": solution,
        })
        idx += 1

    return examples


# IMPORTANT: this is the actual definition that was missing in the broken file.
REASONING_EXAMPLES: List[Dict[str, Any]] = _build_reasoning_examples(EXAMPLE_TARGET)


def _format_reasoning_examples(query: str, top_k: int = 6) -> str:
    """Return the most relevant structured demonstrations for a query."""
    if not REASONING_EXAMPLES:
        return "No demonstrations available."

    embed_model = STATE.get("embed_model")
    if embed_model is None:
        # Safe fallback before model initialization.
        selected = REASONING_EXAMPLES[:top_k]
    else:
        try:
            if STATE.get("reasoning_example_embeddings") is None:
                texts = [
                    f"{x['domain']} {x['skill']} {x['problem']} {x['solution']}"
                    for x in REASONING_EXAMPLES
                ]
                STATE["reasoning_example_embeddings"] = embed_model.encode(
                    texts,
                    normalize_embeddings=True,
                    convert_to_numpy=True,
                    show_progress_bar=False,
                )

            qv = _embed(query)
            matrix = STATE["reasoning_example_embeddings"]
            scores = matrix @ qv
            indices = np.argsort(scores)[::-1][:top_k]
            selected = [REASONING_EXAMPLES[int(i)] for i in indices]
        except Exception as exc:
            logger.warning("Reasoning example retrieval failed: %s", exc)
            selected = REASONING_EXAMPLES[:top_k]

    blocks = []
    for ex in selected:
        blocks.append(
            "EXAMPLE {id} | {domain}/{skill} | difficulty={difficulty}\n"
            "Problem pattern: {problem}\n"
            "Reasoning steps:\n- {steps}\n"
            "Verification:\n- {verification}\n"
            "Solution pattern: {solution}".format(
                id=ex["id"],
                domain=ex["domain"],
                skill=ex["skill"],
                difficulty=ex["difficulty"],
                problem=ex["problem"],
                steps="\n- ".join(ex["reasoning_trace"]),
                verification="\n- ".join(ex["verification"]),
                solution=ex["solution"],
            )
        )
    return "\n\n".join(blocks)


# =====================================================================
# MODEL
# =====================================================================


def initialize_model():
    """Load the language and embedding models once."""
    if STATE.get("model") is not None:
        return

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    logger.info("Loading reasoning model: %s", MODEL_NAME)

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    config = AutoConfig.from_pretrained(MODEL_NAME)
    config.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        config=config,
        torch_dtype=dtype,
        device_map="auto" if device == "cuda" else None,
    ).eval()

    embed_model = SentenceTransformer(
        "sentence-transformers/all-MiniLM-L6-v2",
        device=device,
    )

    STATE.update(
        tokenizer=tokenizer,
        model=model,
        embed_model=embed_model,
        device=device,
    )

    initialize_memory()
    logger.info("Reasoning engine initialized on %s with %d demonstrations", device, len(REASONING_EXAMPLES))


# =====================================================================
# EXPERIENCE MEMORY
# =====================================================================


def initialize_memory():
    MEMORY_DB.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(MEMORY_DB) as db:
        db.execute(
            """
            CREATE TABLE IF NOT EXISTS experiences (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                query TEXT NOT NULL,
                plan TEXT NOT NULL,
                analysis TEXT NOT NULL,
                answer TEXT NOT NULL,
                quality REAL DEFAULT 0.0,
                created REAL NOT NULL
            )
            """
        )
        db.commit()


def _embed(text: str) -> np.ndarray:
    model = STATE["embed_model"]
    return model.encode(
        [text],
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )[0]

def _detect_and_count_characters(query: str) -> str:
    """Programmatically extracts words from the query and builds an infallible character breakdown."""
    words = re.findall(r'["\']?(\b[a-zA-Z]+)\b["\']?', query)
    if not words:
        return ""
    
    lines = ["\n[SYSTEM GROUND-TRUTH LINGUISTIC EVIDENCE]"]
    for word in sorted(list(set(words)), key=len, reverse=True):
        if len(word) >= 3:  # Focus on meaningful potential target words
            w_lower = word.lower()
            counts = {}
            for char in w_lower:
                counts[char] = counts.get(char, 0) + 1
            breakdown = ", ".join([f"'{c}': {cnt}" for c, cnt in sorted(counts.items())])
            lines.append(f"The word '{word}' has a length of {len(word)} letters. Exact character breakdown: {breakdown}")
    return "\n".join(lines) if len(lines) > 1 else ""

def _ensure_experience_embedding_column(conn: sqlite3.Connection) -> None:
    """Ensure the experience table has a BLOB column for precomputed embeddings."""
    columns = {
        row[1] for row in conn.execute("PRAGMA table_info(experiences)").fetchall()
    }
    if "embedding" not in columns:
        conn.execute("ALTER TABLE experiences ADD COLUMN embedding BLOB")
        conn.commit()


def _embedding_to_blob(vector: Optional[np.ndarray]) -> Optional[bytes]:
    if vector is None:
        return None
    arr = np.asarray(vector, dtype=np.float32).reshape(-1)
    return arr.tobytes()


def _blob_to_embedding(blob: Optional[bytes]) -> Optional[np.ndarray]:
    if blob is None:
        return None
    return np.frombuffer(blob, dtype=np.float32).copy()


def save_experience(query: str, plan: str, analysis: str, answer: str, quality: float = 1.0):
    try:
        with sqlite3.connect(MEMORY_DB) as db:
            db.execute(
                """
                INSERT INTO experiences
                (query, plan, analysis, answer, quality, created)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (query, plan, analysis, answer, float(quality), time.time()),
            )
            db.commit()
    except Exception as exc:
        logger.warning("Could not save experience: %s", exc)


def retrieve_experiences(query: str, top_k: int = 3) -> List[Dict[str, Any]]:
    try:
        with sqlite3.connect(MEMORY_DB) as db:
            rows = db.execute(
                """
                SELECT id, query, plan, analysis, answer, quality, created
                FROM experiences
                WHERE quality >= 0.5
                ORDER BY created DESC
                LIMIT 500
                """
            ).fetchall()

        if not rows:
            return []

        # 1. Get the embedding model from STATE
        embed_model = STATE.get("embed_model")
        if embed_model is None:
            # Fallback if model isn't initialized yet
            return []

        # 2. Embed the incoming user query
        qv = _embed(query)

        # 3. Batch embed all historical queries retrieved from the database
        db_queries = [row[1] for row in rows]
        rv_matrix = embed_model.encode(
            db_queries,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False
        )

        # 4. Use matrix multiplication to calculate all scores simultaneously
        scores = rv_matrix @ qv

        # 5. Pair scores back up with their rows and sort
        scored = [(float(score), row) for score, row in zip(scores, rows)]
        scored.sort(key=lambda x: x[0], reverse=True)

        return [
            {
                "similarity": score,
                "query": row[1],
                "plan": row[2],
                "analysis": row[3],
                "answer": row[4],
                "quality": row[5],
            }
            # Slid down to only return the top_k best matches
            for score, row in scored[:top_k]
        ]
    except Exception as exc:
        logger.warning("Experience retrieval failed: %s", exc)
        return []


# =====================================================================
# GENERATION + TRACE EXTRACTION
# =====================================================================
def generate(messages: List[Dict[str, str]], max_new_tokens: int = 1024, temperature: float = 0.15) -> str:
    """Run one model pass, forcing DeepSeek-R1 style reasoning via prompt pre-filling."""
    initialize_model()

    tokenizer = STATE["tokenizer"]
    model = STATE["model"]

    # 1. Apply standard chat template
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # 2. Force the model to start inside the thinking block
    prompt += "<think>\n"
    
    inputs = tokenizer(prompt, return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    generation_kwargs = {
        "max_new_tokens": max_new_tokens,
        "repetition_penalty": 1.05,
        "pad_token_id": tokenizer.eos_token_id,
    }
    if temperature > 0:
        generation_kwargs.update({
            "temperature": temperature,
            "top_p": 0.90,
            "do_sample": True
        })
    else:
        generation_kwargs.update({"do_sample": False})

    with torch.no_grad():
        output = model.generate(**inputs, **generation_kwargs)

    new_tokens = output[0][inputs["input_ids"].shape[1]:]
    generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    
    # 3. Prepend the token we injected so the regex parser handles it perfectly
    return f"<think>\n{generated_text}"


def _extract_think(text: str) -> str:
    """Extract the model's explicit <think> section if it emitted one."""
    match = re.search(r"<think>\s*(.*?)\s*</think>", text, flags=re.IGNORECASE | re.DOTALL)
    return match.group(1).strip() if match else ""


def clean_think_block(text: str) -> str:
    """Return a bounded visible reasoning trace plus the model output.
    We intentionally expose a structured reasoning summary rather than
    pretending that an arbitrary generated <think> block is a faithful dump
    of private internal cognition.
    """
    text = text.strip()
    think = _extract_think(text)
    body = re.sub(r"<think>.*?</think>", "", text, flags=re.IGNORECASE | re.DOTALL).strip()

    if not body:
        body = text.replace("<think>", "").replace("</think>", "").strip()

    if SHOW_REASONING_TRACE and think:
        return f"<think>\n{think}\n</think>\n\n{body}".strip()
    return body


def _make_trace(stage: str, text: str) -> Dict[str, Any]:
    """Create an inspectable stage trace from the explicit model output."""
    think = _extract_think(text)
    body = re.sub(r"<think>.*?</think>", "", text, flags=re.IGNORECASE | re.DOTALL).strip()
    return {
        "stage": stage,
        "reasoning_summary": think,
        "output": body,
    }


# =====================================================================
# COGNITIVE PASSES
# =====================================================================
PLANNER_SYSTEM = r"""
You are the cognitive planner of Infinity GRM.
First, think deeply about the question to yourself inside a <think>...</think> block. Reason completely naturally, fluidly, and without any structured steps, numbers, or headings. Question assumptions, look for pitfalls, and figure out how to approach the answer.
After closing the </think> block, immediately output a JSON object with execution details:
{
  "task": "...",
  "requires_web": true/false,
  "search_queries": ["..."],
  "steps": ["..."],
  "verification": ["..."],
  "difficulty": "easy|medium|hard|expert"
}
"""

ANALYST_SYSTEM = r"""
You are the analysis component of Infinity GRM.
First, think deeply about the question and evidence to yourself inside a <think>...</think> block. Reason completely naturally, fluidly, and without any structured steps, numbers, or headings. Weigh the facts, check calculations, and trace what is true.
After closing the </think> block, provide your analyzed conclusion.
"""

VERIFIER_SYSTEM = r"""
You are the verification component of Infinity GRM.
First, think deeply about the proposed analysis to yourself inside a <think>...</think> block. Reason completely naturally, fluidly, and without any structured steps, numbers, or headings. Play devil's advocate, look for logical gaps, and double-check the calculations. If computer programming is involved make sure all lines make sense and are rendered correctly.
After closing the </think> block, provide a concise verification report.
"""

FINAL_SYSTEM = r"""
You are the final synthesis component of Infinity GRM.
First, think deeply about how to assemble the definitive response to yourself inside a <think>...</think> block. Reason completely naturally, fluidly, and without any structured steps, numbers, or headings. Focus entirely on how to answer the user perfectly, directly, and comprehensively.
After closing the </think> block, present the final answer to the user.
"""

def plan(query: str, history: List[Dict[str, str]]) -> Dict[str, Any]:
    # FIX: Initialize models here so retrieve_experiences has access to the embedding model
    initialize_model()
    
    experiences = retrieve_experiences(query, top_k=3)
    experience_text = json.dumps(experiences, ensure_ascii=False, indent=2)

    messages = [
        {"role": "system", "content": PLANNER_SYSTEM},
        {
            "role": "user",
            "content": (
                f"QUESTION:\n{query}\n\n"
                f"HISTORY:\n{json.dumps(history[-6:], ensure_ascii=False)}\n\n"
                f"PRIOR EXPERIENCES:\n{experience_text}\n\n"
                f"RELEVANT REASONING DEMONSTRATIONS:\n{_format_reasoning_examples(query, top_k=8)}"
            ),
        },
    ]

    raw = generate(messages, max_new_tokens=1400, temperature=0.10)
    think = _extract_think(raw)
    json_match = re.search(r"\{.*\}", raw, flags=re.DOTALL)

    if json_match:
        try:
            data = json.loads(json_match.group(0))
            data["think"] = think
            data["experiences"] = experiences
            return data
        except json.JSONDecodeError:
            pass

    return {
        "think": think,
        "task": query,
        "requires_web": False,
        "search_queries": [],
        "steps": ["Clarify the task", "Solve using the most appropriate method", "Verify the result"],
        "verification": ["Check assumptions and edge cases"],
        "difficulty": "hard",
        "experiences": experiences,
    }

def analyze(query: str, plan_data: Dict[str, Any], evidence: str, history: List[Dict[str, str]]) -> str:
    # Inject programmatic linguistic truth into evidence
    truth = _detect_and_count_characters(query)
    if truth:
        evidence = f"{evidence}\n{truth}"

    messages = [
        {"role": "system", "content": ANALYST_SYSTEM},
        {
            "role": "user",
            "content": (
                f"QUESTION:\n{query}\n\n"
                f"PLAN:\n{json.dumps(plan_data, indent=2)}\n\n"
                f"EVIDENCE:\n{evidence[:20000]}\n\n"
                f"HISTORY:\n{json.dumps(history[-6:], ensure_ascii=False)}\n\n"
                f"RELEVANT REASONING DEMONSTRATIONS:\n{_format_reasoning_examples(query, top_k=6)}"
            ),
        },
    ]
    return generate(messages, max_new_tokens=1600, temperature=0.10)


def verify(query: str, plan_data: Dict[str, Any], analysis: str, evidence: str) -> str:
    """Independent verification pass."""
    # Inject programmatic linguistic truth into evidence for verifier verification
    truth = _detect_and_count_characters(query)
    if truth:
        evidence = f"{evidence}\n{truth}"

    messages = [
        {"role": "system", "content": VERIFIER_SYSTEM},
        {
            "role": "user",
            "content": (
                f"QUESTION:\n{query}\n\n"
                f"PLAN:\n{json.dumps(plan_data, indent=2)}\n\n"
                f"ANALYSIS:\n{analysis[:18000]}\n\n"
                f"EVIDENCE:\n{evidence[:20000]}"
            ),
        },
    ]
    return generate(messages, max_new_tokens=1400, temperature=0.05)


def final_answer(
    query: str,
    plan_data: Dict[str, Any],
    analysis: str,
    evidence: str,
    history: List[Dict[str, str]],
    verification: Optional[str] = None,
) -> str:
    """Final reasoning/synthesis pass."""
    # Ensure final synthesis also sees the correct counts unmistakably
    truth = _detect_and_count_characters(query)
    if truth:
        evidence = f"{evidence}\n{truth}"

    verification_text = verification or "No separate verification pass was supplied."
    messages = [
        {"role": "system", "content": FINAL_SYSTEM},
        {
            "role": "user",
            "content": (
                f"QUESTION:\n{query}\n\n"
                f"PLAN:\n{json.dumps(plan_data, indent=2)}\n\n"
                f"ANALYSIS:\n{analysis[:18000]}\n\n"
                f"VERIFICATION:\n{verification_text[:16000]}\n\n"
                f"EVIDENCE:\n{evidence[:20000]}\n\n"
                f"HISTORY:\n{json.dumps(history[-8:], ensure_ascii=False)}\n\n"
                f"RELEVANT REASONING DEMONSTRATIONS:\n{_format_reasoning_examples(query, top_k=6)}"
            ),
        },
    ]
    return generate(messages, max_new_tokens=2200, temperature=0.10)


def build_reasoning_trace(
    query: str,
    plan_data: Dict[str, Any],
    analysis: str,
    verification: str,
    final: str,
) -> Dict[str, Any]:
    """Return a detailed visible trace of the reasoning pipeline."""
    return {
        "query": query,
        "demonstration_count": len(REASONING_EXAMPLES),
        "stages": [
            {
                "stage": "planning",
                "reasoning_summary": plan_data.get("think", ""),
                "structured_plan": plan_data,
            },
            _make_trace("analysis", analysis),
            _make_trace("verification", verification),
            _make_trace("final_synthesis", final),
        ],
    }


def learn_from_result(
    query: str,
    plan_data: Dict[str, Any],
    analysis: str,
    answer: str,
    quality: float = 1.0,
):
    initialize_model()
    save_experience(
        query=query,
        plan=json.dumps(plan_data, ensure_ascii=False),
        analysis=analysis,
        answer=answer,
        quality=quality,
    )


def initialize():
    initialize_model()


if __name__ == "__main__":
    initialize()
    print("Infinity GRM reasoning engine initialized.")
    print(f"Structured reasoning demonstrations: {len(REASONING_EXAMPLES)}")