core/wakeup.py

"""
core/wakeup.py
Aiko's boot orchestrator — owns parallel subsystem startup and warmup sequencing.
main.py calls AikoWakeup().boot(...) and receives a BootResult with all live
subsystem references; it never needs to know the startup choreography.
Progress is reported through three injected callbacks so wakeup.py stays
completely TUI-ignorant:
    on_loading(key)  — subsystem is starting
    on_done(key)     — subsystem finished successfully
    on_skip(key)     — subsystem skipped
Each module owns its BOOT_LABELS dict; wakeup collects them and exposes
ALL_BOOT_LABELS so the UI can register display text before boot begins.
Usage:
    result = AikoWakeup().boot(
        on_loading = ...,
        on_done    = ...,
        on_skip    = ...,
    )
    think    = result.think
    memorize = result.memorize
Note: the dream scheduler is not used in this Gradio/HF Space demo, so
      `speak` is always None. ASR (`listen`) IS used here — browser-recorded
      audio is transcribed via core.listen.transcribe_file (Modal endpoint
      or local faster-whisper fallback) — but it has no persistent live
      object to hand back, so `listen` remains None in BootResult; its
      warmup is just a cold-start prefill (see _warmup_asr).
All seven warmups (LLM, TTS, ASR, VLM, SearXNG, plus think/memorize init)
fire real inference/synthesis/transcription/search requests rather than
health checks, so CUDA kernels and Modal containers are hot before the
first real user turn.
"""

import os
import time
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Callable
from core.log import get_logger
log = get_logger(__name__)

from core.think    import BOOT_LABELS as _THINK_LABELS
from core.memorize import BOOT_LABELS as _MEM_LABELS

# ── result container ──────────────────────────────────────────────────────────

@dataclass
class BootResult:
    """Holds all live subsystem references produced during boot."""
    think:    object        # AikoThink
    memorize: object        # AikoMemorize
    speak:    None = None   # Not used in Gradio deployment
    listen:   None = None   # No persistent object; see ui.listen.transcribe_file


# ── warmup helpers ────────────────────────────────────────────────────────────

def _warmup_llm(
    on_loading: Callable[[str], None],
    on_done:    Callable[[str], None],
    on_skip:    Callable[[str], None],
) -> None:
    """
    Proper LLM warmup:
      Run a real inference request with the actual soul.md system prompt
      so CUDA kernels are initialized and the system prompt prefix is
      hot in the KV cache — first real user turn will be much faster.
      No /health gate: the inference POST itself rides out a cold Modal
      container start.
    """
    import httpx

    base_url = os.getenv("LLAMA_BASE_URL", "").rstrip("/")
    if not base_url:
        log.debug("LLAMA_BASE_URL not set — skipping LLM warmup")
        on_skip("warmup_llm")
        return

    on_loading("warmup_llm")

    soul_path = Path(__file__).resolve().parent.parent / "persona" / "soul.md"
    try:
        system_prompt = soul_path.read_text(encoding="utf-8").strip() if soul_path.exists() else "You are Aiko."
    except Exception:
        system_prompt = "You are Aiko."

    # Replace template vars with neutral placeholders for warmup
    system_prompt = system_prompt.replace("USER_ID_HERE", "Guest").replace("TODAY_HERE", "today")

    try:
        r = httpx.post(
            base_url,  # /chat endpoint
            json={
                "model":       os.getenv("LLAMA_MODEL", "aiko"),
                "messages":    [
                    {"role": "system", "content": system_prompt},
                    {"role": "user",   "content": "hi"},
                ],
                "max_tokens":  16,    # just enough to force a full forward pass
                "temperature": 0.1,
            },
            timeout=180,  # cold Modal container can take a while to spin up
        )
        log.info("LLM warmup inference: status=%s", r.status_code)
        if r.status_code == 200:
            on_done("warmup_llm")
        else:
            log.warning("LLM warmup got non-200: %s %s", r.status_code, r.text[:200])
            on_skip("warmup_llm")
    except Exception as e:
        log.warning("LLM warmup inference failed (non-fatal): %s", e)
        on_skip("warmup_llm")


def _warmup_tts(
    on_loading: Callable[[str], None],
    on_done:    Callable[[str], None],
    on_skip:    Callable[[str], None],
) -> None:
    """
    Proper TTS warmup:
      Run a real synthesis request so the inner llama-server inside
      MioTTS initializes CUDA kernels and codec generation paths.
      No /health gate: the synthesis POST itself rides out a cold
      Modal container start.
    """
    import httpx

    url = os.getenv("MIOTTS_URL", "").rstrip("/")
    if not url:
        log.debug("MIOTTS_URL not set — skipping TTS warmup")
        on_skip("warmup_tts")
        return

    on_loading("warmup_tts")

    preset_id = os.getenv("MIOTTS_PRESET", "Aiko")
    try:
        r = httpx.post(
            f"{url}/v1/tts/file",
            data={
                "text": "Hello.",
                "reference_preset_id": preset_id,
            },
            timeout=180,
        )
        log.info("TTS warmup synthesis: status=%s bytes=%d", r.status_code, len(r.content))
        if r.status_code == 200:
            on_done("warmup_tts")
        else:
            log.warning("TTS warmup got non-200: %s %s", r.status_code, r.text[:200])
            on_skip("warmup_tts")
    except Exception as e:
        log.warning("TTS warmup synthesis failed (non-fatal): %s", e)
        on_skip("warmup_tts")


def _warmup_asr(
    on_loading: Callable[[str], None],
    on_done:    Callable[[str], None],
    on_skip:    Callable[[str], None],
) -> None:
    """
    Proper ASR warmup:
      Download a short real-speech sample from HuggingFace (LibriSpeech)
      and transcribe it so faster-whisper (Modal large-v3-turbo via
      AIKO_ASR_URL, or local CTranslate2 fallback) has its CUDA kernels
      fully initialized — encoder, decoder, and VAD all get exercised.
      Falls back to a synthetic silent WAV if the download fails.
      No /health gate: the transcription request itself rides out a
      cold Modal container start.
    """
    import wave
    import struct
    import tempfile
    import httpx

    from core.listen import ASR_URL, transcribe_file

    on_loading("warmup_asr")

    # Try downloading a real speech sample from HuggingFace first —
    # exercises the full encoder+decoder+VAD path rather than just silence.
    HF_ASR_SAMPLE = "static/harvard.wav"

    tmp_path = None
    try:
        # Attempt to fetch a real speech sample
        try:
            dl = httpx.get(HF_ASR_SAMPLE, timeout=15, follow_redirects=True)
            dl.raise_for_status()
            suffix = ".flac"
            audio_bytes = dl.content
            log.info("ASR warmup: downloaded HF sample (%d bytes)", len(audio_bytes))
        except Exception as dl_err:
            log.info("ASR warmup: HF download failed (%s), falling back to silent WAV", dl_err)
            # Fallback: ~0.5s of silence at 16kHz mono
            sample_rate = 16000
            num_samples = sample_rate // 2
            silence = struct.pack("<" + "h" * num_samples, *([0] * num_samples))
            import io
            buf = io.BytesIO()
            with wave.open(buf, "wb") as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)
                wf.setframerate(sample_rate)
                wf.writeframes(silence)
            audio_bytes = buf.getvalue()
            suffix = ".wav"

        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            tmp_path = tmp.name
            tmp.write(audio_bytes)

        if ASR_URL:
            mime = "audio/flac" if suffix == ".flac" else "audio/wav"
            with open(tmp_path, "rb") as f:
                resp = httpx.post(
                    f"{ASR_URL}/transcribe",
                    files={"audio": (f"warmup{suffix}", f, mime)},
                    timeout=180,  # cold Modal container can take a while to spin up
                )
            resp.raise_for_status()
            log.info("ASR warmup (Modal): status=%s text=%r",
                     resp.status_code, resp.json().get("text", "")[:80])
        else:
            text = transcribe_file(tmp_path)
            log.info("ASR warmup (local): transcript=%r", text)

        on_done("warmup_asr")
    except Exception as e:
        log.warning("ASR warmup failed (non-fatal): %s", e)
        on_skip("warmup_asr")
    finally:
        if tmp_path:
            try:
                os.unlink(tmp_path)
            except OSError:
                pass


def _warmup_vlm(
    on_loading: Callable[[str], None],
    on_done:    Callable[[str], None],
    on_skip:    Callable[[str], None],
) -> None:
    """
    Proper VLM warmup:
      Send a real image inference request using a public HuggingFace
      sample (the same refract.png used in minicpmv.py's local test)
      so the MiniCPM-V Modal container loads the model, initializes
      CUDA kernels, and runs a full forward pass before the first
      user upload.
      Uses image_url so the Modal container fetches the image from HF
      directly — no local download needed.
    """
    import httpx

    endpoint = os.getenv("VISION_ENDPOINT", "").rstrip("/")
    if not endpoint:
        log.debug("VISION_ENDPOINT not set — skipping VLM warmup")
        on_skip("warmup_vlm")
        return

    on_loading("warmup_vlm")

    # Public HF sample image — same one used in backend/minicpmv.py main()
    HF_VLM_SAMPLE = (
        "https://huggingface.co/datasets/openbmb/DemoCase/"
        "resolve/main/refract.png"
    )

    try:
        r = httpx.post(
            endpoint,
            json={
                "prompt":         "Describe this image briefly.",
                "image_url":      HF_VLM_SAMPLE,
                "max_new_tokens": 16,   # minimal generation — just warm the GPU
            },
            timeout=180,  # cold Modal container with GPU can take a while
        )
        log.info("VLM warmup inference: status=%s text=%r",
                 r.status_code, r.json().get("text", "")[:80])
        if r.status_code == 200:
            on_done("warmup_vlm")
        else:
            log.warning("VLM warmup got non-200: %s %s", r.status_code, r.text[:200])
            on_skip("warmup_vlm")
    except Exception as e:
        log.warning("VLM warmup inference failed (non-fatal): %s", e)
        on_skip("warmup_vlm")


def _warmup_searxng(
    on_loading: Callable[[str], None],
    on_done:    Callable[[str], None],
    on_skip:    Callable[[str], None],
) -> None:
    """
    Proper SearXNG warmup:
      Fire a real search query so the Modal SearXNG container starts
      its internal Flask/Werkzeug server and warms its engine connections
      (DuckDuckGo, Brave, Wikipedia) before the first user search.
      A lightweight "hello" query is enough to trigger a full cold start.
    """
    import httpx

    base_url = os.getenv("SEARXNG_BASE_URL", "").rstrip("/")
    if not base_url:
        log.debug("SEARXNG_BASE_URL not set — skipping SearXNG warmup")
        on_skip("warmup_searxng")
        return

    on_loading("warmup_searxng")

    try:
        r = httpx.get(
            base_url,
            params={"q": "hello", "format": "json", "language": "en"},
            timeout=60,  # SearXNG is CPU-only, spins up faster than GPU containers
        )
        n_results = len(r.json().get("results", []))
        log.info("SearXNG warmup: status=%s results=%d", r.status_code, n_results)
        if r.status_code == 200:
            on_done("warmup_searxng")
        else:
            log.warning("SearXNG warmup got non-200: %s %s", r.status_code, r.text[:200])
            on_skip("warmup_searxng")
    except Exception as e:
        log.warning("SearXNG warmup failed (non-fatal): %s", e)
        on_skip("warmup_searxng")


# ── wakeup ────────────────────────────────────────────────────────────────────

class AikoWakeup:
    """
    Parallel boot orchestrator for Aiko cognitive subsystems.
    Boots AikoThink, AikoMemorize, and the LLM/TTS/ASR/VLM/SearXNG Modal
    servers concurrently with granular progress reporting per step.
    Dream scheduler is excluded from this demo deployment.
    Modal server warmups (LLM + TTS + ASR + VLM + SearXNG) fire immediately
    and run in parallel with think/memorize init so boot time is minimized.
    """

    ALL_BOOT_LABELS: dict[str, str] = {
        **_THINK_LABELS,
        **_MEM_LABELS,
        "warmup_llm":    "🧠 Waking LLM server...",
        "warmup_tts":    "🔊 Waking voice server...",
        "warmup_asr":    "🎙️ Waking ear...",
        "warmup_vlm":    "👁️ Waking vision model...",
        "warmup_searxng": "🔍 Waking search engine...",
        "speak_skip":    "TTS skipped (cloud mode)",
        "listen_skip":   "ASR skipped (cloud mode)",
    }

    def warm_servers_async(self) -> None:
        """Fire-and-forget network warmups for page load in serverless environments."""
        def noop(k): pass
        threading.Thread(target=_warmup_llm,     args=(noop, noop, noop), daemon=True, name="warmup-llm-async").start()
        threading.Thread(target=_warmup_tts,     args=(noop, noop, noop), daemon=True, name="warmup-tts-async").start()
        threading.Thread(target=_warmup_asr,     args=(noop, noop, noop), daemon=True, name="warmup-asr-async").start()
        threading.Thread(target=_warmup_vlm,     args=(noop, noop, noop), daemon=True, name="warmup-vlm-async").start()
        threading.Thread(target=_warmup_searxng, args=(noop, noop, noop), daemon=True, name="warmup-searxng-async").start()

    def boot(
        self,
        on_loading: Callable[[str], None],
        on_done:    Callable[[str], None],
        on_skip:    Callable[[str], None],
    ) -> BootResult:
        """
        Execute boot sequence and return live subsystem references.
        All seven subsystems boot concurrently:
          - AikoThink          (LLM client + persona load + internal warmup)
          - AikoMemorize       (sqlite-vec + fastembed + cleanup)
          - Modal LLM warmup   (real inference prefill)
          - Modal TTS warmup   (real synthesis)
          - ASR warmup         (real transcription with HF speech sample)
          - VLM warmup         (real vision inference with HF image)
          - SearXNG warmup     (real search query)
        Memory backend is injected into think once both are ready.
        All warmups are joined before returning so the first user
        message hits warm containers.
        Args:
            on_loading: Called with a progress key when a subsystem starts.
            on_done:    Called with a progress key when a subsystem finishes.
            on_skip:    Called with a progress key when a subsystem is skipped.
        Returns:
            BootResult with think and memorize references; speak and
            listen are always None in this deployment (see module docstring).
        """
        from core.memorize import AikoMemorize
        from core.think    import AikoThink

        memorize  = [None]
        think_ref = [None]
        mem_ready = threading.Event()

        # ── think ─────────────────────────────────────────────────────────────
        def init_think():
            on_loading("think_start")
            think_ref[0] = AikoThink(None, speak=None)
            on_done("think_start")
            on_loading("think_warmup")
            # AikoThink fires its own internal warmup on __init__; join it here
            # Note: we do NOT fire a second Modal LLM warmup from think —
            # _warmup_llm() below handles that with the full soul.md system prompt.
            think_ref[0].join_warmup()
            on_done("think_warmup")
            mem_ready.wait()                        # hold until memorize is ready
            think_ref[0]._memorize = memorize[0]   # inject memory backend

        # ── memorize ──────────────────────────────────────────────────────────
        def init_memorize():
            try:
                on_loading("mem_sqlite_vec")
                memorize[0] = AikoMemorize(silent=True)
                on_done("mem_sqlite_vec")
                on_loading("mem_embed")
                on_done("mem_embed")
                on_loading("mem_cleanup")
                memorize[0].cleanup()
                on_done("mem_cleanup")
                on_loading("mem_ready")
                on_done("mem_ready")
            except Exception as e:
                log.error("Memory boot failed: %s", e)
            finally:
                mem_ready.set()  # always unblock init_think, even on failure

        # ── launch all seven concurrently ───────────────────────────────────────
        t_think    = threading.Thread(target=init_think,    daemon=True, name="boot-think")
        t_memorize = threading.Thread(target=init_memorize, daemon=True, name="boot-memorize")
        t_llm      = threading.Thread(
            target=_warmup_llm, args=(on_loading, on_done, on_skip),
            daemon=True, name="warmup-llm"
        )
        t_tts      = threading.Thread(
            target=_warmup_tts, args=(on_loading, on_done, on_skip),
            daemon=True, name="warmup-tts"
        )
        t_asr      = threading.Thread(
            target=_warmup_asr, args=(on_loading, on_done, on_skip),
            daemon=True, name="warmup-asr"
        )
        t_vlm      = threading.Thread(
            target=_warmup_vlm, args=(on_loading, on_done, on_skip),
            daemon=True, name="warmup-vlm"
        )
        t_searxng  = threading.Thread(
            target=_warmup_searxng, args=(on_loading, on_done, on_skip),
            daemon=True, name="warmup-searxng"
        )

        t_think.start()
        t_memorize.start()
        t_llm.start()
        t_tts.start()
        t_asr.start()
        t_vlm.start()
        t_searxng.start()

        # Wait for all to complete before returning
        t_think.join()
        t_memorize.join()
        t_llm.join()
        t_tts.join()
        t_asr.join()
        t_vlm.join()
        t_searxng.join()

        on_skip("speak_skip")
        on_skip("listen_skip")

        return BootResult(
            think    = think_ref[0],
            memorize = memorize[0],
        )