scripts/demo_autopilot.py

"""Terminal demo of the Plane Mode Scholar autopilot — fine-tuned Nemotron coach.

Records a real, local FLY autopilot run (no network, no mocks) so the agent
chain and grounded coach output can be captured as a submittable demo.

    PMS_GGUF_DIR=models/nemotron python scripts/demo_autopilot.py
"""

from __future__ import annotations

import json
import os
import sys
import tempfile
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

os.environ.setdefault("PMS_INFERENCE_BACKEND", "llamacpp")
os.environ.setdefault("PMS_LLAMACPP_MODE", "embedded")
os.environ.setdefault("PMS_USE_FINETUNED", "true")
os.environ.setdefault("PMS_GGUF_DIR", str(ROOT / "models" / "nemotron"))
os.environ.setdefault("PMS_LLAMACPP_N_GPU_LAYERS", "0")
os.environ.setdefault("PMS_MAX_NEW_TOKENS", "320")

C = {
    "reset": "\033[0m",
    "dim": "\033[2m",
    "bold": "\033[1m",
    "teal": "\033[38;5;43m",
    "copper": "\033[38;5;173m",
    "amber": "\033[38;5;214m",
    "green": "\033[38;5;48m",
    "blue": "\033[38;5;75m",
    "purple": "\033[38;5;141m",
    "red": "\033[38;5;203m",
    "gray": "\033[38;5;245m",
}

PHASE_COLOR = {
    "BOOT": "teal",
    "PACK": "teal",
    "SESSION": "blue",
    "PLAN": "amber",
    "RETRIEVE": "purple",
    "EXPLAIN": "green",
    "QUIZ": "blue",
    "MEMORY": "copper",
    "DONE": "teal",
    "ERROR": "red",
}

DEMO_NOTES = """# Introduction to Machine Learning

## Chapter 1: Supervised Learning
Supervised learning trains a model on labeled input-output pairs so it can
predict the output for unseen inputs. Classification predicts discrete labels;
regression predicts continuous values.

## Chapter 2: Gradient Descent
Gradient descent minimizes a loss function by repeatedly stepping in the
direction of the negative gradient. The learning rate controls the step size:
too large overshoots the minimum, too small converges slowly.

## Chapter 3: Overfitting and Regularization
A model overfits when it memorizes training noise and fails to generalize.
Regularization (L1/L2) and dropout penalize complexity to improve test accuracy.
"""


def c(text: str, color: str) -> str:
    return f"{C.get(color, '')}{text}{C['reset']}"


def slow_print(text: str, delay: float = 0.012) -> None:
    for ch in text:
        sys.stdout.write(ch)
        sys.stdout.flush()
        time.sleep(delay)
    sys.stdout.write("\n")
    sys.stdout.flush()


def banner() -> None:
    line = "═" * 64
    print(c(line, "teal"))
    print(c("  ✈  PLANE MODE SCHOLAR  //  AUTOPILOT DEMO", "bold"))
    print(c("     Fine-tuned NVIDIA Nemotron 3 Nano 4B  ·  offline study coach", "gray"))
    print(c(line, "teal"))
    print()


def chain_step(phase: str, message: str) -> None:
    color = PHASE_COLOR.get(phase, "gray")
    tag = c(f"[{phase:<8}]", color)
    print(f"  {tag} {c(message, 'gray')}")


def main() -> None:
    banner()

    from plane_mode_scholar.core import llm
    from plane_mode_scholar.core.study_agent import StudyAgent
    from plane_mode_scholar.gradio_ui.layout import STATE, quick_upload_pack

    user_id = "demo_judge"
    tmp = Path(tempfile.mkdtemp()) / "intro_to_ml.md"
    tmp.write_text(DEMO_NOTES, encoding="utf-8")

    print(c("▸ Uploading study material: intro_to_ml.md", "copper"))
    pack = quick_upload_pack([str(tmp)], user_id)
    print(c(f"  ✓ Indexed pack {pack.get('pack_id', '')[:24]}… "
            f"({pack.get('file_count', 1)} file, 3 chapters)", "green"))
    print()
    print(c("▸ Backend:", "copper"),
          c(f"{os.environ['PMS_INFERENCE_BACKEND']} (GGUF + study-coach LoRA)", "gray"))
    print(c("▸ Tap FLY → autonomous monitor → plan → act loop", "copper"))
    print()
    print(c("─" * 64, "teal"))
    print(c("  AGENT CHAIN", "amber"))
    print(c("─" * 64, "teal"))

    agent = StudyAgent(STATE.store, STATE.orchestrator)

    coach_text = ""
    quiz_shown = False
    t_start = time.time()

    for event in agent.run_autopilot(user_id=user_id, pack_id=pack["pack_id"], max_steps=4):
        phase = event.get("phase", "")
        msg = event.get("message", "")
        data = event.get("data", {})

        if phase == "STREAM":
            continue

        if phase == "EXPLAIN":
            chain_step("EXPLAIN", f"Coach explanation ready ({data.get('inference_ms', 0)} ms)")
            coach_text = data.get("response", msg)
            print()
            print(c("  ┌─ COACH ─────────────────────────────────────────────────", "green"))
            for para in coach_text.split("\n"):
                slow_print(c("  │ ", "green") + para, delay=0.004)
            print(c("  └─────────────────────────────────────────────────────────", "green"))
            print()
            continue

        if phase == "QUIZ" and not quiz_shown:
            q = data.get("question", {})
            chain_step("QUIZ", "Generated a comprehension check")
            print()
            print(c("  ┌─ QUIZ ──────────────────────────────────────────────────", "blue"))
            print(c("  │ ", "blue") + c(q.get("question", msg), "bold"))
            for i, opt in enumerate(q.get("options", [])):
                print(c("  │   ", "blue") + f"{'ABCD'[i]}. {opt}")
            print(c("  └─────────────────────────────────────────────────────────", "blue"))
            print()
            quiz_shown = True
            continue

        chain_step(phase, msg)
        time.sleep(0.25)

    elapsed = time.time() - t_start
    print(c("─" * 64, "teal"))
    print(c(f"  ✓ Autopilot complete in {elapsed:.1f}s — "
            f"grounded, offline, fine-tuned.", "green"))
    print(c("─" * 64, "teal"))


if __name__ == "__main__":
    main()