STT-Assistant/src/stt_daemon/config.py

"""Configuration management for STT daemon."""

from dataclasses import dataclass, field
import tomllib
from pathlib import Path


@dataclass
class AudioConfig:
    """Audio capture configuration."""
    backend: str = "parec"
    device: str = "VirtualMicSink.monitor"
    sample_rate: int = 16000
    channels: int = 1
    frame_ms: int = 30
    gain: float = 1.0


@dataclass
class VadConfig:
    """Voice Activity Detection configuration."""
    aggressiveness: int = 2
    pre_roll_ms: int = 400
    speech_start_ms: int = 90
    silence_ms: int = 700
    min_utterance_ms: int = 300
    max_utterance_s: float = 30.0
    chunk_pause_ms: int = 350
    soft_max_s: float = 10.0
    rms_floor: float = 0.008
    trailing_trim_ms: int = 200


@dataclass
class WhisperConfig:
    """Whisper transcription configuration."""
    model: str = "small"
    device: str = "cpu"
    compute_type: str = "int8"
    language: str | None = None
    beam_size: int = 1
    vad_filter: bool = False
    no_speech_prob_max: float = 0.6
    avg_logprob_min: float = -1.0


@dataclass
class OutputConfig:
    """Output configuration."""
    stdout: bool = True
    jsonl_dir: str | None = "transcripts"
    save_audio: bool = False
    audio_dir: str = "audio"


@dataclass
class DaemonConfig:
    """Daemon runtime configuration."""
    queue_size: int = 50
    backlog_warn_at: int = 5


@dataclass
class EventsConfig:
    """Publish-only event bus configuration."""
    socket: str = ""
    heartbeat_ms: int = 1000
    ring_size: int = 200


@dataclass
class AssistantConfig:
    """Assistant (controller) configuration."""
    control_socket: str = ""
    stall_timeout_ms: int = 3000
    worker_stall_ms: int = 15000
    command_arm_timeout_s: float = 10.0
    chat_idle_timeout_s: float = 300.0
    chat_exit_phrases: list[str] = field(
        default_factory=lambda: ["end conversation", "that's all", "stop listening"]
    )
    require_viewer: bool = True
    turn_gap_s: float = 1.5


@dataclass
class AuditConfig:
    """Action audit configuration (daily JSONL, like transcripts)."""
    enabled: bool = True
    dir: str = "audit"


@dataclass
class LlmConfig:
    """llama.cpp (OpenAI-compatible) configuration."""
    enabled: bool = True
    endpoint: str = "http://localhost:9931/v1/chat/completions"
    temperature: float = 0.7
    max_tokens: int = 1024
    timeout_s: float = 120.0
    cache_prompt: bool = True
    max_prompt_tokens: int = 64000
    trim_block_turns: int = 4
    max_tool_iterations: int = 3
    system_prompt_chat: str = (
        "You are a concise voice assistant. Reply in 1-3 sentences."
    )
    system_prompt_command: str = (
        "You translate a spoken request into a single tool call when one applies."
    )


@dataclass
class SinkConfig:
    """Virtual audio sink configuration."""
    sink_name: str = "VirtualMicSink"
    remap_source_name: str = "VirtualMic"
    loopback_latency_msec: int = 1
    monitor_output: bool = False


@dataclass
class Config:
    """Complete daemon configuration."""
    audio: AudioConfig = field(default_factory=AudioConfig)
    vad: VadConfig = field(default_factory=VadConfig)
    whisper: WhisperConfig = field(default_factory=WhisperConfig)
    output: OutputConfig = field(default_factory=OutputConfig)
    sink: SinkConfig = field(default_factory=SinkConfig)
    daemon: DaemonConfig = field(default_factory=DaemonConfig)
    events: EventsConfig = field(default_factory=EventsConfig)
    assistant: AssistantConfig = field(default_factory=AssistantConfig)
    audit: AuditConfig = field(default_factory=AuditConfig)
    llm: LlmConfig = field(default_factory=LlmConfig)


def load_config(path: str | None = None) -> Config:
    """Load configuration from TOML file, with defaults.

    Args:
        path: Path to config.toml. If None, looks for config.toml in cwd,
              then uses all defaults.

    Returns:
        Config with loaded values overlayed on defaults.
    """
    cfg = Config()

    # Try to find and load config file
    config_path = None
    if path:
        config_path = Path(path)
    else:
        default_path = Path("config.toml")
        if default_path.exists():
            config_path = default_path

    if config_path and config_path.exists():
        with open(config_path, "rb") as f:
            data = tomllib.load(f)

        # Audio section
        if "audio" in data:
            audio_data = data["audio"]
            cfg.audio = AudioConfig(
                backend=audio_data.get("backend", cfg.audio.backend),
                device=audio_data.get("device", cfg.audio.device),
                sample_rate=audio_data.get("sample_rate", cfg.audio.sample_rate),
                channels=audio_data.get("channels", cfg.audio.channels),
                frame_ms=audio_data.get("frame_ms", cfg.audio.frame_ms),
                gain=audio_data.get("gain", cfg.audio.gain),
            )

        # VAD section
        if "vad" in data:
            vad_data = data["vad"]
            cfg.vad = VadConfig(
                aggressiveness=vad_data.get("aggressiveness", cfg.vad.aggressiveness),
                pre_roll_ms=vad_data.get("pre_roll_ms", cfg.vad.pre_roll_ms),
                speech_start_ms=vad_data.get("speech_start_ms", cfg.vad.speech_start_ms),
                silence_ms=vad_data.get("silence_ms", cfg.vad.silence_ms),
                min_utterance_ms=vad_data.get("min_utterance_ms", cfg.vad.min_utterance_ms),
                max_utterance_s=vad_data.get("max_utterance_s", cfg.vad.max_utterance_s),
                chunk_pause_ms=vad_data.get("chunk_pause_ms", cfg.vad.chunk_pause_ms),
                soft_max_s=vad_data.get("soft_max_s", cfg.vad.soft_max_s),
                rms_floor=vad_data.get("rms_floor", cfg.vad.rms_floor),
                trailing_trim_ms=vad_data.get("trailing_trim_ms", cfg.vad.trailing_trim_ms),
            )

        # Whisper section
        if "whisper" in data:
            whisper_data = data["whisper"]
            language = whisper_data.get("language", cfg.whisper.language)
            # Normalize language: "", "auto", "none" -> None
            if language in ("", "auto", "none"):
                language = None
            cfg.whisper = WhisperConfig(
                model=whisper_data.get("model", cfg.whisper.model),
                device=whisper_data.get("device", cfg.whisper.device),
                compute_type=whisper_data.get("compute_type", cfg.whisper.compute_type),
                language=language,
                beam_size=whisper_data.get("beam_size", cfg.whisper.beam_size),
                vad_filter=whisper_data.get("vad_filter", cfg.whisper.vad_filter),
                no_speech_prob_max=whisper_data.get(
                    "no_speech_prob_max", cfg.whisper.no_speech_prob_max
                ),
                avg_logprob_min=whisper_data.get(
                    "avg_logprob_min", cfg.whisper.avg_logprob_min
                ),
            )

        # Output section
        if "output" in data:
            output_data = data["output"]
            cfg.output = OutputConfig(
                stdout=output_data.get("stdout", cfg.output.stdout),
                jsonl_dir=output_data.get("jsonl_dir", cfg.output.jsonl_dir),
                save_audio=output_data.get("save_audio", cfg.output.save_audio),
                audio_dir=output_data.get("audio_dir", cfg.output.audio_dir),
            )

        # Sink section
        if "sink" in data:
            sink_data = data["sink"]
            cfg.sink = SinkConfig(
                sink_name=sink_data.get("sink_name", cfg.sink.sink_name),
                remap_source_name=sink_data.get(
                    "remap_source_name", cfg.sink.remap_source_name
                ),
                loopback_latency_msec=sink_data.get(
                    "loopback_latency_msec", cfg.sink.loopback_latency_msec
                ),
                monitor_output=sink_data.get("monitor_output", cfg.sink.monitor_output),
            )

        # Daemon section
        if "daemon" in data:
            daemon_data = data["daemon"]
            cfg.daemon = DaemonConfig(
                queue_size=daemon_data.get("queue_size", cfg.daemon.queue_size),
                backlog_warn_at=daemon_data.get(
                    "backlog_warn_at", cfg.daemon.backlog_warn_at
                ),
            )

        # Events section
        if "events" in data:
            events_data = data["events"]
            cfg.events = EventsConfig(
                socket=events_data.get("socket", cfg.events.socket),
                heartbeat_ms=events_data.get("heartbeat_ms", cfg.events.heartbeat_ms),
                ring_size=events_data.get("ring_size", cfg.events.ring_size),
            )

        # Assistant section
        if "assistant" in data:
            assistant_data = data["assistant"]
            cfg.assistant = AssistantConfig(
                control_socket=assistant_data.get(
                    "control_socket", cfg.assistant.control_socket
                ),
                stall_timeout_ms=assistant_data.get(
                    "stall_timeout_ms", cfg.assistant.stall_timeout_ms
                ),
                worker_stall_ms=assistant_data.get(
                    "worker_stall_ms", cfg.assistant.worker_stall_ms
                ),
                command_arm_timeout_s=assistant_data.get(
                    "command_arm_timeout_s", cfg.assistant.command_arm_timeout_s
                ),
                chat_idle_timeout_s=assistant_data.get(
                    "chat_idle_timeout_s", cfg.assistant.chat_idle_timeout_s
                ),
                chat_exit_phrases=list(
                    assistant_data.get(
                        "chat_exit_phrases", cfg.assistant.chat_exit_phrases
                    )
                ),
                require_viewer=assistant_data.get(
                    "require_viewer", cfg.assistant.require_viewer
                ),
                turn_gap_s=assistant_data.get("turn_gap_s", cfg.assistant.turn_gap_s),
            )

        # Audit section
        if "audit" in data:
            audit_data = data["audit"]
            cfg.audit = AuditConfig(
                enabled=audit_data.get("enabled", cfg.audit.enabled),
                dir=audit_data.get("dir", cfg.audit.dir),
            )

        # LLM section
        if "llm" in data:
            llm_data = data["llm"]
            cfg.llm = LlmConfig(
                enabled=llm_data.get("enabled", cfg.llm.enabled),
                endpoint=llm_data.get("endpoint", cfg.llm.endpoint),
                temperature=llm_data.get("temperature", cfg.llm.temperature),
                max_tokens=llm_data.get("max_tokens", cfg.llm.max_tokens),
                timeout_s=llm_data.get("timeout_s", cfg.llm.timeout_s),
                cache_prompt=llm_data.get("cache_prompt", cfg.llm.cache_prompt),
                max_prompt_tokens=llm_data.get(
                    "max_prompt_tokens", cfg.llm.max_prompt_tokens
                ),
                trim_block_turns=llm_data.get(
                    "trim_block_turns", cfg.llm.trim_block_turns
                ),
                max_tool_iterations=llm_data.get(
                    "max_tool_iterations", cfg.llm.max_tool_iterations
                ),
                system_prompt_chat=llm_data.get(
                    "system_prompt_chat", cfg.llm.system_prompt_chat
                ),
                system_prompt_command=llm_data.get(
                    "system_prompt_command", cfg.llm.system_prompt_command
                ),
            )

    return cfg