Trixy7/plugins/stt_vosk/main.py

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

Unterstützt:
- Offline-Erkennung
- Schnell und ressourcenschonend
- Mehrere Sprachen
- Streaming
- Auto-Download von Modellen
"""

import asyncio
import json
import time
from pathlib import Path
from typing import Any, AsyncIterator

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


# Vosk Modell-Katalog mit Download-URLs
VOSK_MODELS = {
    "de": {
        "name": "vosk-model-de-0.21",
        "url": "https://alphacephei.com/vosk/models/vosk-model-de-0.21.zip",
        "size": "1.9 GB",
        "description": "Großes deutsches Modell",
    },
    "de-small": {
        "name": "vosk-model-small-de-0.15",
        "url": "https://alphacephei.com/vosk/models/vosk-model-small-de-0.15.zip",
        "size": "45 MB",
        "description": "Kleines deutsches Modell (schnell)",
    },
    "en": {
        "name": "vosk-model-en-us-0.22",
        "url": "https://alphacephei.com/vosk/models/vosk-model-en-us-0.22.zip",
        "size": "1.8 GB",
        "description": "Großes englisches Modell (US)",
    },
    "en-small": {
        "name": "vosk-model-small-en-us-0.15",
        "url": "https://alphacephei.com/vosk/models/vosk-model-small-en-us-0.15.zip",
        "size": "40 MB",
        "description": "Kleines englisches Modell (US)",
    },
    "fr": {
        "name": "vosk-model-fr-0.22",
        "url": "https://alphacephei.com/vosk/models/vosk-model-fr-0.22.zip",
        "size": "1.4 GB",
        "description": "Großes französisches Modell",
    },
    "fr-small": {
        "name": "vosk-model-small-fr-0.22",
        "url": "https://alphacephei.com/vosk/models/vosk-model-small-fr-0.22.zip",
        "size": "41 MB",
        "description": "Kleines französisches Modell",
    },
    "es": {
        "name": "vosk-model-es-0.42",
        "url": "https://alphacephei.com/vosk/models/vosk-model-es-0.42.zip",
        "size": "1.4 GB",
        "description": "Großes spanisches Modell",
    },
    "it": {
        "name": "vosk-model-it-0.22",
        "url": "https://alphacephei.com/vosk/models/vosk-model-it-0.22.zip",
        "size": "1.2 GB",
        "description": "Großes italienisches Modell",
    },
    "ru": {
        "name": "vosk-model-ru-0.42",
        "url": "https://alphacephei.com/vosk/models/vosk-model-ru-0.42.zip",
        "size": "1.8 GB",
        "description": "Großes russisches Modell",
    },
    "cn": {
        "name": "vosk-model-cn-0.22",
        "url": "https://alphacephei.com/vosk/models/vosk-model-cn-0.22.zip",
        "size": "1.3 GB",
        "description": "Großes chinesisches Modell",
    },
}


class VoskSTTProvider(STTProvider):
    """Vosk-basierter STT-Provider."""

    def __init__(
        self,
        config: STTConfig | None = None,
        models_dir: Path | None = None,
        model_name: str = "de",
        auto_download: bool = True,
        event_manager: Any = None,
    ):
        super().__init__(config)
        self._models_dir = models_dir
        self._model_name = model_name
        self._auto_download = auto_download
        self._event_manager = event_manager
        self._model = None
        self._recognizer = None
        self._model_path: Path | None = None

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

    @property
    def supported_languages(self) -> list[str]:
        return ["de-DE", "en-US", "fr-FR", "es-ES", "it-IT", "ru-RU", "zh-CN"]

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

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

    async def initialize(self) -> None:
        """Lädt das Vosk-Modell."""
        try:
            from vosk import Model, KaldiRecognizer, SetLogLevel

            # Vosk-Logging reduzieren
            SetLogLevel(-1)

            self._state = STTState.PROCESSING

            # Modell-Info
            if self._model_name not in VOSK_MODELS:
                raise RuntimeError(
                    f"Unbekanntes Modell: {self._model_name}. "
                    f"Verfügbar: {', '.join(VOSK_MODELS.keys())}"
                )

            model_info = VOSK_MODELS[self._model_name]
            model_dir_name = model_info["name"]
            self._model_path = self._models_dir / model_dir_name

            # Auto-Download wenn nicht vorhanden
            if not self._model_path.exists():
                if self._auto_download:
                    await self._download_model()
                else:
                    raise RuntimeError(
                        f"Vosk-Modell nicht gefunden: {self._model_path}. "
                        "Auto-Download ist deaktiviert."
                    )

            # Modell laden (in Thread)
            loop = asyncio.get_event_loop()
            self._model = await loop.run_in_executor(
                None,
                lambda: Model(str(self._model_path))
            )

            # Recognizer erstellen
            self._recognizer = KaldiRecognizer(
                self._model,
                self._config.sample_rate
            )
            self._recognizer.SetWords(True)

            self._model_loaded = True
            self._state = STTState.READY

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

    async def _download_model(self) -> None:
        """Lädt das Modell herunter und extrahiert es."""
        from trixy_core.utils.debug import pinfo
        from trixy_core.utils.download import download_and_extract

        model_info = VOSK_MODELS[self._model_name]
        url = model_info["url"]
        size = model_info["size"]

        self._models_dir.mkdir(parents=True, exist_ok=True)

        pinfo(f"Lade Vosk-Modell: {self._model_name} ({size})...")

        # Progressbar wird automatisch von DownloadProgress angezeigt
        success = await download_and_extract(
            url=url,
            dest_dir=self._models_dir,
            event_manager=self._event_manager,
            download_id=f"vosk-{self._model_name}",
        )

        if not success:
            raise RuntimeError(f"Download des Modells fehlgeschlagen: {self._model_name}")

        pinfo(f"Vosk-Modell heruntergeladen: {self._model_name}")

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

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

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

        try:
            from vosk import KaldiRecognizer

            # Neuen Recognizer für diese Anfrage
            recognizer = KaldiRecognizer(self._model, self._config.sample_rate)
            recognizer.SetWords(True)

            # Audio verarbeiten (in Thread)
            loop = asyncio.get_event_loop()

            def do_recognize():
                recognizer.AcceptWaveform(audio_data)
                return recognizer.FinalResult()

            result_json = await loop.run_in_executor(None, do_recognize)
            result = json.loads(result_json)

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

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

            self._state = STTState.READY

            return STTResult(
                text=result.get("text", "").strip(),
                confidence=self._calculate_confidence(result),
                language=language or self._config.language,
                duration_seconds=self.get_audio_duration(audio_data),
                processing_time_ms=processing_time,
                word_timings=word_timings,
                provider="vosk",
                model=self._model_name,
                raw_result=result,
            )

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

    def _calculate_confidence(self, result: dict) -> float:
        """Berechnet Durchschnitts-Confidence."""
        if "result" not in result or not result["result"]:
            return 0.0

        confidences = [w.get("conf", 1.0) for w in result["result"]]
        return sum(confidences) / len(confidences)

    async def transcribe_stream(
        self,
        audio_stream: AsyncIterator[bytes],
        language: str | None = None,
    ) -> AsyncIterator[STTResult]:
        """Transkribiert Audio-Stream mit Vosk."""
        if not self._model:
            raise RuntimeError("Vosk-Modell nicht geladen")

        from vosk import KaldiRecognizer

        recognizer = KaldiRecognizer(self._model, self._config.sample_rate)
        recognizer.SetWords(True)

        start_time = time.time()
        all_audio = b""

        async for chunk in audio_stream:
            all_audio += chunk

            # Partial result
            if recognizer.AcceptWaveform(chunk):
                result_json = recognizer.Result()
                result = json.loads(result_json)

                if result.get("text"):
                    yield STTResult(
                        text=result["text"].strip(),
                        confidence=self._calculate_confidence(result),
                        language=language or self._config.language,
                        is_final=False,
                        provider="vosk",
                        model=self._model_name,
                    )

        # Final result
        result_json = recognizer.FinalResult()
        result = json.loads(result_json)
        processing_time = (time.time() - start_time) * 1000

        word_timings = []
        if "result" in result:
            for word_info in result["result"]:
                word_timings.append(WordTiming(
                    word=word_info["word"],
                    start_time=word_info["start"],
                    end_time=word_info["end"],
                    confidence=word_info.get("conf", 1.0),
                ))

        yield STTResult(
            text=result.get("text", "").strip(),
            confidence=self._calculate_confidence(result),
            language=language or self._config.language,
            duration_seconds=self.get_audio_duration(all_audio),
            processing_time_ms=processing_time,
            word_timings=word_timings,
            is_final=True,
            provider="vosk",
            model=self._model_name,
            raw_result=result,
        )


class VoskSTTPlugin(TrixyPlugin):
    """Vosk STT Plugin für Trixy."""

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

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

        # Konfiguration
        model_name = self.config.get("model", "de-small")
        language = self.config.get("language", "de-DE")
        auto_download = self.config.get("auto_download", True)

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

        stt_config = STTConfig(
            language=language,
        )

        # Provider erstellen
        # Provider erstellen (mit EventManager für Download-Events)
        self._provider = VoskSTTProvider(
            stt_config,
            models_dir=models_dir,
            model_name=model_name,
            auto_download=auto_download,
            event_manager=self.application.events,
        )

        # 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("Vosk STT Plugin: Extension registriert")

        # Event-Handler registrieren
        self._register_event_handlers()

        pinfo(f"Vosk STT Plugin: Geladen (Modell: {model_name})")

    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

            if isinstance(audio_data, str):
                audio_data = bytes.fromhex(audio_data)

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

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

                await em.emit("stt_completed", {
                    "text": result.text,
                    "confidence": result.confidence,
                    "language": result.language,
                    "provider": "vosk",
                    "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"Vosk 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"Vosk STT Fehler: {e}")

                await em.emit("stt_error", {
                    "error": str(e),
                    "provider": "vosk",
                    "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("Vosk STT Plugin: Entladen")

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


# Plugin-Export
Plugin = VoskSTTPlugin