Trixy7/plugins/stt_whisper/main.py

# -*- coding: utf-8 -*-
"""
Whisper STT Plugin - Speech-to-Text mit OpenAI Whisper.

Unterstützt:
- Lokale Modelle (tiny, base, small, medium, large)
- GPU-Beschleunigung
- Mehrere Sprachen
- Wort-Timings
- Automatischer Modell-Download
"""

import asyncio
import os
import time
from pathlib import Path
from typing import Any

from trixy_core.plugins import TrixyPlugin
from trixy_core.audio.stt import STTProvider, STTConfig, STTResult, STTState, WordTiming


# Whisper Modell-Größen
WHISPER_MODELS = {
    "tiny": {"size": "39 MB", "description": "Schnellstes Modell, geringste Genauigkeit"},
    "tiny.en": {"size": "39 MB", "description": "Nur Englisch, schnell"},
    "base": {"size": "74 MB", "description": "Gute Balance zwischen Geschwindigkeit und Genauigkeit"},
    "base.en": {"size": "74 MB", "description": "Nur Englisch, gute Balance"},
    "small": {"size": "244 MB", "description": "Bessere Genauigkeit, langsamer"},
    "small.en": {"size": "244 MB", "description": "Nur Englisch, bessere Genauigkeit"},
    "medium": {"size": "769 MB", "description": "Hohe Genauigkeit"},
    "medium.en": {"size": "769 MB", "description": "Nur Englisch, hohe Genauigkeit"},
    "large": {"size": "1550 MB", "description": "Höchste Genauigkeit, langsam"},
    "large-v2": {"size": "1550 MB", "description": "Verbesserte Version von large"},
    "large-v3": {"size": "1550 MB", "description": "Neueste Version, beste Genauigkeit"},
}


class WhisperSTTProvider(STTProvider):
    """Whisper-basierter STT-Provider."""

    def __init__(
        self,
        config: STTConfig | None = None,
        models_dir: Path | None = None,
        model_name: str = "base",
    ):
        super().__init__(config)
        self._models_dir = models_dir
        self._model_name = model_name
        self._model = None
        self._device = "cuda" if config and config.use_gpu else "cpu"

    @property
    def name(self) -> str:
        return "whisper"

    @property
    def supported_languages(self) -> list[str]:
        return [
            "de-DE", "en-US", "en-GB", "fr-FR", "es-ES", "it-IT",
            "pt-PT", "nl-NL", "pl-PL", "ru-RU", "ja-JP", "zh-CN",
            "ko-KR", "ar-SA", "hi-IN", "tr-TR", "vi-VN", "th-TH",
        ]

    @property
    def supports_word_timings(self) -> bool:
        return True

    async def initialize(self) -> None:
        """Lädt das Whisper-Modell."""
        try:
            import whisper
            from trixy_core.utils.debug import pinfo

            self._state = STTState.PROCESSING

            # Modell-Verzeichnis setzen (Whisper cached hier)
            if self._models_dir:
                self._models_dir.mkdir(parents=True, exist_ok=True)
                # Whisper nutzt XDG_CACHE_HOME oder ~/.cache
                os.environ["XDG_CACHE_HOME"] = str(self._models_dir)

            # Prüfen ob Modell existiert
            model_info = WHISPER_MODELS.get(self._model_name, {})
            size = model_info.get("size", "unbekannt")

            pinfo(f"Lade Whisper-Modell: {self._model_name} ({size})...")
            pinfo("Bei erstem Start wird das Modell heruntergeladen...")

            # Modell laden (in Thread wegen blocking)
            loop = asyncio.get_event_loop()
            self._model = await loop.run_in_executor(
                None,
                lambda: whisper.load_model(
                    self._model_name,
                    device=self._device,
                    download_root=str(self._models_dir) if self._models_dir else None
                )
            )

            self._model_loaded = True
            self._state = STTState.READY
            pinfo(f"Whisper-Modell geladen: {self._model_name}")

        except ImportError:
            raise RuntimeError(
                "Whisper nicht installiert. "
                "Installieren mit: pip install openai-whisper"
            )
        except Exception as e:
            self._state = STTState.ERROR
            raise RuntimeError(f"Fehler beim Laden des Whisper-Modells: {e}")

    async def shutdown(self) -> None:
        """Gibt Ressourcen frei."""
        self._model = None
        self._model_loaded = False
        self._state = STTState.UNINITIALIZED

    async def transcribe(
        self,
        audio_data: bytes,
        language: str | None = None,
    ) -> STTResult:
        """Transkribiert Audio mit Whisper."""
        if not self._model:
            raise RuntimeError("Whisper-Modell nicht geladen")

        self._state = STTState.PROCESSING
        start_time = time.time()

        try:
            import numpy as np
            import whisper

            # Audio-Daten zu numpy array
            audio_np = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0

            # Sprache bestimmen
            lang = None
            if language:
                lang = language.split("-")[0]  # "de-DE" -> "de"

            # Transkription in Thread
            loop = asyncio.get_event_loop()

            def do_transcribe():
                return self._model.transcribe(
                    audio_np,
                    language=lang,
                    word_timestamps=True,
                    fp16=(self._device == "cuda"),
                )

            result = await loop.run_in_executor(None, do_transcribe)

            processing_time = (time.time() - start_time) * 1000

            # Wort-Timings extrahieren
            word_timings = []
            if "segments" in result:
                for segment in result["segments"]:
                    if "words" in segment:
                        for word_info in segment["words"]:
                            word_timings.append(WordTiming(
                                word=word_info["word"].strip(),
                                start_time=word_info["start"],
                                end_time=word_info["end"],
                                confidence=word_info.get("probability", 1.0),
                            ))

            self._state = STTState.READY

            return STTResult(
                text=result["text"].strip(),
                confidence=1.0,  # Whisper gibt keine Confidence zurück
                language=result.get("language", language or "de"),
                duration_seconds=self.get_audio_duration(audio_data),
                processing_time_ms=processing_time,
                word_timings=word_timings,
                provider="whisper",
                model=self._model_name,
                raw_result=result,
            )

        except Exception as e:
            self._state = STTState.ERROR
            raise RuntimeError(f"Whisper-Transkription fehlgeschlagen: {e}")


class WhisperSTTPlugin(TrixyPlugin):
    """Whisper STT Plugin für Trixy."""

    def __init__(self, application, plugin_path, config: dict | None = None):
        super().__init__(application, plugin_path, config)
        self._provider: WhisperSTTProvider | None = None

    async def on_load(self) -> None:
        """Plugin wird geladen."""
        from trixy_core.utils.debug import pinfo
        pinfo("Whisper STT Plugin: Lade...")

        # Konfiguration
        model = self.config.get("model", "base")
        use_gpu = self.config.get("use_gpu", False)
        language = self.config.get("language", "de-DE")

        # Models-Verzeichnis im Plugin-Ordner
        models_dir = self.plugin_path / "models"
        models_dir.mkdir(parents=True, exist_ok=True)

        stt_config = STTConfig(
            language=language,
            use_gpu=use_gpu,
            model_name=model,
        )

        # Provider erstellen
        self._provider = WhisperSTTProvider(
            stt_config,
            models_dir=models_dir,
            model_name=model,
        )

        # Modell laden (und ggf. downloaden)
        await self._provider.initialize()

        # Extension registrieren
        if hasattr(self.application, "extension_points"):
            ext_point = self.application.extension_points.get("conversation.stt")
            if ext_point:
                ext_point.register(self._provider)
                pinfo("Whisper STT Plugin: Extension registriert")

        # Event-Handler registrieren
        self._register_event_handlers()

        pinfo(f"Whisper STT Plugin: Geladen (Modell: {model})")

    def _register_event_handlers(self) -> None:
        """Registriert Event-Handler."""
        em = self.application.events

        @em.on("raw_audio_received")
        async def on_audio_received(event_name: str, data: dict) -> None:
            """Verarbeitet empfangene Audio-Daten."""
            if not self._provider or not self._provider.is_ready:
                return

            audio_data = data.get("audio_data")
            satellite_id = data.get("satellite_id")
            session_id = data.get("session_id")

            if not audio_data:
                return

            # Audio in bytes konvertieren falls hex-encoded
            if isinstance(audio_data, str):
                audio_data = bytes.fromhex(audio_data)

            from trixy_core.utils.debug import pinfo
            pinfo(f"Whisper STT: Verarbeite Audio ({len(audio_data)} bytes)")

            try:
                # Transkribieren
                result = await self._provider.transcribe(audio_data)

                # Event emittieren
                await em.emit("stt_completed", {
                    "text": result.text,
                    "confidence": result.confidence,
                    "language": result.language,
                    "provider": "whisper",
                    "satellite_id": satellite_id,
                    "session_id": session_id,
                    "duration_seconds": result.duration_seconds,
                    "processing_time_ms": result.processing_time_ms,
                    "word_timings": [w.__dict__ for w in result.word_timings],
                })

                pinfo(f"Whisper STT: '{result.text}' ({result.processing_time_ms:.0f}ms)")

                # Event-Brücke: speech_recognized für NLP-Pipeline
                if result.text.strip():
                    from trixy_core.events.event_data.basic import SpeechRecognized

                    speech_event = SpeechRecognized(
                        satellite_id=satellite_id or "",
                        text=result.text,
                        confidence=result.confidence,
                        language=result.language or "de",
                        is_final=True,
                        source="stt",
                    )
                    speech_event.metadata["session_id"] = session_id or ""

                    await em.trigger("speech_recognized", speech_event)

            except Exception as e:
                from trixy_core.utils.debug import perror
                perror(f"Whisper STT Fehler: {e}")

                await em.emit("stt_error", {
                    "error": str(e),
                    "provider": "whisper",
                    "satellite_id": satellite_id,
                    "session_id": session_id,
                })

    async def on_unload(self) -> None:
        """Plugin wird entladen."""
        if self._provider:
            await self._provider.shutdown()
            self._provider = None

        from trixy_core.utils.debug import pinfo
        pinfo("Whisper STT Plugin: Entladen")

    @property
    def provider(self) -> WhisperSTTProvider | None:
        """STT-Provider."""
        return self._provider


# Plugin-Export
Plugin = WhisperSTTPlugin