Update reasoning_engine.py

reasoning_engine.py

@@ -2,31 +2,26 @@
 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
-- Free-form reasoning -> independent final synthesis
+- Planning -> analysis -> verification -> final synthesis
 - Persistent SQLite experience memory
-- Optional detailed, user-visible reasoning TRACE
+- 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.
-
-Reasoning behavior
-------------------
-The reasoning pass is intentionally NOT a fixed step-by-step procedure.
-The model is encouraged to explore the problem freely, reconsider
-assumptions, try alternatives, correct mistakes, and verify conclusions.
-
-The final answer is generated in a separate model pass so the model's
-reasoning is not simply copied into the final response.
 """
 
 import os
@@ -43,68 +38,12 @@ 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"}
-)
-
-# Larger budget for the free-form reasoning pass.
-THINK_MAX_NEW_TOKENS = int(
-    os.getenv(
-        "GRM_THINK_MAX_NEW_TOKENS",
-        "8192",
-    )
-)
-
-# Separate, smaller budget for the clean final answer.
-FINAL_MAX_NEW_TOKENS = int(
-    os.getenv(
-        "GRM_FINAL_MAX_NEW_TOKENS",
-        "2200",
-    )
-)
-
-THINK_TEMPERATURE = float(
-    os.getenv(
-        "GRM_THINK_TEMPERATURE",
-        "0.10",
-    )
-)
-
-FINAL_TEMPERATURE = float(
-    os.getenv(
-        "GRM_FINAL_TEMPERATURE",
-        "0.10",
-    )
-)
-
+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,
@@ -125,34 +64,27 @@ _REASONING_SKILLS = [
     ("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"),
@@ -185,6 +117,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -201,6 +135,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -224,6 +160,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -242,6 +180,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -253,6 +193,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -268,6 +210,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -279,6 +223,8 @@ _REASONING_SKILLS = [
     ("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"),
@@ -291,7 +237,6 @@ _REASONING_SKILLS = [
     ("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",
@@ -306,36 +251,20 @@ _VARIANTS = [
 ]
 
 
-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.
+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
-        ]
+        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. "
@@ -357,118 +286,79 @@ def _build_reasoning_examples(
         ]
 
         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,
-            }
+            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
 
 
-REASONING_EXAMPLES: List[Dict[str, Any]] = (
-    _build_reasoning_examples(EXAMPLE_TARGET)
-)
+# 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.
-    """
-
+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']} "
-                        f"{x['skill']} "
-                        f"{x['problem']} "
-                        f"{x['solution']}"
-                    )
+                    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,
-                    )
+                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"
-            ]
-
+            matrix = STATE["reasoning_example_embeddings"]
             scores = matrix @ qv
-
             indices = np.argsort(scores)[::-1][:top_k]
-
-            selected = [
-                REASONING_EXAMPLES[int(i)]
-                for i in indices
-            ]
-
+            selected = [REASONING_EXAMPLES[int(i)] for i in indices]
         except Exception as exc:
-            logger.warning(
-                "Reasoning example retrieval failed: %s",
-                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 guidance:\n- {steps}\n"
-            "Verification guidance:\n- {verification}\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"]
-                ),
+                steps="\n- ".join(ex["reasoning_trace"]),
+                verification="\n- ".join(ex["verification"]),
                 solution=ex["solution"],
             )
         )
-
     return "\n\n".join(blocks)
 
 
@@ -478,52 +368,27 @@ def _format_reasoning_examples(
 
 
 def initialize_model():
-    """
-    Load the language and embedding models once.
-    """
-
+    """Load the language and embedding models once."""
     if STATE.get("model") is not None:
         return
 
-    device = (
-        "cuda"
-        if torch.cuda.is_available()
-        else "cpu"
-    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if device == "cuda" else torch.float32
 
-    dtype = (
-        torch.float16
-        if device == "cuda"
-        else torch.float32
-    )
-
-    logger.info(
-        "Loading reasoning model: %s",
-        MODEL_NAME,
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME
-    )
+    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 = 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
-        ),
+        device_map="auto" if device == "cuda" else None,
     ).eval()
 
     embed_model = SentenceTransformer(
@@ -539,12 +404,7 @@ def initialize_model():
     )
 
     initialize_memory()
-
-    logger.info(
-        "Reasoning engine initialized on %s with %d demonstrations",
-        device,
-        len(REASONING_EXAMPLES),
-    )
+    logger.info("Reasoning engine initialized on %s with %d demonstrations", device, len(REASONING_EXAMPLES))
 
 
 # =====================================================================
@@ -553,11 +413,7 @@ def initialize_model():
 
 
 def initialize_memory():
-    MEMORY_DB.parent.mkdir(
-        parents=True,
-        exist_ok=True,
-    )
-
+    MEMORY_DB.parent.mkdir(parents=True, exist_ok=True)
     with sqlite3.connect(MEMORY_DB) as db:
         db.execute(
             """
@@ -572,13 +428,11 @@ def initialize_memory():
             )
             """
         )
-
         db.commit()
 
 
 def _embed(text: str) -> np.ndarray:
     model = STATE["embed_model"]
-
     return model.encode(
         [text],
         normalize_embeddings=True,
@@ -587,64 +441,30 @@ def _embed(text: str) -> np.ndarray:
     )[0]
 
 
-def _ensure_experience_embedding_column(
-    conn: sqlite3.Connection,
-) -> None:
-    """
-    Ensure the experience table has a BLOB column for precomputed embeddings.
-    """
-
+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()
+        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.execute("ALTER TABLE experiences ADD COLUMN embedding BLOB")
         conn.commit()
 
 
-def _embedding_to_blob(
-    vector: Optional[np.ndarray],
-) -> Optional[bytes]:
-
+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)
-
+    arr = np.asarray(vector, dtype=np.float32).reshape(-1)
     return arr.tobytes()
 
 
-def _blob_to_embedding(
-    blob: Optional[bytes],
-) -> Optional[np.ndarray]:
-
+def _blob_to_embedding(blob: Optional[bytes]) -> Optional[np.ndarray]:
     if blob is None:
         return None
+    return np.frombuffer(blob, dtype=np.float32).copy()
 
-    return np.frombuffer(
-        blob,
-        dtype=np.float32,
-    ).copy()
 
-
-def save_experience(
-    query: str,
-    plan: str,
-    analysis: str,
-    answer: str,
-    quality: float = 1.0,
-):
+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(
@@ -653,42 +473,19 @@ def save_experience(
                 (query, plan, analysis, answer, quality, created)
                 VALUES (?, ?, ?, ?, ?, ?)
                 """,
-                (
-                    query,
-                    plan,
-                    analysis,
-                    answer,
-                    float(quality),
-                    time.time(),
-                ),
+                (query, plan, analysis, answer, float(quality), time.time()),
             )
-
             db.commit()
-
     except Exception as exc:
-        logger.warning(
-            "Could not save experience: %s",
-            exc,
-        )
-
+        logger.warning("Could not save experience: %s", exc)
 
-def retrieve_experiences(
-    query: str,
-    top_k: int = 3,
-) -> List[Dict[str, Any]]:
 
+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
+                SELECT id, query, plan, analysis, answer, quality, created
                 FROM experiences
                 WHERE quality >= 0.5
                 ORDER BY created DESC
@@ -699,42 +496,30 @@ def retrieve_experiences(
         if not rows:
             return []
 
-        embed_model = STATE.get(
-            "embed_model"
-        )
-
+        # 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)
 
-        db_queries = [
-            row[1]
-            for row in rows
-        ]
-
+        # 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,
+            show_progress_bar=False
         )
 
+        # 4. Use matrix multiplication to calculate all scores simultaneously
         scores = rv_matrix @ qv
 
-        scored = [
-            (
-                float(score),
-                row,
-            )
-            for score, row
-            in zip(scores, rows)
-        ]
-
-        scored.sort(
-            key=lambda x: x[0],
-            reverse=True,
-        )
+        # 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 [
             {
@@ -745,178 +530,75 @@ def retrieve_experiences(
                 "answer": row[4],
                 "quality": row[5],
             }
-            for score, row
-            in scored[:top_k]
+            # 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,
-        )
-
+        logger.warning("Experience retrieval failed: %s", exc)
         return []
 
 
 # =====================================================================
-# GENERATION
+# GENERATION + TRACE EXTRACTION
 # =====================================================================
 
 
-def generate(
-    messages: List[Dict[str, str]],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.15,
-) -> str:
-    """
-    Run one model pass.
-
-    This function intentionally does not impose a reasoning format.
-    The caller decides whether the pass is a thinking pass or a final
-    answer pass.
-    """
-
+def generate(messages: List[Dict[str, str]], max_new_tokens: int = 1024, temperature: float = 0.15) -> str:
+    """Run one model pass."""
     initialize_model()
 
     tokenizer = STATE["tokenizer"]
     model = STATE["model"]
 
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-    )
-
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-    )
-
-    inputs = {
-        k: v.to(model.device)
-        for k, v in inputs.items()
-    }
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    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,
-            }
-        )
-
+        generation_kwargs.update({"temperature": temperature, "top_p": 0.90, "do_sample": True})
     else:
-        generation_kwargs.update(
-            {
-                "do_sample": False,
-            }
-        )
+        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]:
-    ]
-
-    return tokenizer.decode(
-        new_tokens,
-        skip_special_tokens=True,
-    ).strip()
-
-
-# =====================================================================
-# THINKING / OUTPUT UTILITIES
-# =====================================================================
-
-
-def _extract_think(
-    text: str,
-) -> str:
-    """
-    Extract an explicit <think> section if one exists.
-    """
-
-    match = re.search(
-        r"<think>\s*(.*?)\s*</think>",
-        text,
-        flags=re.IGNORECASE | re.DOTALL,
-    )
-
-    return (
-        match.group(1).strip()
-        if match
-        else ""
-    )
+        output = model.generate(**inputs, **generation_kwargs)
 
+    new_tokens = output[0][inputs["input_ids"].shape[1]:]
+    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
-def _remove_think(
-    text: str,
-) -> str:
-    return re.sub(
-        r"<think>.*?</think>",
-        "",
-        text,
-        flags=re.IGNORECASE | re.DOTALL,
-    ).strip()
 
+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:
-    """
-    Preserve a visible reasoning block when the model explicitly emits one.
 
-    This function does not invent reasoning that the model did not generate.
+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 = _remove_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()
+        body = text.replace("<think>", "").replace("</think>", "").strip()
 
     if SHOW_REASONING_TRACE and think:
-        return (
-            f"<think>\n"
-            f"{think}\n"
-            f"</think>\n\n"
-            f"{body}"
-        ).strip()
-
+        return f"<think>\n{think}\n</think>\n\n{body}".strip()
     return body
 
 
-def _make_trace(
-    stage: str,
-    text: str,
-) -> Dict[str, Any]:
-
+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 = _remove_think(text)
-
+    body = re.sub(r"<think>.*?</think>", "", text, flags=re.IGNORECASE | re.DOTALL).strip()
     return {
         "stage": stage,
         "reasoning_summary": think,
@@ -925,373 +607,164 @@ def _make_trace(
 
 
 # =====================================================================
-# FREE-FORM THINKING PASS
+# COGNITIVE PASSES
 # =====================================================================
 
-
-THINK_SYSTEM = r"""
-You are the free-form reasoning engine of Infinity GRM.
-
-Your job is to think deeply about the user's problem before a separate
-model pass produces the final answer.
-
-Do NOT follow a predefined step-by-step procedure.
-
-Do NOT use a fixed checklist.
-
-Do NOT force your reasoning into stages such as:
-1. interpretation
-2. planning
-3. analysis
-4. verification
-5. synthesis
-
-There is no required order.
-
-Instead, reason naturally according to whatever the problem requires.
-
-You may:
-- explore an idea;
-- abandon an approach;
-- reconsider an assumption;
-- notice a contradiction;
-- try a different method;
-- compare alternatives;
-- perform calculations;
-- test an intermediate result;
-- discover a useful relationship;
-- correct an earlier mistake;
-- backtrack;
-- verify something;
-- have an insight that changes the direction of the solution.
-
-The reasoning should adapt to the problem.
-
-Easy problems may require very little reasoning.
-Difficult problems may require substantial exploration.
-
-Do not artificially make the reasoning longer just to appear intelligent.
-
-Your output is the reasoning/workspace for a separate final-answer
-generation pass.
-
-Do NOT write a polished final answer.
-
-Do NOT repeat the user's question unnecessarily.
-
-Do NOT end with a separate "final answer" section.
-
-Think freely and deeply, then stop when you have enough information
-for another model pass to produce the answer.
+PLANNER_SYSTEM = r"""
+You are the planning component of Infinity GRM.
+Produce a detailed but concise, inspectable reasoning trace. This is a
+user-visible explanation of decisions, not a claim that it exposes private
+hidden cognition.
+You must:
+- identify the actual task;
+- identify constraints and assumptions;
+- determine whether external information is required;
+- identify useful subproblems;
+- choose appropriate reasoning methods;
+- identify what should be verified;
+- avoid inventing facts.
+Return:
+<think>
+1. Task interpretation
+2. Key assumptions/constraints
+3. Chosen method and why
+4. Subproblem decomposition
+5. Verification plan
+</think>
+Then return JSON with:
+{
+  "task": "...",
+  "requires_web": true/false,
+  "search_queries": ["..."],
+  "steps": ["..."],
+  "verification": ["..."],
+  "difficulty": "easy|medium|hard|expert"
+}
 """
 
+ANALYST_SYSTEM = r"""
+You are the evidence-analysis component of Infinity GRM.
+Give a detailed, inspectable reasoning summary. Do not merely restate the
+answer. Explicitly identify evidence, assumptions, contradictions, missing
+information, calculations, and verification checks.
+Return:
+<think>
+1. What the evidence establishes
+2. What it does not establish
+3. Important calculations/inferences
+4. Contradictions or uncertainty
+5. Verification result
+</think>
+Then provide the analyzed conclusion.
+"""
 
-def think(
-    query: str,
-    history: List[Dict[str, str]],
-    context: str = "",
-) -> str:
-    """
-    Free-form reasoning pass.
-
-    This is deliberately independent of plan/analyze/verify stages.
-    """
-
-    initialize_model()
-
-    experiences = retrieve_experiences(
-        query,
-        top_k=3,
-    )
-
-    experience_text = json.dumps(
-        experiences,
-        ensure_ascii=False,
-        indent=2,
-    )
-
-    demonstrations = _format_reasoning_examples(
-        query,
-        top_k=8,
-    )
-
-    messages = [
-        {
-            "role": "system",
-            "content": THINK_SYSTEM,
-        },
-        {
-            "role": "user",
-            "content": (
-                f"USER QUESTION:\n"
-                f"{query}\n\n"
-                f"CONTEXT / EVIDENCE:\n"
-                f"{context[:20000]}\n\n"
-                f"CONVERSATION HISTORY:\n"
-                f"{json.dumps(history[-8:], ensure_ascii=False)}\n\n"
-                f"RELEVANT PRIOR EXPERIENCES:\n"
-                f"{experience_text}\n\n"
-                f"RELEVANT REASONING EXAMPLES:\n"
-                f"{demonstrations}"
-            ),
-        },
-    ]
-
-    raw = generate(
-        messages,
-        max_new_tokens=THINK_MAX_NEW_TOKENS,
-        temperature=THINK_TEMPERATURE,
-    )
-
-    # If the model voluntarily used <think>, expose only the contents.
-    explicit_think = _extract_think(raw)
-
-    if explicit_think:
-        return explicit_think
-
-    # Otherwise the entire thinking-pass output is the reasoning workspace.
-    return _remove_think(raw)
-
-
-# =====================================================================
-# FINAL ANSWER PASS
-# =====================================================================
-
+VERIFIER_SYSTEM = r"""
+You are the verification component of Infinity GRM.
+Audit the proposed analysis. Check correctness, unsupported claims,
+calculations, edge cases, source consistency, and whether the conclusion
+actually follows.
+Return:
+<think>
+1. Claim-by-claim checks
+2. Error/edge-case checks
+3. Independent sanity check
+4. Corrections required
+5. Final verification status
+</think>
+Then provide a concise verification report.
+"""
 
 FINAL_SYSTEM = r"""
-You are the final answer component of Infinity GRM.
-
-Answer the user's original question directly.
-
-A separate reasoning model has already explored the problem. Its output is
-provided below as INTERNAL REASONING.
-
-Use that reasoning to improve accuracy, but DO NOT copy it verbatim.
-
-DO NOT reproduce the reasoning process.
-
-DO NOT describe the reasoning process unless the user explicitly asks for
-an explanation.
-
-DO NOT say "the reasoning above", "my thought process", or similar phrases.
-
-DO NOT include <think> tags.
-
-DO NOT turn the internal reasoning into a numbered chain of thought.
-
-Instead, produce a clean, useful answer appropriate to the user's question.
-
-You may include concise explanations, calculations, code, examples, or
-justification when they are useful to answer the question.
-
-If the internal reasoning contains an error, correct it rather than copying
-it.
-
-If information is missing or uncertain, say so rather than inventing it.
-
-The final response should stand on its own and should not require the user
-to read the internal reasoning.
+You are the final synthesis component of Infinity GRM.
+Answer the user directly using the plan, analysis, verification, and evidence.
+Do not invent facts. If uncertainty remains, state it precisely.
+Return:
+<think>
+1. What was established
+2. How the verification affected the result
+3. What answer follows
+4. Any material caveat
+</think>
+Then give the final answer in a useful format.
 """
 
 
-def final_from_thinking(
-    query: str,
-    thinking: str,
-    history: List[Dict[str, str]],
-    evidence: str = "",
-) -> str:
-    """
-    Generate the final answer in a separate model pass.
-
-    This separation prevents the thinking pass from simply becoming the
-    user's final answer.
-    """
+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": FINAL_SYSTEM,
-        },
+        {"role": "system", "content": PLANNER_SYSTEM},
         {
             "role": "user",
             "content": (
-                f"ORIGINAL USER QUESTION:\n"
-                f"{query}\n\n"
-                f"INTERNAL REASONING:\n"
-                f"{thinking[:30000]}\n\n"
-                f"EVIDENCE / CONTEXT:\n"
-                f"{evidence[:20000]}\n\n"
-                f"CONVERSATION HISTORY:\n"
-                f"{json.dumps(history[-8:], ensure_ascii=False)}\n\n"
-                f"Now produce only the clean final answer."
+                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=FINAL_MAX_NEW_TOKENS,
-        temperature=FINAL_TEMPERATURE,
-    )
-
-    # The final pass should not contain reasoning tags.
-    final = _remove_think(raw)
-
-    return final.strip()
-
-
-# =====================================================================
-# BACKWARD-COMPATIBLE API
-# =====================================================================
-
+    raw = generate(messages, max_new_tokens=1400, temperature=0.10)
+    think = _extract_think(raw)
+    json_match = re.search(r"\{.*\}", raw, flags=re.DOTALL)
 
-def plan(
-    query: str,
-    history: List[Dict[str, str]],
-) -> Dict[str, Any]:
-    """
-    Backward-compatible plan() API.
-
-    The old implementation generated a rigid numbered planning trace.
-    This version instead uses the free-form reasoning pass.
-
-    The returned structure is preserved so existing callers do not break.
-    """
-
-    initialize_model()
-
-    experiences = retrieve_experiences(
-        query,
-        top_k=3,
-    )
-
-    reasoning = think(
-        query=query,
-        history=history,
-        context="",
-    )
+    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": reasoning,
+        "think": think,
         "task": query,
         "requires_web": False,
         "search_queries": [],
-        "steps": [],
-        "verification": [],
+        "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:
-    """
-    Backward-compatible analyze() API.
-
-    Rather than forcing another predefined analysis stage, this performs
-    another free-form reasoning pass using the available evidence.
-    """
-
-    existing_reasoning = plan_data.get(
-        "think",
-        "",
-    )
-
+def analyze(query: str, plan_data: Dict[str, Any], evidence: str, history: List[Dict[str, str]]) -> str:
     messages = [
-        {
-            "role": "system",
-            "content": THINK_SYSTEM,
-        },
+        {"role": "system", "content": ANALYST_SYSTEM},
         {
             "role": "user",
             "content": (
-                f"USER QUESTION:\n"
-                f"{query}\n\n"
-                f"EXISTING REASONING:\n"
-                f"{existing_reasoning[:24000]}\n\n"
-                f"EVIDENCE:\n"
-                f"{evidence[:24000]}\n\n"
-                f"HISTORY:\n"
-                f"{json.dumps(history[-8:], ensure_ascii=False)}\n\n"
-                "Continue reasoning freely. Reconsider the existing "
-                "reasoning if necessary. Correct errors, explore better "
-                "approaches, and determine what conclusion the evidence "
-                "actually supports. Do not produce a polished final answer."
+                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)
 
-    return generate(
-        messages,
-        max_new_tokens=THINK_MAX_NEW_TOKENS,
-        temperature=THINK_TEMPERATURE,
-    )
-
-
-def verify(
-    query: str,
-    plan_data: Dict[str, Any],
-    analysis: str,
-    evidence: str,
-) -> str:
-    """
-    Backward-compatible verify() API.
-
-    Verification is still available, but it is no longer presented to the
-    model as a fixed numbered checklist.
-    """
 
+def verify(query: str, plan_data: Dict[str, Any], analysis: str, evidence: str) -> str:
+    """Independent verification pass."""
     messages = [
-        {
-            "role": "system",
-            "content": r"""
-You are a reasoning verification engine.
-
-Inspect the proposed reasoning and determine whether it is actually correct.
-
-Do not follow a fixed checklist.
-
-Reason naturally about whatever could be wrong.
-
-You may:
-- recalculate something;
-- test an assumption;
-- find a counterexample;
-- check an edge case;
-- compare an alternative interpretation;
-- identify unsupported claims;
-- discover a contradiction;
-- confirm that the conclusion follows.
-
-If the reasoning is correct, explain briefly why it survives verification.
-
-If it is incorrect, identify the actual problem and work out the correction.
-
-Do not produce a polished final answer for the user.
-""",
-        },
+        {"role": "system", "content": VERIFIER_SYSTEM},
         {
             "role": "user",
             "content": (
-                f"QUESTION:\n"
-                f"{query}\n\n"
-                f"REASONING:\n"
-                f"{analysis[:26000]}\n\n"
-                f"EVIDENCE:\n"
-                f"{evidence[:24000]}"
+                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=THINK_MAX_NEW_TOKENS,
-        temperature=0.05,
-    )
+    return generate(messages, max_new_tokens=1400, temperature=0.05)
 
 
 def final_answer(
@@ -1302,103 +775,27 @@ def final_answer(
     history: List[Dict[str, str]],
     verification: Optional[str] = None,
 ) -> str:
+    """Final reasoning/synthesis pass.
+    `verification` is optional to preserve compatibility with the previous
+    five-argument search pipeline.
     """
-    Backward-compatible final_answer() API.
-
-    The final response is always generated separately from the reasoning.
-    """
-
-    verification_text = (
-        verification
-        or "No separate verification pass was supplied."
-    )
-
-    thinking_material = (
-        f"PRIMARY REASONING:\n"
-        f"{analysis[:22000]}\n\n"
-        f"VERIFICATION:\n"
-        f"{verification_text[:18000]}"
-    )
-
-    return final_from_thinking(
-        query=query,
-        thinking=thinking_material,
-        history=history,
-        evidence=evidence,
-    )
-
-
-# =====================================================================
-# FULL REASONING PIPELINE
-# =====================================================================
-
-
-def reason(
-    query: str,
-    history: Optional[List[Dict[str, str]]] = None,
-    evidence: str = "",
-) -> Dict[str, Any]:
-    """
-    Main reasoning pipeline.
-
-    The important architecture is:
-
-        question
-            ↓
-        free-form thinking
-            ↓
-        separate final-answer generation
-            ↓
-        clean answer
-
-    No mandatory numbered reasoning procedure is imposed.
-    """
-
-    if history is None:
-        history = []
-
-    initialize_model()
-
-    start = time.time()
-
-    thinking = think(
-        query=query,
-        history=history,
-        context=evidence,
-    )
-
-    final = final_from_thinking(
-        query=query,
-        thinking=thinking,
-        history=history,
-        evidence=evidence,
-    )
-
-    elapsed = time.time() - start
-
-    result = {
-        "query": query,
-        "thinking": thinking,
-        "answer": final,
-        "elapsed_seconds": elapsed,
-    }
-
-    if SHOW_REASONING_TRACE:
-        result["response"] = (
-            f"<think>\n"
-            f"{thinking}\n"
-            f"</think>\n\n"
-            f"{final}"
-        )
-    else:
-        result["response"] = final
-
-    return result
-
-
-# =====================================================================
-# TRACE
-# =====================================================================
+    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(
@@ -1408,48 +805,23 @@ def build_reasoning_trace(
     verification: str,
     final: str,
 ) -> Dict[str, Any]:
-    """
-    Return a detailed visible trace of the reasoning pipeline.
-
-    Kept for compatibility with existing callers.
-    """
-
+    """Return a detailed visible trace of the reasoning pipeline."""
     return {
         "query": query,
-        "demonstration_count": len(
-            REASONING_EXAMPLES
-        ),
+        "demonstration_count": len(REASONING_EXAMPLES),
         "stages": [
             {
-                "stage": "free_form_reasoning",
-                "reasoning_summary": plan_data.get(
-                    "think",
-                    "",
-                ),
+                "stage": "planning",
+                "reasoning_summary": plan_data.get("think", ""),
                 "structured_plan": plan_data,
             },
-            _make_trace(
-                "continued_reasoning",
-                analysis,
-            ),
-            _make_trace(
-                "verification",
-                verification,
-            ),
-            {
-                "stage": "final_synthesis",
-                "reasoning_summary": "",
-                "output": _remove_think(final),
-            },
+            _make_trace("analysis", analysis),
+            _make_trace("verification", verification),
+            _make_trace("final_synthesis", final),
         ],
     }
 
 
-# =====================================================================
-# LEARNING
-# =====================================================================
-
-
 def learn_from_result(
     query: str,
     plan_data: Dict[str, Any],
@@ -1458,51 +830,20 @@ def learn_from_result(
     quality: float = 1.0,
 ):
     initialize_model()
-
     save_experience(
         query=query,
-        plan=json.dumps(
-            plan_data,
-            ensure_ascii=False,
-        ),
+        plan=json.dumps(plan_data, ensure_ascii=False),
         analysis=analysis,
         answer=answer,
         quality=quality,
     )
 
 
-# =====================================================================
-# INITIALIZATION
-# =====================================================================
-
-
 def initialize():
     initialize_model()
 
 
-# =====================================================================
-# MAIN
-# =====================================================================
-
-
 if __name__ == "__main__":
     initialize()
-
-    print(
-        "Infinity GRM reasoning engine initialized."
-    )
-
-    print(
-        f"Structured reasoning demonstrations: "
-        f"{len(REASONING_EXAMPLES)}"
-    )
-
-    print(
-        f"Free-form thinking budget: "
-        f"{THINK_MAX_NEW_TOKENS} tokens"
-    )
-
-    print(
-        f"Final answer budget: "
-        f"{FINAL_MAX_NEW_TOKENS} tokens"
-    )
+    print("Infinity GRM reasoning engine initialized.")
+    print(f"Structured reasoning demonstrations: {len(REASONING_EXAMPLES)}")