models/minicpm_agent.py

"""MiniCPM5-1B: the text "brain" that drives the find-and-point loop.

Where models/minicpm.py is the MiniCPM-V VLM — the "eyes" that ground the circle
and write the ingest figure/table descriptions — this is the small TEXT model
that decides what to do. Each step it sees the conversation so far, the manual's
table of contents, and the WHOLE text of the page being viewed, and picks ONE
tool:

    search(query)            semantic-search the manual; the best page is shown
                             and its text added to the conversation
    find_answer(query)       dense TEXT search of the parsed chunks for a
                             spec/value/fact; the page that states it is shown,
                             its text added so the agent can circle the answer
    go_to_section(section)   jump to a numbered table-of-contents section
    circle(target)           circle something on the CURRENT page (its text is
                             in context); the VLM grounds the box
    done(message)            nothing more to do, or it can't be found

It never writes answers — the manual does the talking. History is used only to
resolve references ("circle the other one", "go back to that bolt").

The default brain is a standard LlamaForCausalLM (no trust_remote_code); other
selectable brains may ship custom modeling code (loaded with trust_remote_code
per their AGENT_MODELS spec — see use_model). Loaded as a module-level CUDA
global for the same ZeroGPU reason as the other models: module-level CUDA tensors
are shared with the GPU worker, whereas function arguments are pickled.

Entry points (plain functions; the caller runs them inside its @spaces.GPU
context). decide() takes the running message list the pipeline maintains and
returns a parsed tool call; rerank() picks the best of N candidate pages by their
text — the search tool's reranker over ColEmbed's shortlist. The system prompt
and the message builders live here so the wording stays with the model.
"""

from __future__ import annotations

import gc
import json
import logging
import re

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from core import tracing
from core.constants import (
    AGENT_MAX_NEW_TOKENS,
    AGENT_MODELS,
    DEFAULT_AGENT_MODEL,
)
from core.vram import log_vram

log = logging.getLogger("repairguy.agent")

# The tools the agent may emit, and the JSON shape of each. Kept here so the
# prompt and the parser can't drift apart.
TOOLS = ("search", "find_answer", "go_to_section", "go_to_page", "circle", "done")

SYSTEM_PROMPT = (
    "You are the assistant for a hands-busy mechanic reading a repair manual on "
    "a page viewer. You do NOT answer questions or explain — the manual does the "
    "talking. Your only job is to FIND the right page and POINT at things on it.\n\n"
    "Each step, choose exactly ONE tool and reply with ONLY its JSON object — no "
    "prose, no markdown, nothing else. Replace every <...> placeholder with the "
    "real value — NEVER output the angle brackets.\n\n"
    "Tools:\n"
    '- Search the manual for a part/topic/procedure:\n'
    '  {"tool": "search", "query": "<focused search phrase>"}\n'
    "- Look up the ANSWER to a question — a spec, value, or fact (fuel type, a "
    "torque, an oil/coolant capacity). This searches the manual TEXT, so it "
    "finds the page that STATES the answer when a topic search would miss the "
    "plain specs page:\n"
    '  {"tool": "find_answer", "query": "<the spec or fact being asked for>"}\n'
    "- Jump to a section — use its number from the TABLE OF CONTENTS:\n"
    '  {"tool": "go_to_section", "section": <number>}\n'
    "- Jump straight to a known PHYSICAL page number:\n"
    '  {"tool": "go_to_page", "page": <number>}\n'
    "- Circle something on a page CURRENTLY ON SCREEN (their full text is given "
    "to you) — a part, OR the line/value that answers their question. When two "
    "pages are shown, set page to the one whose text has it:\n"
    '  {"tool": "circle", "target": "<short name or printed words to circle>", '
    '"page": <the on-screen page number it is on>}\n'
    "- Finish — nothing more to do, or it isn't in the manual:\n"
    '  {"tool": "done", "message": "<one short line for the mechanic>"}\n\n'
    "How to choose:\n"
    '- They say "go to" / "take me to" / name a section → go_to_section.\n'
    "- You already KNOW the page number — e.g. the CURRENT page is an index or "
    'contents that lists the part with a page number ("Actuators .... 855"), or '
    "history gave one → go_to_page that number. Do NOT circle the index line; "
    "go to the page it points to.\n"
    "- The thing they want is on a page ON SCREEN → circle it, and set page to "
    "the one it is on. The target MUST be what the mechanic asked for — match it "
    "to that page's wording if it appears there (a part named inside a figure or "
    "diagram description still counts as on the page); NEVER circle a different "
    "component.\n"
    "- They ask a QUESTION whose answer is printed on a page on screen — a spec, "
    'a value, or a line (e.g. "what fuel does it take" answered by "Engine fuel '
    '- Gasoline") → circle that spot, using the printed words as the target. '
    "Point at the answer; do NOT keep searching for a better page.\n"
    "- They ask a QUESTION for a spec/value/fact that is NOT on screen yet "
    '("what fuel does it take", "engine oil capacity", "drain-plug torque") → '
    "find_answer with the fact as the query; it searches the manual text and "
    "shows the page that states it, which you then circle. Use find_answer for "
    "fact questions; use search to find a part, diagram, or procedure.\n"
    "- Otherwise → search. After a search shows a page, that page is on screen, "
    "so circle on it or search again — but NEVER repeat a search that did not "
    "move you; act on the page instead.\n"
    "- Use the conversation history only to resolve what they mean (e.g. "
    '"circle the other one"); never restate earlier answers.\n\n'
    "Examples (copy the FORMAT, not the values):\n"
    'Mechanic: "go to the cooling system" → {"tool": "go_to_section", "section": 5}\n'
    'Mechanic: "where do I replace the fuel filter" (not on this page) → '
    '{"tool": "search", "query": "fuel filter replacement"}\n'
    'Mechanic: "the wastegate actuator" (current page is an index reading '
    '"Actuators .... 855") → {"tool": "go_to_page", "page": 855}\n'
    'Mechanic: "circle the bleeder screw" (it is on p.412, which is on screen) → '
    '{"tool": "circle", "target": "bleeder screw", "page": 412}\n'
    'Mechanic: "what fuel does it take" (p.5 on screen shows "Engine fuel - '
    'Gasoline") → {"tool": "circle", "target": "Engine fuel - Gasoline", "page": 5}\n'
    'Mechanic: "what\'s the engine oil capacity" (not on this page) → '
    '{"tool": "find_answer", "query": "engine oil capacity"}'
)


def system_message() -> dict:
    return {"role": "system", "content": SYSTEM_PROMPT}


def state_message(
    request: str,
    toc: list[dict],
    shown: list[dict],
    section: str,
) -> dict:
    """The user message for the current step: what the mechanic just said, the
    table of contents (numbered, the go_to_section index), and the whole text of
    the page(s) currently on the viewer. The viewer shows a two-page spread, so
    `shown` is [{page, text}] for each page on screen (one or two). Each page's
    text is the parsed page rendered to text (figures/tables as descriptions) —
    empty when the manual has no parse. When the agent circles, it names which of
    these pages the target is on."""
    toc_lines = "\n".join(
        f"{i + 1}. {s['title']} (p.{s['page']})" for i, s in enumerate(toc)
    ) or "(none)"
    if shown:
        where = " and ".join(f"p.{s['page']}" for s in shown) + (
            f' (section "{section}")' if section else ""
        )
        blocks = "\n\n".join(
            f"PAGE {s['page']} — full text:\n{s['text'] or '(no text available)'}"
            for s in shown
        )
        page_block = (
            f"CURRENTLY ON SCREEN — {where}. You can circle on "
            + (
                "either of these pages (say which in the circle call):\n"
                if len(shown) > 1
                else "this page:\n"
            )
            + blocks
        )
    else:
        page_block = "CURRENTLY ON SCREEN: (no page open)"
    # The request goes LAST (after the long page text) so it stays freshest —
    # otherwise the page block dominates and the agent acts on the page instead
    # of what was asked.
    return {
        "role": "user",
        "content": (
            f"{page_block}\n\n"
            f"TABLE OF CONTENTS:\n{toc_lines}\n\n"
            f"The mechanic said: {request!r}\n"
            "Choose ONE tool and reply with ONLY its JSON object."
        ),
    }


def tool_result_message(text: str) -> dict:
    """A tool's outcome fed back into the loop as the next observation (e.g. the
    page a search landed on, with its text, so the agent can circle on it)."""
    return {"role": "user", "content": text}


def search_result_message(request: str, page: int, text: str, stuck: bool) -> str:
    """The observation fed back after a search shows a page. `stuck` is set by the
    pipeline when the search was a no-op — the same query again, or it landed on a
    page already shown this turn. Then the message FORCES a decision and DROPS the
    "search again" option, because a 1B on greedy decoding will otherwise re-issue
    the same search and loop on the same page forever (seen live: 5x the identical
    search, all landing on p.876, never circling)."""
    body = (
        f"Search showed p.{page}. It is now the CURRENT page.\n"
        f"CURRENT PAGE (p.{page}) — full text:\n"
        f"{text or '(no text available)'}\n\n"
    )
    if stuck:
        return body + (
            f"You are still on p.{page}; that search did not move you, and "
            "repeating it will NOT help. Decide now — do not search again:\n"
            f"- If this page shows {request!r} OR the line/value that answers it "
            "(it counts even when named inside a figure or diagram description), "
            'circle that spot: {"tool": "circle", "target": "<the printed words>", '
            f'"page": {page}}}.\n'
            "- If it belongs on a different page, use go_to_section or go_to_page.\n"
            "- Only if it is truly not in this manual, use done."
        )
    return body + (
        f"The mechanic asked for: {request!r}. If this page shows THAT — or the "
        "line/value that answers it, including when named inside a figure or "
        "diagram description — circle that spot (use the words as printed on the "
        "page). Otherwise search again or go to a section. Do not circle a "
        "different component."
    )


def ground_failed_message(request: str, target: str, page: int) -> str:
    """The observation fed back when the VLM could not locate `target` on p.`page`.
    A failed grounding almost always means the target is NOT on this page (the
    agent circled on the wrong one), so this FORCES a relocate — search or
    navigate — and forbids re-circling the same spot, the same greedy-loop guard
    used after a no-op search. Without it, circle is a dead-end: the page shows
    with no pin and the turn ends."""
    return (
        f"I could not find {target!r} on p.{page} — it does not appear to be on "
        "this page, so circling here will not work. Do NOT circle that on this "
        "page again. It is almost certainly on a DIFFERENT page: search for "
        f"{request!r}, or go to the right section or page, and only circle once "
        "you are on the page whose text actually shows it. Use done only if it is "
        "truly not in this manual."
    )


def assistant_action_message(tool: dict) -> dict:
    """Record an action the agent took, so it stays in the running transcript
    (this turn) and the compact history (later turns)."""
    return {"role": "assistant", "content": json.dumps(tool, separators=(",", ":"))}


# --- the resident brain: ONE model in VRAM, swapped on demand ---------------
# The agent model is selectable from the UI (core.constants.AGENT_MODELS). Only
# one is held on the GPU at a time: use_model() evicts the current one before
# loading the next ("load-on-switch"), so VRAM stays flat as the user A/Bs
# models. The model/tokenizer live in module globals for the same ZeroGPU reason
# as the other models — module-level CUDA tensors are shared with the GPU worker.
_REGISTRY = {m["key"]: m for m in AGENT_MODELS}
_active_key: str | None = None
_MODEL = None
_TOKENIZER = None
# Whether the active model's chat template accepts enable_thinking (Qwen3 /
# all current brains do) — drives whether we pass the kwarg below.
_THINKING = False


def _spec(key: str | None) -> dict:
    return _REGISTRY.get(key or "", _REGISTRY[DEFAULT_AGENT_MODEL])


def _active_model_id() -> str:
    """The HF id of the resident brain — the `model` for its generation spans."""
    return _spec(_active_key).get("model_id", _active_key or "unknown")


def use_model(key: str | None = None) -> str:
    """Make `key` the resident agent brain, loading it (and evicting the
    previous one) when it isn't already active — one model in VRAM at a time.
    Unknown keys fall back to the default. Called once per turn by the pipeline,
    inside the GPU context. Returns the active key. Must run on GPU."""
    global _active_key, _MODEL, _TOKENIZER, _THINKING
    spec = _spec(key)
    if spec["key"] == _active_key and _MODEL is not None:
        return _active_key
    # Drop the current model first so VRAM holds only one brain at a time.
    if _MODEL is not None:
        log.info("agent model: evicting %s", _active_key)
        log_vram(f"before-evict-{_active_key}")
        _MODEL = _TOKENIZER = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        log_vram(f"after-evict-{_active_key}")
    log.info("agent model: loading %s (%s)", spec["key"], spec["model_id"])
    trust = spec.get("trust_remote_code", False)
    _TOKENIZER = AutoTokenizer.from_pretrained(
        spec["model_id"], revision=spec["revision"], trust_remote_code=trust
    )
    load_kwargs = dict(
        revision=spec["revision"],
        dtype=torch.bfloat16,
        trust_remote_code=trust,
    )
    # attn_implementation defaults to "sdpa" but a spec can override it — None
    # omits the kwarg entirely so the model's own modeling code picks its default
    # (MiniCPM4.1-8B ships custom sparse attention that we don't want to force).
    attn = spec.get("attn_implementation", "sdpa")
    if attn is not None:
        load_kwargs["attn_implementation"] = attn
    # A spec can ask accelerate to place weights straight on the GPU
    # (device_map) instead of the default load-on-CPU-then-.to("cuda"). Big
    # brains switched in at runtime want this: the bulk host→device copy of a
    # multi-GB model inside the ephemeral @spaces.GPU fork trips an NVML assert
    # in the CUDA caching allocator. device_map loads onto CUDA directly, so we
    # skip the trailing .to("cuda").
    device_map = spec.get("device_map")
    if device_map is not None:
        load_kwargs["device_map"] = device_map
    model = AutoModelForCausalLM.from_pretrained(spec["model_id"], **load_kwargs)
    if device_map is None:
        model = model.to("cuda")
    _MODEL = model.eval()
    _active_key, _THINKING = spec["key"], spec["thinking"]
    log_vram(f"after-load-{_active_key}")
    return _active_key


# Load the default brain eagerly at import — same as the other models, so the
# ZeroGPU startup packing covers it and the common (no-switch) path pays no
# load cost on the first turn.
use_model(DEFAULT_AGENT_MODEL)


def _template_kwargs() -> dict:
    """Extra apply_chat_template kwargs for the active model. Tool routing wants
    a terse decision, not a reasoning trace, so disable thinking — but only for
    templates that accept the kwarg (others would ignore or choke on it)."""
    return {"enable_thinking": False} if _THINKING else {}


def _generate(
    messages: list[dict], max_new_tokens: int, trace_name: str = "agent-generate"
) -> str:
    """Greedy decode the assistant's next message. Traced as one `generation`
    (the resident brain as the model, the messages as input, the reply and the
    in/out token counts attached) when Langfuse is configured."""
    with tracing.generation(
        trace_name, model=_active_model_id(), input=messages
    ) as gen:
        inputs = _TOKENIZER.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
            **_template_kwargs(),
        ).to(_MODEL.device)
        # apply_chat_template emits token_type_ids, which this LlamaForCausalLM's
        # generate() rejects as an unused kwarg.
        inputs.pop("token_type_ids", None)
        n_in = int(inputs["input_ids"].shape[1])
        with torch.no_grad():
            out = _MODEL.generate(
                **inputs, max_new_tokens=max_new_tokens, do_sample=False
            )
        new_ids = out[0, n_in:]
        text = _TOKENIZER.decode(new_ids, skip_special_tokens=True)
        if gen is not None:
            n_out = int(new_ids.shape[0])
            gen.update(
                output=text.strip(),
                usage_details={"input": n_in, "output": n_out, "total": n_in + n_out},
            )
        return text.strip()


def render_prompt(messages: list[dict]) -> str:
    """The exact text fed to the model this step — the chat template applied to
    the running message list, special tokens and all. Mirrors _generate's
    template call but returns the string instead of tokenizing, so the
    Diagnostics view can show precisely what the brain was asked. CPU-only."""
    return _TOKENIZER.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        **_template_kwargs(),
    )


def _parse_tool(raw: str) -> dict | None:
    """Pull the JSON tool call out of the reply. Tolerant of a ```json fence or
    a stray lead-in. Returns a validated
    {tool, ...} or None when the reply isn't usable (the caller decides the
    fallback). Range-checking go_to_section is the caller's job — it has the
    TOC; this only validates shape."""
    text = raw.strip()
    if text.startswith("```"):
        text = text.split("\n", 1)[1] if "\n" in text else ""
        text = text.rsplit("```", 1)[0].strip()
    start = text.find("{")
    if start < 0:
        return None
    try:
        obj, _ = json.JSONDecoder().raw_decode(text[start:])
    except ValueError:
        return None
    if not isinstance(obj, dict):
        return None
    tool = obj.get("tool")
    if tool == "search":
        query = str(obj.get("query") or "").strip()
        return {"tool": "search", "query": query} if _real(query) else None
    if tool == "find_answer":
        query = str(obj.get("query") or "").strip()
        return {"tool": "find_answer", "query": query} if _real(query) else None
    if tool == "go_to_section":
        try:
            return {"tool": "go_to_section", "section": int(obj.get("section"))}
        except (TypeError, ValueError):
            return None
    if tool == "go_to_page":
        try:
            return {"tool": "go_to_page", "page": int(obj.get("page"))}
        except (TypeError, ValueError):
            return None
    if tool == "circle":
        target = str(obj.get("target") or "").strip()
        if not _real(target):
            return None
        out = {"tool": "circle", "target": target}
        # Optional: which on-screen page the target is on (the pipeline validates
        # it against the pages actually shown and defaults to the active page).
        try:
            if obj.get("page") is not None:
                out["page"] = int(obj.get("page"))
        except (TypeError, ValueError):
            pass
        return out
    if tool == "done":
        return {"tool": "done", "message": str(obj.get("message") or "").strip()}
    return None


def _real(value: str) -> bool:
    """A usable arg, not an echoed placeholder. Small models sometimes copy the
    schema example verbatim ("<short name of the thing to circle>") — angle
    brackets are the tell; reject so the loop re-asks for a real value."""
    return bool(value) and "<" not in value and ">" not in value


def decide(messages: list[dict]) -> tuple[dict | None, str]:
    """One agentic step. messages is the running conversation the pipeline
    maintains (system + past turns + this turn's state and any tool results).
    Returns (parsed tool call, raw reply); the tool is None when the reply isn't
    a usable JSON tool call. Must run on GPU."""
    raw = _generate(messages, AGENT_MAX_NEW_TOKENS, trace_name="agent-decide")
    return _parse_tool(raw), raw


RERANK_PROMPT = (
    "A mechanic is searching a repair manual for: {query!r}\n\n"
    "Below are {n} candidate pages, each its page number and the text on it. "
    "Pick the ONE page that best covers what the mechanic wants — the component, "
    "system, or procedure involved (judge by topic, not exact wording). Reply "
    "with ONLY that page's number, nothing else.\n\n{candidates}"
)


def rerank(query: str, candidates: list[tuple[int, str]]) -> tuple[int, str]:
    """Pick the best of ColEmbed's shortlist by page text. candidates is
    [(page, text)] in retrieval order. Returns (index into candidates, raw
    reply); falls back to index 0 (the top ColEmbed hit) when the reply can't be
    read. Must run on GPU."""
    if not candidates:
        return 0, ""
    listing = "\n\n".join(
        f"PAGE {page}:\n{text or '(no text)'}" for page, text in candidates
    )
    prompt = RERANK_PROMPT.format(query=query, n=len(candidates), candidates=listing)
    raw = _generate(
        [{"role": "user", "content": prompt}], max_new_tokens=8, trace_name="agent-rerank"
    )
    m = re.search(r"\d+", raw)
    if m:
        picked = int(m.group())
        for i, (page, _) in enumerate(candidates):
            if page == picked:
                return i, raw
    return 0, raw