watch_sceen_ai/wsai/pipeline.py

"""Orchestrator: wires the perception loop and the conversation loop together."""

from __future__ import annotations

import asyncio
import logging

from .interfaces import (
    Brain,
    FrameSource,
    Reply,
    SpeechToText,
    TextChannel,
    TextToSpeech,
    Utterance,
    VisionBackend,
)
from .monitor import Monitor
from .state import SharedScreenContext

log = logging.getLogger("wsai.pipeline")


class Pipeline:
    """Runs two concurrent loops:

    * perception: FrameSource -> VisionBackend -> SharedScreenContext
    * conversation: (SpeechToText | TextChannel) -> Brain -> (TextToSpeech | TextChannel)

    Any half can be omitted. With no source/vision it runs eyes-free as a pure
    voice loop (STT -> Brain -> TTS); with no stt/tts it runs text-only; with no
    conversation it is a headless "just watch" configuration.
    """

    def __init__(
        self,
        *,
        source: FrameSource | None = None,
        vision: VisionBackend | None = None,
        brain: Brain,
        stt: SpeechToText | None = None,
        tts: TextToSpeech | None = None,
        text_channel: TextChannel | None = None,
        history_turns: int = 12,
        monitor: Monitor | None = None,
    ) -> None:
        self.source = source
        self.vision = vision
        self.brain = brain
        self.stt = stt
        self.tts = tts
        self.text_channel = text_channel
        self.monitor = monitor
        self.context = SharedScreenContext()
        self._history: list[tuple[str, str]] = []
        self._history_turns = history_turns

    # -- perception -------------------------------------------------------- #
    async def _perceive(self) -> None:
        if self.source is None or self.vision is None:
            return  # eyes-free (voice-only) configuration
        async for frame in self.source.frames():
            try:
                obs = await self.vision.describe(frame)
            except Exception as exc:  # a single bad frame must not kill the loop
                log.exception("vision.describe failed")
                if self.monitor is not None:
                    self.monitor.log("error", f"화면 이해 실패: {exc}")
                continue
            await self.context.update(obs)
            log.debug("screen: %s", obs.text[:120])

    # -- conversation ------------------------------------------------------ #
    async def _handle(self, utt: Utterance) -> None:
        if self.monitor is None:
            screen = await self.context.latest()
            reply = await self.brain.respond(utt.text, screen, self._history)
            self._remember(utt.text, reply.text)
            await self._emit(reply)
            return

        # Same work, but each stage is timed and streamed to the dashboard so a
        # viewer can see what was heard, what the brain answered, how long each
        # step took, and whether anything errored.
        turn = self.monitor.turn(source=utt.source)
        turn.heard(utt.text)
        try:
            async with turn.step("화면 맥락"):
                screen = await self.context.latest()
            async with turn.step("두뇌(생각)"):
                reply = await self.brain.respond(utt.text, screen, self._history)
            turn.replied(reply.text)
            self._remember(utt.text, reply.text)
            async with turn.step("응답(TTS/전송)"):
                await self._emit(reply)
        except Exception as exc:
            # finish() records the error and emits the single error-level log
            # event that bumps errors_total, so don't log the same failure twice.
            turn.finish(error=f"{type(exc).__name__}: {exc}")
            raise
        else:
            turn.finish()

    def _remember(self, user: str, ai: str) -> None:
        self._history.append((user, ai))
        if len(self._history) > self._history_turns:
            self._history = self._history[-self._history_turns :]

    async def _emit(self, reply: Reply) -> None:
        tasks = []
        if self.tts is not None:
            tasks.append(self.tts.speak(reply))
        if self.text_channel is not None:
            tasks.append(self.text_channel.send(reply))
        if not tasks:
            log.info("AI: %s", reply.text)
        else:
            await asyncio.gather(*tasks)

    async def _listen_voice(self) -> None:
        if self.stt is None:
            return
        if self.monitor is not None:
            self.monitor.set_status(listening=True)
            self.monitor.log("info", "음성 수신 시작 — 발화 대기 중")
        try:
            async for utt in self.stt.utterances():
                await self._handle(utt)
        finally:
            if self.monitor is not None:
                self.monitor.set_status(listening=False)

    async def _listen_text(self) -> None:
        if self.text_channel is None:
            return
        async for utt in self.text_channel.messages():
            await self._handle(utt)

    # -- lifecycle --------------------------------------------------------- #
    async def _prewarm(self) -> None:
        """Load slow-to-start backends before the loops accept input.

        Real STT/TTS backends (faster-whisper, MeloTTS) load a model into a
        persistent worker on first use — several seconds on CPU. Warming them
        here means the first real utterance is answered warm (~1s) instead of
        paying the cold model load mid-conversation."""
        warmers = []
        for comp, name in ((self.tts, "tts"), (self.stt, "stt")):
            warmup = getattr(comp, "warmup", None)
            if callable(warmup):
                warmers.append((name, warmup))
        if not warmers:
            return
        for name, warmup in warmers:
            try:
                await warmup()
            except Exception as exc:  # a warm failure must not abort startup
                log.warning("prewarm %s failed: %s", name, exc)
                if self.monitor is not None:
                    self.monitor.log("error", f"{name} 예열 실패: {exc}")

    async def run(self) -> None:
        # A TaskGroup (not bare gather) so that if ONE loop raises, the others
        # are cancelled and awaited before teardown. With plain gather the
        # failing loop propagated while the siblings kept running detached, and
        # aclose() in the finally then closed a source/stt out from under a
        # still-live loop (close-during-use).
        if self.monitor is not None:
            self.monitor.set_status(running=True)
            self.monitor.log("info", "파이프라인 시작")
        await self._prewarm()
        try:
            async with asyncio.TaskGroup() as tg:
                tg.create_task(self._perceive())
                tg.create_task(self._listen_voice())
                tg.create_task(self._listen_text())
        except* Exception as eg:
            if self.monitor is not None:
                for exc in eg.exceptions:
                    self.monitor.log("error", f"루프 예외: {type(exc).__name__}: {exc}")
            raise
        finally:
            if self.monitor is not None:
                self.monitor.set_status(running=False, listening=False)
                self.monitor.log("info", "파이프라인 종료")
            await self.aclose()

    async def aclose(self) -> None:
        # tts is included because a real TTS (e.g. MeloTTS) owns a worker
        # subprocess that must be torn down; mock backends have no aclose.
        for closer in (self.source, self.stt, self.text_channel, self.tts):
            if closer is not None and hasattr(closer, "aclose"):
                try:
                    await closer.aclose()
                except Exception:
                    log.exception("error closing %s", closer)