mind/planner.py

"""The Planner — NEMOCITY one-shot mind (no turn loop, ever).

grant(): ONE backend generation turns a petition into the BUILD permit JSON
(mind/prompts.py), parsed leniently (mind/validate.py), retried once at
temperature 0 ("!strict" tag) with the parse error in context, then a
deterministic fallback (one house) — a petition NEVER builds nothing. The
engine executes every act and owns all placement; the planner never touches
world state. ONE @spaces.GPU call per petition in the normal path.

Interface (mirrors godseed's planner/queue contract — the queue worker calls
grant positionally with (wish, world_summary, act, emit)):

    await grant(wish, world_summary, act, emit, wish_id=None) -> trace dict
        act:  async (building: dict) -> observation str. Called once per
              building (1..3) with the stable 5-key shape
              {"kind": str, "name": str, "near": str|None,
               "floors": int|None, "hue": int|None}.
        emit: async (event: dict) -> None. Emitted, in order:
              {"type": "plan", "plan": {...validated BUILD JSON...}}  (once)
              {"type": "thought_token", "text": ...}  (blurb + per-act lines)
              wish_id is included in the plan event only when the caller
              passes it; otherwise the worker injects it before broadcast
              (the godseed convention).

    await fix(stats, candidates, act_fix, emit, wish_id=None) -> trace dict
        act_fix: async ({"choice","diagnosis","blurb"}) -> observation str,
        called exactly once. Parse failure / invalid / missing choice falls
        back to the candidate with the lowest predicted_index plus a canned
        diagnosis built from stats.

Trace shape (the queue worker merges wish_id/text/submitted_at/moderation/
feature_ids, exactly like godseed):

    {"reading": blurb, "turns": [{"thought","call","observation"}...],
     "epitaph": blurb, "ms_total": int, "model": str, "backend": str,
     "raw_json": validated plan dict, "raw_text": raw model output,
     "fallback": bool, "declined": bool, "decline_reason": str|None}

For fix traces reading=diagnosis and epitaph=blurb. A declined petition has
turns=[] and declined=True — the worker should emit wish_rejected with the
decline_reason instead of wish_granted.
"""

from __future__ import annotations

import re
import time
from typing import Awaitable, Callable

from . import prompts
from .backends import make_backend
from .validate import BUILD_TAG, FIX_TAG, STRICT_SUFFIX, parse_build, parse_fix

Act = Callable[[dict], Awaitable[str]]
Emit = Callable[[dict], Awaitable[None]]

FALLBACK_BLURB = "The city granted a home."

_RAW_MAX_CHARS = 6000
_OBSERVATION_MAX_CHARS = 500

_STOP_WORDS = {
    "a", "an", "the", "please", "build", "make", "made", "add", "want",
    "wants", "need", "needs", "near", "for", "with", "of", "in", "on", "to",
    "and", "my", "me", "us", "we", "i", "new", "some", "little", "big",
}


def _default_name(wish: str) -> str:
    """A building name derived from the petition's own words."""
    words = [w for w in re.findall(r"[A-Za-z]+", str(wish or "")) if w.lower() not in _STOP_WORDS]
    base = " ".join(w.capitalize() for w in words[:2])
    name = f"{base} House".strip()
    return (name[:24].rstrip() or "Petition House")[:24]


def _fallback_plan(wish: str) -> dict:
    """The safety net: a petition must never build nothing."""
    return {
        "intent": "build",
        "blurb": FALLBACK_BLURB,
        "buildings": [
            {"kind": "house", "name": _default_name(wish), "near": None, "floors": None, "hue": None}
        ],
        "decline_reason": None,
    }


def _thought_for(building: dict) -> str:
    kind = str((building or {}).get("kind") or "building")
    name = str((building or {}).get("name") or "").strip()
    return f"Permit issued: {name} ({kind})." if name else f"Permit issued: a {kind}."


def _canned_fix(stats: dict, candidate: dict | None) -> tuple[str, str]:
    """Deterministic diagnosis + blurb citing whatever real numbers we have."""
    worst = " ".join(str(stats.get("worst") or "the worst crossing").split())[:48]
    ti = stats.get("traffic_index")
    pred = (candidate or {}).get("predicted_index")
    desc = " ".join(str((candidate or {}).get("desc") or "").split())[:80]
    if isinstance(ti, (int, float)) and not isinstance(ti, bool):
        diagnosis = f"Traffic index {int(ti)}; {worst} is over capacity."
    else:
        diagnosis = f"The network strains at {worst}."
    if isinstance(pred, (int, float)) and not isinstance(pred, bool):
        diagnosis += f" {desc or 'The fix'} brings the predicted index to {int(pred)}."
    blurb = f"City Engineer orders relief at {worst}."
    return diagnosis[:160], blurb[:120]


class Planner:
    """Grants one petition (or one traffic fix) at a time against an injected
    backend. Owns prompting, parsing, and pacing; never raises for backend or
    act misbehavior — it always returns a complete trace."""

    BUILD_MAX_TOKENS = 700
    FIX_MAX_TOKENS = 300

    def __init__(self, backend=None):
        self.backend = backend if backend is not None else make_backend()

    # ------------------------------------------------------------------ emit

    @staticmethod
    async def _safe_emit(emit: Emit, event: dict) -> None:
        """Broadcast failures must never kill a grant."""
        try:
            await emit(event)
        except Exception:
            pass

    async def _emit_text(self, emit: Emit, text: str, *, end: str = "\n") -> None:
        """Stream already-known text as small chunks (the ticker stays alive)."""
        words = text.split()
        step = 3
        for i in range(0, len(words), step):
            chunk = " ".join(words[i : i + step])
            if i + step < len(words):
                chunk += " "
            await self._safe_emit(emit, {"type": "thought_token", "text": chunk})
        if end:
            await self._safe_emit(emit, {"type": "thought_token", "text": end})

    async def _emit_plan(self, emit: Emit, plan: dict, wish_id: str | None) -> None:
        event = {"type": "plan", "plan": plan}
        if wish_id is not None:
            event["wish_id"] = wish_id
        await self._safe_emit(emit, event)

    # --------------------------------------------------------------- collect

    async def _collect(self, prompt: str, tag: str, max_tokens: int) -> str:
        parts: list[str] = []
        total = 0
        stream = self.backend.generate_stream(prompt, tag, max_tokens)
        try:
            async for chunk in stream:
                if not chunk:
                    continue
                parts.append(chunk)
                total += len(chunk)
                if total >= _RAW_MAX_CHARS:
                    break
        finally:
            try:
                await stream.aclose()
            except Exception:
                pass
        return "".join(parts)

    async def _attempt(self, prompt: str, tag: str, max_tokens: int) -> str:
        try:
            return await self._collect(prompt, tag, max_tokens)
        except Exception:
            return ""  # a dead backend falls through to the parser/fallback

    async def _plan_with_retry(
        self, prompt: str, tag: str, max_tokens: int, parser
    ) -> tuple[dict | None, str]:
        """One generation + one temperature-0 retry. Returns (obj|None, raw)."""
        raw = await self._attempt(prompt, tag, max_tokens)
        obj, err = parser(raw)
        if err is None:
            return obj, raw
        retry_prompt = (
            f"{prompt}\n{prompts.REASK_NOTICE} ({err}). Output ONLY the JSON object."
        )
        raw2 = await self._attempt(retry_prompt, tag + STRICT_SUFFIX, max_tokens)
        obj, err = parser(raw2)
        if err is None:
            return obj, raw2
        return None, raw or raw2

    def _trace(self, *, reading, turns, epitaph, started, raw_json, raw_text,
               fallback=False, declined=False, decline_reason=None) -> dict:
        return {
            "reading": reading,
            "turns": turns,
            "epitaph": epitaph,
            "ms_total": int((time.perf_counter() - started) * 1000),
            "model": str(getattr(self.backend, "model_id", "unknown")),
            "backend": str(getattr(self.backend, "name", "unknown")),
            "raw_json": raw_json,
            "raw_text": str(raw_text or "")[:_RAW_MAX_CHARS],
            "fallback": fallback,
            "declined": declined,
            "decline_reason": decline_reason,
        }

    # ----------------------------------------------------------------- grant

    async def grant(
        self, wish: str, world_summary: str, act: Act, emit: Emit, *, wish_id: str | None = None
    ) -> dict:
        """Grant one petition. See the module docstring for the contract."""
        started = time.perf_counter()
        wish = str(wish or "")
        world_summary = str(world_summary or "")

        prompt = prompts.build_build_prompt(wish, world_summary)
        plan, raw = await self._plan_with_retry(
            prompt, BUILD_TAG, self.BUILD_MAX_TOKENS, parse_build
        )
        fallback = plan is None
        if fallback:
            plan = _fallback_plan(wish)

        if plan["intent"] == "decline":
            await self._emit_plan(emit, plan, wish_id)
            reason = plan.get("decline_reason") or "the petition could not be read as a building"
            blurb = plan.get("blurb") or "City hall declines this petition."
            await self._emit_text(emit, blurb, end="\n\n")
            return self._trace(
                reading=blurb, turns=[], epitaph=blurb, started=started,
                raw_json=plan, raw_text=raw, declined=True, decline_reason=reason,
            )

        for building in plan["buildings"]:
            if not building.get("name"):
                building["name"] = _default_name(wish)
        blurb = plan.get("blurb") or FALLBACK_BLURB
        plan["blurb"] = blurb

        await self._emit_plan(emit, plan, wish_id)
        await self._emit_text(emit, blurb, end="\n\n")

        turns: list[dict] = []
        for building in plan["buildings"][:3]:
            thought = _thought_for(building)
            await self._emit_text(emit, thought)
            try:
                observation = await act(building)
            except Exception as exc:
                observation = f"error: act failed ({type(exc).__name__})"
            turns.append({
                "thought": thought,
                "call": dict(building),
                "observation": str(observation or "")[:_OBSERVATION_MAX_CHARS],
            })

        return self._trace(
            reading=blurb, turns=turns, epitaph=blurb, started=started,
            raw_json=plan, raw_text=raw, fallback=fallback,
        )

    # ------------------------------------------------------------------- fix

    async def fix(
        self, stats, candidates, act_fix: Act, emit: Emit, *, wish_id: str | None = None
    ) -> dict:
        """One traffic fix: {stats, candidates} -> ONE generation -> one
        act_fix call. Engine pre-validated the candidates; an invalid or
        missing model choice falls back to the best predicted_index."""
        started = time.perf_counter()
        stats = stats if isinstance(stats, dict) else {}
        candidates = [c for c in (candidates or ()) if isinstance(c, dict)]

        prompt = prompts.build_fix_prompt(stats, candidates)
        obj, raw = await self._plan_with_retry(
            prompt, FIX_TAG, self.FIX_MAX_TOKENS, parse_fix
        )

        by_id = {str(c.get("id")): c for c in candidates if c.get("id")}
        choice = None
        if obj and obj.get("choice"):
            wanted = str(obj["choice"]).strip().lower()
            for cid in by_id:
                if cid.lower() == wanted:
                    choice = cid
                    break

        fallback = choice is None
        if fallback:
            def _pred(c: dict) -> float:
                value = c.get("predicted_index")
                ok = isinstance(value, (int, float)) and not isinstance(value, bool)
                return float(value) if ok else float("inf")

            best = min(candidates, key=_pred) if candidates else None
            choice = str(best.get("id")) if best and best.get("id") else None
            diagnosis, blurb = _canned_fix(stats, best)
        else:
            canned = _canned_fix(stats, by_id[choice])
            diagnosis = (obj.get("diagnosis") or "").strip()[:160] or canned[0]
            blurb = (obj.get("blurb") or "").strip()[:120] or canned[1]

        plan = {"diagnosis": diagnosis, "choice": choice, "blurb": blurb}
        await self._emit_plan(emit, plan, wish_id)
        await self._emit_text(emit, diagnosis, end="\n\n")

        try:
            observation = await act_fix(plan)
        except Exception as exc:
            observation = f"error: act failed ({type(exc).__name__})"
        turns = [{
            "thought": diagnosis,
            "call": dict(plan),
            "observation": str(observation or "")[:_OBSERVATION_MAX_CHARS],
        }]

        return self._trace(
            reading=diagnosis, turns=turns, epitaph=blurb, started=started,
            raw_json=plan, raw_text=raw, fallback=fallback,
        )