atom-audio-engine 0.1.0__py3-none-any.whl

asr/__init__.py

@@ -0,0 +1,45 @@
+"""ASR (Speech-to-Text) providers."""
+
+from core.config import ASRConfig
+
+from .base import BaseASR
+from .deepgram import DeepgramASR
+from .cartesia import CartesiaASR
+
+__all__ = ["BaseASR", "DeepgramASR", "CartesiaASR", "get_asr_from_config"]
+
+
+def get_asr_from_config(config: ASRConfig) -> BaseASR:
+    """
+    Instantiate ASR provider from config.
+
+    Args:
+        config: ASRConfig object with provider name and settings
+
+    Returns:
+        Initialized BaseASR provider instance
+
+    Raises:
+        ValueError: If provider name is not recognized
+    """
+    provider_name = config.provider.lower()
+
+    if provider_name == "deepgram":
+        return DeepgramASR(
+            api_key=config.api_key,
+            model=config.model or "nova-2",
+            language=config.language,
+            **config.extra,
+        )
+    elif provider_name == "cartesia":
+        return CartesiaASR(
+            api_key=config.api_key,
+            model=config.model or "ink-whisper",
+            language=config.language,
+            **config.extra,
+        )
+    else:
+        raise ValueError(
+            f"Unknown ASR provider: {config.provider}. "
+            f"Supported: deepgram, cartesia"
+        )

asr/base.py

@@ -0,0 +1,89 @@
+"""Abstract base class for ASR (Speech-to-Text) providers."""
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterator, Optional
+
+from core.types import AudioChunk, TranscriptChunk
+
+
+class BaseASR(ABC):
+    """
+    Abstract base class for Speech-to-Text providers.
+
+    All ASR implementations must inherit from this class and implement
+    the required methods for both batch and streaming transcription.
+    """
+
+    def __init__(self, api_key: Optional[str] = None, **kwargs):
+        """
+        Initialize the ASR provider.
+
+        Args:
+            api_key: API key for the provider (if required)
+            **kwargs: Additional provider-specific configuration
+        """
+        self.api_key = api_key
+        self.config = kwargs
+
+    @abstractmethod
+    async def transcribe(self, audio: bytes, sample_rate: int = 16000) -> str:
+        """
+        Transcribe a complete audio buffer to text.
+
+        Args:
+            audio: Raw audio bytes (PCM format expected)
+            sample_rate: Sample rate of the audio in Hz
+
+        Returns:
+            Transcribed text string
+        """
+        pass
+
+    @abstractmethod
+    async def transcribe_stream(
+        self, audio_stream: AsyncIterator[AudioChunk]
+    ) -> AsyncIterator[TranscriptChunk]:
+        """
+        Transcribe streaming audio in real-time.
+
+        Args:
+            audio_stream: Async iterator yielding AudioChunk objects
+
+        Yields:
+            TranscriptChunk objects with partial and final transcriptions
+        """
+        pass
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.disconnect()
+
+    async def connect(self):
+        """
+        Establish connection to the ASR service.
+        Override in subclasses if needed.
+        """
+        pass
+
+    async def disconnect(self):
+        """
+        Close connection to the ASR service.
+        Override in subclasses if needed.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the name of this ASR provider."""
+        pass
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Whether this provider supports real-time streaming."""
+        return True

asr/cartesia.py

@@ -0,0 +1,356 @@
+"""Cartesia API implementation for ASR (Speech-to-Text) via WebSocket."""
+
+import asyncio
+import json
+import logging
+from typing import AsyncIterator, Optional
+from urllib.parse import quote
+
+import websockets
+
+from core.types import AudioChunk, TranscriptChunk
+from .base import BaseASR
+
+logger = logging.getLogger(__name__)
+
+# Cartesia API version (required header)
+CARTESIA_VERSION = "2025-04-16"
+
+
+class CartesiaASR(BaseASR):
+    """
+    Cartesia API client for speech-to-text transcription via WebSocket.
+
+    Supports both batch transcription and real-time streaming.
+    Uses Cartesia's Whisper model (ink-whisper) for high-accuracy transcription.
+
+    Approach:
+    1. Batch mode: collect audio, send via WebSocket, wait for final result
+    2. Streaming mode: send audio chunks as they arrive, yield results immediately
+    3. Background receive task queues responses from server
+    4. VAD (Voice Activity Detection) configurable via min_volume and max_silence_duration_secs
+
+    Example:
+        asr = CartesiaASR(api_key="sk_...")
+
+        # Batch transcription
+        text = await asr.transcribe(audio_bytes)
+
+        # Streaming transcription
+        async for chunk in asr.transcribe_stream(audio_stream):
+            print(chunk.text, end="", flush=True)
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "ink-whisper",
+        language: str = "en",
+        encoding: str = "pcm_s16le",
+        sample_rate: int = 16000,
+        min_volume: float = 0.0,
+        max_silence_duration_secs: float = 30.0,
+        **kwargs,
+    ):
+        """
+        Initialize Cartesia ASR provider.
+
+        Args:
+            api_key: Cartesia API key
+            model: Model to use (default: ink-whisper)
+            language: Language code in ISO-639-1 format (default: en)
+            encoding: Audio encoding format (default: pcm_s16le)
+            sample_rate: Sample rate in Hz (default: 16000)
+            min_volume: VAD threshold 0.0-1.0, higher = more aggressive (default: 0.0)
+            max_silence_duration_secs: Max silence before endpointing (default: 30.0)
+            **kwargs: Additional config (stored in self.config)
+        """
+        super().__init__(api_key=api_key, **kwargs)
+        self.model = model
+        self.language = language
+        self.encoding = encoding
+        self.sample_rate = sample_rate
+        self.min_volume = min_volume
+        self.max_silence_duration_secs = max_silence_duration_secs
+
+        self.websocket = None
+        self._receive_task: Optional[asyncio.Task] = None
+        self._response_queue: asyncio.Queue = asyncio.Queue()
+
+    @property
+    def name(self) -> str:
+        """Return provider name."""
+        return "cartesia"
+
+    async def connect(self):
+        """
+        Initialize WebSocket connection to Cartesia STT endpoint.
+
+        Approach:
+        1. Construct WebSocket URL with parameters (model, language, encoding, sample_rate, VAD)
+        2. Connect to wss://api.cartesia.ai/stt/websocket
+        3. Launch background receive task to collect server responses
+        4. Log initialization status
+
+        Rationale: Lazy connection on first transcription; background task ensures
+        responses are queued even if caller is temporarily blocked.
+        """
+        if self.websocket:
+            return
+
+        try:
+            if not self.api_key:
+                # Try to get from environment
+                import os
+
+                self.api_key = os.getenv("CARTESIA_API_KEY") or os.getenv("ASR_API_KEY")
+
+            if not self.api_key:
+                raise ValueError("Cartesia API key not provided")
+
+            # Construct WebSocket URL with properly encoded parameters
+            # API key must be URL-encoded to handle special characters
+            url = (
+                f"wss://api.cartesia.ai/stt/websocket?"
+                f"model={quote(str(self.model))}"
+                f"&language={quote(str(self.language))}"
+                f"&encoding={quote(str(self.encoding))}"
+                f"&sample_rate={quote(str(self.sample_rate))}"
+                f"&min_volume={quote(str(self.min_volume))}"
+                f"&max_silence_duration_secs={quote(str(self.max_silence_duration_secs))}"
+                f"&api_key={quote(str(self.api_key))}"
+            )
+            logger.debug(f"Cartesia WebSocket URL: {url}")
+
+            # Connect to WebSocket with required Cartesia-Version header
+            try:
+                # Build URL with all parameters
+                # Add Cartesia-Version header via subprotocol or connection params
+                self.websocket = await asyncio.wait_for(
+                    websockets.connect(
+                        url, additional_headers=[("Cartesia-Version", CARTESIA_VERSION)]
+                    ),
+                    timeout=30.0,  # Increase timeout to 30s for initial connection
+                )
+                logger.debug("Cartesia WebSocket connected")
+            except asyncio.TimeoutError:
+                logger.error(f"WebSocket connection timeout to {url}")
+                raise TimeoutError(
+                    "Failed to connect to Cartesia WebSocket within 30s timeout"
+                )
+
+            # Start background receive task
+            self._receive_task = asyncio.create_task(self._receive_loop())
+
+        except Exception as e:
+            logger.error(f"Failed to initialize Cartesia WebSocket: {e}")
+            raise
+
+    async def disconnect(self):
+        """Close WebSocket connection and cleanup."""
+        try:
+            if self._receive_task:
+                self._receive_task.cancel()
+                try:
+                    await self._receive_task
+                except asyncio.CancelledError:
+                    pass
+
+            if self.websocket:
+                await self.websocket.close()
+                logger.debug("Cartesia WebSocket closed")
+
+        except Exception as e:
+            logger.error(f"Error disconnecting Cartesia: {e}")
+
+    async def _receive_loop(self):
+        """
+        Background task: continuously receive messages from WebSocket.
+
+        Parses JSON responses and queues them for retrieval by transcribe methods.
+        Handles: transcript, flush_done, done, error message types.
+        """
+        try:
+            if not self.websocket:
+                return
+
+            async for message in self.websocket:
+                try:
+                    # Parse JSON response
+                    response = json.loads(message)
+                    await self._response_queue.put(response)
+
+                except json.JSONDecodeError as e:
+                    logger.error(f"Failed to parse Cartesia response: {e}")
+                except Exception as e:
+                    logger.error(f"Error in receive loop: {e}")
+
+        except asyncio.CancelledError:
+            logger.debug("Receive loop cancelled")
+        except Exception as e:
+            logger.error(f"Unexpected error in receive loop: {e}")
+
+    async def transcribe(self, audio: bytes, sample_rate: int = 16000) -> str:
+        """
+        Transcribe complete audio buffer to text.
+
+        Approach:
+        1. Initialize WebSocket if needed
+        2. Send audio in chunks (100ms intervals)
+        3. Send 'done' command to finalize
+        4. Collect all responses until 'done' received
+        5. Extract and return transcript text
+
+        Rationale: Batch mode for complete audio files; simple sequential flow.
+
+        Args:
+            audio: Raw PCM audio bytes
+            sample_rate: Sample rate in Hz (default 16000)
+
+        Returns:
+            Transcribed text
+        """
+        if not self.websocket:
+            await self.connect()
+
+        try:
+            logger.debug(f"Transcribing {len(audio)} bytes at {sample_rate}Hz")
+
+            # Send audio in chunks (100ms intervals at 16kHz = 3200 bytes)
+            chunk_size = int(self.sample_rate * 0.1 * 2)  # 100ms in bytes
+            offset = 0
+
+            while offset < len(audio):
+                chunk = audio[offset : offset + chunk_size]
+                await self.websocket.send(chunk)
+                offset += chunk_size
+
+            # Send 'done' command to finalize
+            await self.websocket.send("done")
+
+            # Collect responses until 'done' received
+            transcript_parts = []
+            while True:
+                try:
+                    response = await asyncio.wait_for(
+                        self._response_queue.get(), timeout=10.0
+                    )
+
+                    if response.get("type") == "transcript":
+                        text = response.get("text", "")
+                        if text:
+                            transcript_parts.append(text)
+
+                    elif response.get("type") == "done":
+                        break
+
+                    elif response.get("type") == "error":
+                        error_msg = response.get("error", "Unknown error")
+                        raise RuntimeError(f"Cartesia error: {error_msg}")
+
+                except asyncio.TimeoutError:
+                    logger.warning("Timeout waiting for Cartesia response")
+                    break
+
+            return "".join(transcript_parts)
+
+        except Exception as e:
+            logger.error(f"Cartesia transcription error: {e}")
+            raise
+
+    async def transcribe_stream(
+        self, audio_stream: AsyncIterator[AudioChunk]
+    ) -> AsyncIterator[TranscriptChunk]:
+        """
+        Transcribe streaming audio in real-time.
+
+        Approach:
+        1. Initialize WebSocket if needed
+        2. For each audio chunk from stream:
+           - Send binary audio via WebSocket
+           - Check response queue for server responses (non-blocking)
+           - Yield TranscriptChunk for each response
+        3. On final audio chunk (is_final=True), send 'done' command
+        4. Continue yielding responses until 'done' received
+        5. Signal stream end
+
+        Rationale: Streaming yields results immediately; low latency;
+        background task queues responses so we don't block on receives.
+
+        Args:
+            audio_stream: Async iterator yielding AudioChunk objects
+
+        Yields:
+            TranscriptChunk objects with partial and final transcriptions
+        """
+        if not self.websocket:
+            await self.connect()
+
+        try:
+            done_sent = False
+
+            async for audio_chunk in audio_stream:
+                # Send audio via WebSocket
+                await self.websocket.send(audio_chunk.data)
+
+                # Try to get responses (non-blocking)
+                while not self._response_queue.empty():
+                    response = self._response_queue.get_nowait()
+
+                    if response.get("type") == "transcript":
+                        text = response.get("text", "")
+                        is_final = response.get("is_final", False)
+                        if text:
+                            yield TranscriptChunk(
+                                text=text,
+                                confidence=None,  # Cartesia doesn't return confidence
+                                is_final=is_final,
+                            )
+
+                    elif response.get("type") == "error":
+                        error_msg = response.get("error", "Unknown error")
+                        logger.error(f"Cartesia error: {error_msg}")
+
+                # If this is the final audio chunk, send 'done' command
+                if audio_chunk.is_final and not done_sent:
+                    await self.websocket.send("done")
+                    done_sent = True
+
+            # Continue collecting responses until 'done' received
+            if done_sent:
+                while True:
+                    try:
+                        response = await asyncio.wait_for(
+                            self._response_queue.get(), timeout=5.0
+                        )
+
+                        if response.get("type") == "transcript":
+                            text = response.get("text", "")
+                            is_final = response.get("is_final", False)
+                            if text:
+                                yield TranscriptChunk(
+                                    text=text,
+                                    confidence=None,
+                                    is_final=is_final,
+                                )
+
+                        elif response.get("type") == "done":
+                            # Yield final chunk to signal stream end
+                            yield TranscriptChunk(
+                                text="",
+                                confidence=None,
+                                is_final=True,
+                            )
+                            break
+
+                        elif response.get("type") == "error":
+                            error_msg = response.get("error", "Unknown error")
+                            logger.error(f"Cartesia error: {error_msg}")
+
+                    except asyncio.TimeoutError:
+                        logger.warning("Timeout waiting for Cartesia final response")
+                        break
+
+        except Exception as e:
+            logger.error(f"Cartesia streaming transcription error: {e}")
+            raise

asr/deepgram.py

@@ -0,0 +1,196 @@
+"""Deepgram API implementation for ASR (Speech-to-Text)."""
+
+import logging
+from typing import AsyncIterator, Optional
+
+from deepgram import DeepgramClient
+
+from core.types import AudioChunk, TranscriptChunk
+from .base import BaseASR
+
+logger = logging.getLogger(__name__)
+
+
+class DeepgramASR(BaseASR):
+    """
+    Deepgram API client for speech-to-text transcription.
+
+    Supports both batch transcription and real-time streaming.
+    Outputs high-accuracy transcripts using Deepgram's Nova-2 model by default.
+
+    Example:
+        asr = DeepgramASR(api_key="dg_...")
+
+        # Batch transcription
+        text = await asr.transcribe(audio_bytes)
+
+        # Streaming transcription
+        async for chunk in asr.transcribe_stream(audio_stream):
+            print(chunk.text, end="", flush=True)
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "nova-2",
+        language: str = "en",
+        **kwargs,
+    ):
+        """
+        Initialize Deepgram ASR provider.
+
+        Args:
+            api_key: Deepgram API key
+            model: Model to use (e.g., "nova-2", "nova", "enhanced")
+            language: Language code (e.g., "en", "es", "fr")
+            **kwargs: Additional config (stored in self.config)
+        """
+        super().__init__(api_key=api_key, **kwargs)
+        self.model = model
+        self.language = language
+        self.client = None
+
+    @property
+    def name(self) -> str:
+        """Return provider name."""
+        return "deepgram"
+
+    async def connect(self):
+        """
+        Initialize Deepgram client.
+
+        Approach:
+        1. Create client with API key (from param or env var DEEPGRAM_API_KEY)
+        2. Log initialization status
+
+        Rationale: Client reuse for multiple transcription requests.
+        """
+        try:
+            if self.api_key:
+                self.client = DeepgramClient(api_key=self.api_key)
+            else:
+                # Fallback to env var DEEPGRAM_API_KEY
+                self.client = DeepgramClient()
+
+            logger.debug("Deepgram client initialized")
+        except Exception as e:
+            logger.error(f"Failed to initialize Deepgram client: {e}")
+            raise
+
+    async def disconnect(self):
+        """Close Deepgram client connection."""
+        if self.client:
+            try:
+                pass
+            except Exception as e:
+                logger.error(f"Error disconnecting Deepgram: {e}")
+
+    async def transcribe(self, audio: bytes, sample_rate: int = 16000) -> str:
+        """
+        Transcribe complete audio buffer to text.
+
+        Approach:
+        1. Initialize client if needed
+        2. Send audio to Deepgram prerecorded API
+        3. Extract and return transcript text
+
+        Rationale: Batch mode for complete audio files with standard latency.
+
+        Args:
+            audio: Raw PCM audio bytes
+            sample_rate: Sample rate in Hz (default 16000)
+
+        Returns:
+            Transcribed text
+        """
+        if not self.client:
+            await self.connect()
+
+        try:
+            logger.debug(f"Transcribing {len(audio)} bytes at {sample_rate}Hz")
+
+            # Call Deepgram API - using synchronous client
+            response = self.client.listen.prerecorded.v(
+                {
+                    "model": self.model,
+                    "language": self.language,
+                    "encoding": "linear16",
+                    "sample_rate": sample_rate,
+                }
+            ).transcribe_file({"buffer": audio})
+
+            # Extract transcript
+            if response and response.results:
+                if response.results.channels:
+                    channel = response.results.channels[0]
+                    if channel.alternatives:
+                        transcript = channel.alternatives[0].transcript
+                        logger.debug(f"Transcribed to: {transcript[:100]}...")
+                        return transcript
+
+            return ""
+
+        except Exception as e:
+            logger.error(f"Deepgram transcription error: {e}")
+            raise
+
+    async def transcribe_stream(
+        self, audio_stream: AsyncIterator[AudioChunk]
+    ) -> AsyncIterator[TranscriptChunk]:
+        """
+        Transcribe streaming audio in real-time.
+
+        Approach:
+        1. Collect audio chunks from stream until is_final flag
+        2. Send buffered audio to Deepgram API
+        3. Yield transcription results
+
+        Rationale: Simple buffering approach; Deepgram SDK doesn't expose
+        native streaming in current version, so we batch on is_final signals.
+
+        Args:
+            audio_stream: Async iterator yielding AudioChunk objects
+
+        Yields:
+            TranscriptChunk objects with partial and final transcriptions
+        """
+        if not self.client:
+            await self.connect()
+
+        try:
+            buffer = bytearray()
+
+            async for chunk in audio_stream:
+                buffer.extend(chunk.data)
+
+                if chunk.is_final:
+                    # Transcribe accumulated buffer
+                    if buffer:
+                        response = self.client.listen.prerecorded.v(
+                            {
+                                "model": self.model,
+                                "language": self.language,
+                                "encoding": "linear16",
+                                "sample_rate": 16000,
+                            }
+                        ).transcribe_file({"buffer": bytes(buffer)})
+
+                        if response and response.results:
+                            if response.results.channels:
+                                channel = response.results.channels[0]
+                                if channel.alternatives:
+                                    transcript = channel.alternatives[0].transcript
+
+                                    yield TranscriptChunk(
+                                        text=transcript,
+                                        is_final=True,
+                                        confidence=getattr(
+                                            channel.alternatives[0], "confidence", None
+                                        ),
+                                    )
+
+                        buffer = bytearray()
+
+        except Exception as e:
+            logger.error(f"Deepgram streaming error: {e}")
+            raise