Owner/src/owner_voice_pet/models.py

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Mapping


class PipelineState(str, Enum):
    IDLE = "idle"
    WAKE_LISTENING = "wake_listening"
    SPEECH_DETECTING = "speech_detecting"
    RECORDING = "recording"
    TRANSCRIBING = "transcribing"
    THINKING = "thinking"
    SPEAKING = "speaking"
    INTERRUPTED = "interrupted"
    ERROR_RECOVERING = "error_recovering"


class ErrorCode(str, Enum):
    AUDIO_INPUT_DEVICE_MISSING = "AUDIO_INPUT_DEVICE_MISSING"
    AUDIO_OUTPUT_DEVICE_MISSING = "AUDIO_OUTPUT_DEVICE_MISSING"
    AUDIO_PERMISSION_DENIED = "AUDIO_PERMISSION_DENIED"
    AUDIO_STREAM_UNDERRUN = "AUDIO_STREAM_UNDERRUN"
    AUDIO_FORMAT_UNSUPPORTED = "AUDIO_FORMAT_UNSUPPORTED"
    CONFIG_MISSING_VALUE = "CONFIG_MISSING_VALUE"
    WAKE_MODEL_MISSING = "WAKE_MODEL_MISSING"
    WAKE_MODEL_LOAD_FAILED = "WAKE_MODEL_LOAD_FAILED"
    WAKE_AUDIO_FORMAT_INVALID = "WAKE_AUDIO_FORMAT_INVALID"
    VAD_MODEL_LOAD_FAILED = "VAD_MODEL_LOAD_FAILED"
    VAD_TIMEOUT_NO_SPEECH = "VAD_TIMEOUT_NO_SPEECH"
    VAD_MAX_RECORDING_REACHED = "VAD_MAX_RECORDING_REACHED"
    STT_MODEL_MISSING = "STT_MODEL_MISSING"
    STT_TRANSCRIBE_FAILED = "STT_TRANSCRIBE_FAILED"
    STT_EMPTY_TRANSCRIPT = "STT_EMPTY_TRANSCRIPT"
    LLM_API_KEY_MISSING = "LLM_API_KEY_MISSING"
    LLM_REQUEST_TIMEOUT = "LLM_REQUEST_TIMEOUT"
    LLM_RATE_LIMITED = "LLM_RATE_LIMITED"
    LLM_NETWORK_ERROR = "LLM_NETWORK_ERROR"
    LLM_EMPTY_REPLY = "LLM_EMPTY_REPLY"
    TTS_MODEL_MISSING = "TTS_MODEL_MISSING"
    TTS_SYNTHESIS_FAILED = "TTS_SYNTHESIS_FAILED"
    TTS_EMPTY_AUDIO = "TTS_EMPTY_AUDIO"
    ASSET_MISSING = "ASSET_MISSING"
    VALIDATION_FAILED = "VALIDATION_FAILED"


@dataclass(slots=True)
class ProviderError(Exception):
    code: ErrorCode
    message: str
    retryable: bool
    provider: str
    stage: str

    def __str__(self) -> str:
        return f"{self.code.value} [{self.stage}/{self.provider}]: {self.message}"


@dataclass(frozen=True, slots=True)
class AudioFrame:
    pcm: bytes
    sample_rate: int
    channels: int
    timestamp_ms: int
    frame_id: int
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.sample_rate <= 0:
            raise ValueError("sample_rate must be positive")
        if self.channels <= 0:
            raise ValueError("channels must be positive")
        if self.timestamp_ms < 0:
            raise ValueError("timestamp_ms must be non-negative")
        if self.frame_id < 0:
            raise ValueError("frame_id must be non-negative")


@dataclass(frozen=True, slots=True)
class AudioSegment:
    pcm: bytes
    sample_rate: int
    channels: int
    start_time_ms: int
    end_time_ms: int
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.sample_rate <= 0:
            raise ValueError("sample_rate must be positive")
        if self.channels <= 0:
            raise ValueError("channels must be positive")
        if self.start_time_ms < 0:
            raise ValueError("start_time_ms must be non-negative")
        if self.end_time_ms < self.start_time_ms:
            raise ValueError("end_time_ms must be >= start_time_ms")

    @property
    def duration_ms(self) -> int:
        return self.end_time_ms - self.start_time_ms


@dataclass(frozen=True, slots=True)
class WakeEvent:
    keyword: str
    confidence: float
    timestamp_ms: int


@dataclass(frozen=True, slots=True)
class VadResult:
    is_speech: bool
    confidence: float
    speech_ms: int
    silence_ms: int
    end_reason: str | None = None


@dataclass(frozen=True, slots=True)
class Transcript:
    text: str
    language: str
    confidence: float | None
    duration_ms: int
    provider: str
    raw_metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def normalized_text(self) -> str:
        return self.text.strip()


@dataclass(frozen=True, slots=True)
class Message:
    role: str
    content: str
    created_at: float


@dataclass(frozen=True, slots=True)
class ReplyDelta:
    text_delta: str
    is_sentence_boundary: bool = False
    finish_reason: str | None = None


@dataclass(frozen=True, slots=True)
class PlaybackResult:
    played: bool
    duration_ms: int
    error: ProviderError | None = None


@dataclass(frozen=True, slots=True)
class TransportHealth:
    input_available: bool
    output_available: bool
    message: str = ""