run_tiny.py

"""NEXUS-R1D harness — real-tiny live runner.

Runs a tiny open model (default Qwen3-0.6B; `--model` accepts any HF id)
through the NEXUS-R1D reasoning→tool→verify→answer loop. The NEXUS
controllers (budget/rounds/repetition, tool schema validation, expected-vs-
actual completion) are the real code from nexus_r1d/; the 26B backbone is
substituted by the tiny model so the loop can actually run on CPU.

Usage:
  python3 run_tiny.py                          # interactive chat
  python3 run_tiny.py --task demos/tasks.json   # run task file
  python3 run_tiny.py --model Qwen/Qwen3-0.6B   # other tiny model
"""
from __future__ import annotations

import argparse
import json
import math
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from nexus_r1d.config import NexusR1DConfig
from nexus_r1d.reasoning import ReasoningController, ReasoningState
from nexus_r1d.tools import Tool, ToolCall, ToolInterface, ToolPolicy


# --------------------------------------------------------------------------- #
# Built-in tools (safe, local, deterministic)                                   #
# --------------------------------------------------------------------------- #
def _register_tools(ti: ToolInterface) -> None:
    ti.register(Tool(
        "calculator", "evaluate arithmetic expressions like (2+3)*7",
        {"type": "object", "properties": {"expression": {"type": "string"}},
         "required": ["expression"]},
        handler=lambda expression: eval(expression, {"__builtins__": {}}, {})
        if re.fullmatch(r"[0-9+\-*/(). ]+", expression) else (_ for _ in ()).throw(
            ValueError("only digits and + - * / ( ) . allowed")),
        read_only=True,
    ))
    ti.register(Tool(
        "read_file", "read a local text file (workspace sandbox)",
        {"type": "object", "properties": {"path": {"type": "string"}},
         "required": ["path"]},
        handler=lambda path: Path(path).read_text()[:4000] if Path(path).is_file()
        else f"ERROR: file not found: {path}",
        read_only=True,
    ))
    ti.register(Tool(
        "glob_files", "list files matching a pattern under the workspace",
        {"type": "object", "properties": {"pattern": {"type": "string"}},
         "required": ["pattern"]},
        handler=lambda pattern: json.dumps(
            sorted(str(p) for p in Path(".").glob(pattern))[:50]),
        read_only=True,
    ))
    ti.register(Tool(
        "grep_files", "regex search in workspace files",
        {"type": "object", "properties": {"pattern": {"type": "string"}},
         "required": ["pattern"]},
        handler=lambda pattern: "\n".join(
            f"{p}:{i}:{l.strip()[:120]}"
            for p in Path(".").rglob("*") if p.is_file() and p.suffix in {".py", ".md", ".json", ".txt"}
            for i, l in enumerate(p.read_text(errors="ignore").splitlines(), 1)
            if re.search(pattern, l))[:2000] or "no matches",
        read_only=True,
    ))


# --------------------------------------------------------------------------- #
# Generation loop with NEXUS reasoning budgets                                 #
# --------------------------------------------------------------------------- #
SYSTEM_PROMPT = (
    "You are NEXUS-R1D (tiny probe build). Answer concisely.\n"
    "TOOLS: {schemas}\n"
    "HOW TO CALL A TOOL: end your reply with exactly one flat JSON object. "
    "Example — if the user asks 'what is 2+2?', your reply ends with:\n"
    '{"name": "calculator", "arguments": {"expression": "2+2"}}\n'
    "The JSON must be flat: 'name' is the tool name string, 'arguments' holds the "
    "tool's parameters directly. Never nest JSON inside arguments.\n"
    "ONLY call a tool when truly needed: simple facts (capitals, definitions, "
    "greetings) get NO tool — answer directly. If a tool failed, do not repeat it."
)


def _split_think(text: str) -> tuple[str, str]:
    """R1-style channel separation: return (thinking, answer).

    Qwen3 thinking mode emits reasoning first, then closes and emits the
    final answer -- structurally identical to R1's discipline: reasoning is
    generated BEFORE the answer and kept in a separate channel.
    """
    CLOSE = "</" + "think" + ">"
    OPEN = "<" + "think" + ">"
    if CLOSE in text:
        thinking, answer = text.split(CLOSE, 1)
        return thinking.replace(OPEN, "").strip(), answer.strip()
    return text.replace(OPEN, "").strip(), ""


def build_system_prompt(schemas: str) -> str:
    """Escape braces so .format-style substitution is safe with JSON examples."""
    return SYSTEM_PROMPT.replace("{schemas}", schemas)

NAME_ARGS_RE = re.compile(r'\{"name"\s*:\s*"[^"]+"\s*,\s*"arguments"\s*:\s*\{.*?\}\s*\}', re.S)


class TinyProbe:
    def __init__(self, model_id: str, budget: str = "medium"):
        print(f"Loading {model_id} (CPU, bf16)…")
        self.tok = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id, dtype=torch.bfloat16, device_map="cpu")
        self.model.eval()
        self.cfg = NexusR1DConfig()
        self.rc = ReasoningController(self.cfg)
        self.ti = ToolInterface(self.cfg, ToolPolicy(max_retries=2))
        _register_tools(self.ti)
        self.budget = budget
        self.think_enabled = True
        self.rc_max_think = self.rc.initial_budget(budget) // 4  # words ~ tokens/1.3
        self.show_thinking = False   # spec 10: latent by default, toggle to show
        self._failed_calls: set[str] = set()

    def generate(self, prompt: str, max_new_tokens: int) -> str:
        msgs = [{"role": "system", "content": build_system_prompt(
            json.dumps([t["name"] + ": " + t["description"] for t in self.ti.schemas()]))},
            {"role": "user", "content": prompt}]
        enc = self.tok.apply_chat_template(
            msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True,
            enable_thinking=self.think_enabled)
        ids = enc["input_ids"]
        with torch.no_grad():
            out = self.model.generate(
                **enc, max_new_tokens=max_new_tokens, do_sample=False,
                pad_token_id=self.tok.eos_token_id)
        text = self.tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True)
        if not self.think_enabled:
            return "", text
        thinking, answer = _split_think(text)
        # NEXUS ReasoningController: enforce the token budget on the thinking
        # channel (R1-style adaptive compute, spec 11 / docs/02).
        if self.rc_max_think and len(thinking.split()) > self.rc_max_think:
            thinking = "[budget-cut] " + " ".join(thinking.split()[:self.rc_max_think]) + " ..."
        return thinking, answer

    def ask(self, question: str) -> dict:
        state = ReasoningState(token_budget=self.rc.initial_budget(self.budget) // 4)
        history = question
        trace = []
        for round_idx in range(3):
            thinking, text = self.generate(history, max_new_tokens=896)
            # NEXUS ReasoningController: budget/round/repetition guardrails
            # (should_stop appends to the transcript and advances the round).
            stop, reason = self.rc.should_stop(state, text)
            trace.append({"round": round_idx, "thinking": thinking[:400],
                          "output": text[:400], "stop": reason or None})
            # scan BOTH channels (answer first, then thinking) for the LAST
            # tool-call JSON; NEXUS treats it as an action either way
            m = None
            for chan in (text, thinking):
                for m_ in NAME_ARGS_RE.finditer(chan):
                    m = m_
            if m is None:
                return {"answer": (text or "(model still thinking at token cap)")[:300],
                        "thinking": thinking, "tool_used": None,
                        "rounds": state.round_idx, "trace": trace}
            try:
                obj = json.loads(m.group(0))
                call = ToolCall(name=obj["name"], arguments=obj.get("arguments", {}))
            except json.JSONDecodeError:
                continue
            # §25: no repeated identical failed calls (error-recovery budget)
            key = call.name + ":" + json.dumps(call.arguments, sort_keys=True)
            if key in self._failed_calls:
                trace.append({"tool": call.name, "status": "blocked-repeat"})
                # force a final no-tool answer instead of looping
                history = (
                    f"{question}\n\nThe tool call keeps failing. Do NOT call any tool. "
                    "Give your final plain-text answer now, explaining what happened."
                )
                thinking2, text = self.generate(history, max_new_tokens=384)
                trace.append({"round": round_idx, "output": text[:300]})
                return {"answer": (text or "(no final answer produced)")[:300],
                        "thinking": (thinking + "\n--forced--\n" + thinking2)[:600],
                        "tool_used": call.name,
                        "rounds": state.round_idx, "trace": trace}
            res = self.ti.execute(call)
            trace.append({"tool": f"{res.name}({json.dumps(call.arguments)})",
                          "status": res.status,
                          "result": str(res.result)[:300] if res.result else res.error})
            if res.status == "success":
                if stop:
                    # Budget exhausted but a tool result is in hand: force a
                    # immediate synthesis round instead of another full loop.
                    history = (
                        f"{question}\n\nYou called tool `{res.name}` and got result:\n{res.result}\n\n"
                        "Your reasoning budget is exhausted. Give the final answer NOW. "
                        "Plain text only, NO tool call."
                    )
                    thinking2, text = self.generate(history, max_new_tokens=256)
                    trace.append({"round": round_idx, "output": text[:300]})
                    return {"answer": (text or "(no final answer produced)")[:300],
                            "thinking": (thinking + "\n--budget-stop--\n" + thinking2)[:600],
                            "tool_used": res.name,
                            "rounds": state.round_idx, "trace": trace}
                history = (
                    f"{question}\n\nYou called tool `{res.name}` and got result:\n{res.result}\n\n"
                    "Now give the final answer to the user. Plain text only, NO tool call."
                )
            else:
                self._failed_calls.add(call.name + ":" + json.dumps(call.arguments, sort_keys=True))
                history = (
                    f"{question}\n\nYour tool call failed: {res.error}\n"
                    "Do NOT repeat this same call. Answer without tools, or fix the arguments."
                )
        return {"answer": (trace[-1].get("output", "") or "max rounds reached")[:300],
                "thinking": trace[-1].get("thinking", ""),
                "tool_used": "max rounds", "rounds": state.round_idx, "trace": trace}


# --------------------------------------------------------------------------- #
# Entry                                                                        #
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen3-0.6B")
    ap.add_argument("--budget", default="medium", choices=["low", "medium", "high"])
    ap.add_argument("--task", default=None, help="path to task JSON file")
    ap.add_argument("--show-thinking", action="store_true",
                    help="print the R1-style thinking channel (default: hidden, spec 10)")
    ap.add_argument("--no-thinking", action="store_true",
                    help="disable the thinking channel entirely (fast mode)")
    args = ap.parse_args()

    probe = TinyProbe(args.model, args.budget)
    if args.no_thinking:
        probe.think_enabled = False
    probe.show_thinking = args.show_thinking

    if args.task:
        tasks = json.loads(Path(args.task).read_text())
        results = [probe.ask(t["question"]) for t in tasks]
        print("\n=== RESULTS ===")
        for t, r in zip(tasks, results):
            print(f"\n[task {t.get('id', '?')}] {t['question'][:60]}…")
            print(f"  answer: {str(r['answer'])[:200]}")
            if probe.show_thinking and r.get("thinking"):
                print(f"  thinking: {r['thinking'][:200]}")
            for s in r["trace"]:
                if "tool" in s:
                    print(f"  tool: {s['tool']} -> {s['status']}")
        return

    print("\nNEXUS-R1D tiny probe — interactive (empty line to quit)")
    while True:
        try:
            q = input("\nyou> ").strip()
        except EOFError:
            break
        if not q:
            break
        r = probe.ask(q)
        print(f"\nrounds={r['rounds']} | tool={r['tool_used']}")
        if probe.show_thinking and r.get("thinking"):
            print(f"think> {r['thinking'][:400]}")
        print(f"nexus> {r['answer']}")


if __name__ == "__main__":
    main()