cosmo-voice-assistant/cosmo/main.py

"""
Cosmo — локальный голосовой ассистент.
Стек: openWakeWord → RealtimeSTT (Whisper + Silero VAD) → smolagents + Ollama → RealtimeTTS (Silero)

Запуск:
    bash run.sh
    python cosmo/main.py
"""

import sys
import os
import argparse
import threading

import yaml

# Указываем pydub где искать ffmpeg (установлен через imageio-ffmpeg)
try:
    import imageio_ffmpeg
    from pydub import AudioSegment
    AudioSegment.converter = imageio_ffmpeg.get_ffmpeg_exe()
except Exception:
    pass
from loguru import logger

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from cosmo.wake_word import WakeWordDetector
from cosmo.transcriber import Transcriber
from cosmo.memory import Memory
from cosmo.agent import Agent
from cosmo.tts import TTS



def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def setup_logging(config: dict):
    log_cfg = config.get("logging", {})
    level = log_cfg.get("level", "INFO")
    log_file = log_cfg.get("file", "logs/cosmo.log")
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    logger.remove()
    logger.add(
        sys.stderr, level=level, colorize=True,
        format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | {message}",
    )
    logger.add(log_file, level="DEBUG", rotation="10 MB", retention="7 days", encoding="utf-8")


class Cosmo:
    def __init__(self, config: dict):
        self.config = config
        self.name = config["assistant"]["name"]
        self._running = False
        self._command_event = threading.Event()

        logger.info("Инициализирую модули...")

        self.tts = TTS(config)
        self.transcriber = Transcriber(config)
        self.memory = Memory()
        self.agent = Agent(config, self.memory)
        self.wake_word = WakeWordDetector(config, on_detected_callback=self._on_wake_word)

    def _on_wake_word(self):
        logger.info(f"=== {self.name} активирован! ===")
        self._command_event.set()

    def _process_command(self):
        # Синхронно — сначала говорим, потом слушаем (иначе TTS мешает записи)
        self.tts.say("Слушаю")

        # Partial results — печатаем в лог что слышим в реальном времени
        def on_partial(text):
            logger.debug(f"[partial] {text}")

        text = self.transcriber.record_and_transcribe(on_partial=on_partial)

        if not text.strip():
            self.tts.say_async("Не расслышал, попробуй ещё раз")
            return

        # Агент обрабатывает команду и возвращает ответ для TTS
        response = self.agent.run(text)
        self.tts.say(response)

    def run(self):
        self._running = True
        self.wake_word.start()

        logger.info("=" * 50)
        logger.info(f"  {self.name} запущен!")
        logger.info(f"  Скажи '{self.name}' чтобы активировать.")
        logger.info(f"  Ctrl+C для выхода.")
        logger.info("=" * 50)

        self.tts.say(f"{self.name} запущен")

        try:
            while self._running:
                triggered = self._command_event.wait(timeout=0.5)
                if triggered and self._running:
                    self._command_event.clear()
                    try:
                        self._process_command()
                    finally:
                        self.wake_word.resume()
        except KeyboardInterrupt:
            logger.info("Ctrl+C — завершаю...")
        finally:
            self.shutdown()

    def shutdown(self):
        logger.info("Завершаю работу...")
        self._running = False
        self.wake_word.stop()
        self.transcriber.shutdown()
        logger.info(f"{self.name} остановлен")


def main():
    parser = argparse.ArgumentParser(description="Cosmo — голосовой ассистент")
    parser.add_argument("--config", default="config/config.yaml")
    args = parser.parse_args()

    config_path = args.config
    if not os.path.isabs(config_path):
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        config_path = os.path.join(project_root, config_path)

    if not os.path.exists(config_path):
        print(f"Конфиг не найден: {config_path}")
        sys.exit(1)

    config = load_config(config_path)
    setup_logging(config)
    Cosmo(config).run()


if __name__ == "__main__":
    main()