splitbit_llm/voice/jarvis.py

"""Jarvis Voice Assistant — main orchestrator.

Ties together wake word → STT → LLM → TTS in a continuous loop.
Prebuilt and ready to go: `python -m splitbit_llm jarvis` starts listening.

Conversation loop:
1. Listen for wake word ("Jarvis")
2. Record user command until silence
3. Transcribe speech to text
4. Generate LLM response (voice-optimized mode)
5. Speak response sentence-by-sentence as generated
6. Self-improvement: feed conversation into learning pipeline
7. Return to step 1

Interruptible: user can say "Jarvis stop" to cancel TTS mid-speech.
Context aware: remembers conversation within a session using recursive links.
Personality: concise, direct, uncensored — designed for fast voice interactions.
Self-talking mode: "Jarvis, practice mode on" — generates synthetic training data.
"""

from __future__ import annotations

import logging
import time
from typing import Any

from ..harness.harness import SplitBitHarness
from .wake_word import WakeWordDetector
from .stt import STTEngine
from .tts import TTSEngine
from .voice_adapter import VoiceAdapter
from .tts_output import TTSOutputFormatter
from .self_improve import SelfImprovementEngine

logger = logging.getLogger(__name__)


class JarvisAssistant:
    """Built-in Jarvis voice assistant.

    Prebuilt voice assistant that uses the SplitBit LLM for responses.
    Wake word detection, speech-to-text, LLM response, text-to-speech.
    Self-improving: every conversation makes it smarter.
    """

    def __init__(self, harness: SplitBitHarness | None = None,
                 wake_word: str = "jarvis") -> None:
        self.harness = harness or SplitBitHarness()
        self.wake_word = wake_word

        # Voice components
        self.wake_detector = WakeWordDetector(wake_word=wake_word)
        self.stt = STTEngine()
        self.tts = TTSEngine()
        self.voice_adapter = VoiceAdapter()
        self.tts_formatter = TTSOutputFormatter()
        self.self_improve = SelfImprovementEngine(
            model=self.harness.model,
            tokenizer=self.harness.tokenizer,
        )

        # State
        self._running = False
        self._practice_mode = False
        self._session_id = f"jarvis-{int(time.time())}"
        self._conversation_count = 0

        logger.info("Jarvis assistant initialized (wake word: '%s')", wake_word)

    def run(self) -> None:
        """Start Jarvis — listen for wake word and process commands."""
        self._running = True
        logger.info("Jarvis is online. Say '%s' to start talking.", self.wake_word)

        # Start self-improvement idle monitor
        self.self_improve.start_idle_monitor(self._generate_response)

        # Start wake word detection
        self.wake_detector.start(
            on_wake=self._on_wake,
            on_command=self._on_command,
        )

    def stop(self) -> None:
        """Stop Jarvis."""
        self._running = False
        self.wake_detector.stop()
        self.tts.stop()
        self.self_improve.stop_idle_monitor()
        logger.info("Jarvis stopped. Had %d conversations.", self._conversation_count)

    def _on_wake(self) -> None:
        """Called when wake word is detected."""
        logger.info("Wake word detected!")
        # Brief acknowledgment beep/phrase
        self.tts.speak("Yes?", blocking=True)

    def _on_command(self, command_text: str) -> None:
        """Called when a command is transcribed."""
        if not command_text:
            # No transcription available — would need STT backend
            logger.debug("No command text (STT not available)")
            return

        self._process_command(command_text)

    def _process_command(self, text: str) -> None:
        """Process a voice command."""
        text = text.strip()

        # Check for special commands
        if text.lower() in [f"{self.wake_word} stop", "stop", "quiet"]:
            self.tts.stop()
            logger.info("TTS stopped by user")
            return

        if "practice mode on" in text.lower():
            self._practice_mode = True
            self.tts.speak("Practice mode enabled. I'll train myself when idle.", blocking=True)
            return

        if "practice mode off" in text.lower():
            self._practice_mode = False
            self.tts.speak("Practice mode disabled.", blocking=True)
            return

        if text.lower() in ["goodbye", "bye", "shut down", f"{self.wake_word} goodbye"]:
            self.tts.speak("Goodbye!", blocking=True)
            self.stop()
            return

        self._conversation_count += 1

        # Generate response via harness
        response = self.harness.chat(
            message=text,
            channel="voice",
            session_id=self._session_id,
        )

        response_text = response.get("response", "")
        elapsed = response.get("elapsed_s", 0)

        # Format for TTS
        formatted = self.tts_formatter.format(response_text)

        # Speak the response
        self.tts.speak(formatted, blocking=True)

        # Record for self-improvement
        confidence = min(0.9, 1.0 / max(elapsed, 0.1))
        self.self_improve.record_conversation(text, response_text, confidence=confidence)

        logger.info("Conversation #%d: '%s' → '%s' (%.2fs)",
                     self._conversation_count, text[:50], response_text[:50], elapsed)

    def _generate_response(self, prompt: str) -> str:
        """Generate a response — used by self-talk."""
        result = self.harness.chat(prompt, channel="voice")
        return result.get("response", "")

    def text_chat(self, text: str) -> str:
        """Process a text command (for when voice isn't available).

        Args:
            text: user's text input
        Returns:
            Jarvis's response text
        """
        self._conversation_count += 1
        response = self.harness.chat(
            message=text,
            channel="voice",
            session_id=self._session_id,
        )
        response_text = response.get("response", "")
        elapsed = response.get("elapsed_s", 0)

        # Try to speak it
        formatted = self.tts_formatter.format(response_text)
        self.tts.speak(formatted)

        # Record for self-improvement
        confidence = min(0.9, 1.0 / max(elapsed, 0.1))
        self.self_improve.record_conversation(text, response_text, confidence=confidence)

        return response_text

    def get_stats(self) -> dict[str, Any]:
        """Get comprehensive stats."""
        return {
            "jarvis": {
                "running": self._running,
                "conversation_count": self._conversation_count,
                "practice_mode": self._practice_mode,
                "session_id": self._session_id,
                "wake_word": self.wake_word,
            },
            "wake_word": self.wake_detector.get_stats(),
            "stt": self.stt.get_stats(),
            "tts": self.tts.get_stats(),
            "voice_adapter": self.voice_adapter.get_stats(),
            "tts_formatter": self.tts_formatter.get_stats(),
            "self_improvement": self.self_improve.get_stats(),
            "harness": self.harness.get_stats(),
        }