atom-audio-engine 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

audio_engine/streaming/websocket_server.py

@@ -0,0 +1,333 @@
+"""WebSocket server for real-time audio streaming."""
+
+import asyncio
+import json
+import logging
+from typing import Optional, Callable, Any
+
+import websockets
+
+from ..core.pipeline import Pipeline
+from ..core.types import AudioChunk, AudioFormat
+from ..core.config import AudioEngineConfig
+
+logger = logging.getLogger(__name__)
+
+# Type alias for WebSocket connection
+WebSocketServerProtocol = Any
+
+
+class WebSocketServer:
+    """
+    WebSocket server for real-time audio-to-audio streaming.
+
+    Protocol:
+        Client sends:
+        - Binary messages: Raw audio chunks (PCM 16-bit, 16kHz mono)
+        - JSON messages: Control commands {"type": "end_of_speech"} or {"type": "reset"}
+
+        Server sends:
+        - Binary messages: Response audio chunks
+        - JSON messages: Events {"type": "transcript", "text": "..."} etc.
+
+    Example:
+        ```python
+        server = WebSocketServer(
+            pipeline=pipeline,
+            host="0.0.0.0",
+            port=8765
+        )
+        await server.start()
+        ```
+    """
+
+    def __init__(
+        self,
+        pipeline: Pipeline,
+        host: str = "0.0.0.0",
+        port: int = 8765,
+        input_sample_rate: int = 16000,
+        on_connect: Optional[Callable[[str], Any]] = None,
+        on_disconnect: Optional[Callable[[str], Any]] = None,
+    ):
+        """
+        Initialize the WebSocket server.
+
+        Args:
+            pipeline: Configured Pipeline instance
+            host: Host to bind to
+            port: Port to listen on
+            input_sample_rate: Expected sample rate of input audio
+            on_connect: Callback when client connects
+            on_disconnect: Callback when client disconnects
+        """
+        if websockets is None:
+            raise ImportError("websockets package required. Install with: pip install websockets")
+
+        self.pipeline = pipeline
+        self.host = host
+        self.port = port
+        self.input_sample_rate = input_sample_rate
+        self.on_connect = on_connect
+        self.on_disconnect = on_disconnect
+
+        self._server = None
+        self._clients: dict[str, WebSocketServerProtocol] = {}
+
+    async def start(self):
+        """Start the WebSocket server."""
+        await self.pipeline.connect()
+
+        self._server = await websockets.serve(
+            self._handle_client,
+            self.host,
+            self.port,
+        )
+
+        logger.info(f"WebSocket server started on ws://{self.host}:{self.port}")
+
+    async def stop(self):
+        """Stop the WebSocket server."""
+        if self._server:
+            self._server.close()
+            await self._server.wait_closed()
+            self._server = None
+
+        await self.pipeline.disconnect()
+        logger.info("WebSocket server stopped")
+
+    async def _handle_client(self, websocket: WebSocketServerProtocol):
+        """Handle a single client connection."""
+        client_id = str(id(websocket))
+        self._clients[client_id] = websocket
+
+        logger.info(f"Client connected: {client_id}")
+        if self.on_connect:
+            self.on_connect(client_id)
+
+        # Send welcome message
+        await websocket.send(
+            json.dumps(
+                {
+                    "type": "connected",
+                    "client_id": client_id,
+                    "providers": self.pipeline.providers,
+                }
+            )
+        )
+
+        try:
+            await self._process_client_stream(websocket, client_id)
+        except websockets.exceptions.ConnectionClosed:
+            logger.info(f"Client disconnected: {client_id}")
+        except Exception as e:
+            logger.error(f"Error handling client {client_id}: {e}")
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "error",
+                        "message": str(e),
+                    }
+                )
+            )
+        finally:
+            del self._clients[client_id]
+            if self.on_disconnect:
+                self.on_disconnect(client_id)
+
+    async def _process_client_stream(self, websocket: WebSocketServerProtocol, client_id: str):
+        """Process streaming audio from a client."""
+        audio_queue: asyncio.Queue[AudioChunk] = asyncio.Queue()
+        end_of_speech = asyncio.Event()
+
+        async def audio_stream():
+            """Yield audio chunks from the queue."""
+            while True:
+                if end_of_speech.is_set() and audio_queue.empty():
+                    break
+                try:
+                    chunk = await asyncio.wait_for(audio_queue.get(), timeout=0.1)
+                    yield chunk
+                    if chunk.is_final:
+                        break
+                except asyncio.TimeoutError:
+                    if end_of_speech.is_set():
+                        break
+                    continue
+
+        async def receive_audio():
+            """Receive audio from WebSocket and queue it."""
+            async for message in websocket:
+                if isinstance(message, bytes):
+                    # Binary audio data
+                    chunk = AudioChunk(
+                        data=message,
+                        sample_rate=self.input_sample_rate,
+                        format=AudioFormat.PCM_16K,
+                    )
+                    await audio_queue.put(chunk)
+
+                elif isinstance(message, str):
+                    # JSON control message
+                    try:
+                        data = json.loads(message)
+                        msg_type = data.get("type")
+
+                        if msg_type == "end_of_speech":
+                            # Mark final chunk
+                            final_chunk = AudioChunk(
+                                data=b"",
+                                is_final=True,
+                            )
+                            await audio_queue.put(final_chunk)
+                            end_of_speech.set()
+                            break
+
+                        elif msg_type == "reset":
+                            self.pipeline.reset_context()
+                            await websocket.send(
+                                json.dumps(
+                                    {
+                                        "type": "context_reset",
+                                    }
+                                )
+                            )
+
+                    except json.JSONDecodeError:
+                        logger.warning(f"Invalid JSON from client: {message}")
+
+        async def send_response():
+            """Stream response audio back to client."""
+            # Set up callbacks to send events
+            original_on_transcript = self.pipeline.on_transcript
+            original_on_llm_response = self.pipeline.on_llm_response
+
+            async def send_transcript(text: str):
+                await websocket.send(
+                    json.dumps(
+                        {
+                            "type": "transcript",
+                            "text": text,
+                        }
+                    )
+                )
+                if original_on_transcript:
+                    original_on_transcript(text)
+
+            async def send_llm_response(text: str):
+                await websocket.send(
+                    json.dumps(
+                        {
+                            "type": "response_text",
+                            "text": text,
+                        }
+                    )
+                )
+                if original_on_llm_response:
+                    original_on_llm_response(text)
+
+            # Temporarily override callbacks
+            self.pipeline.on_transcript = lambda t: asyncio.create_task(send_transcript(t))
+            self.pipeline.on_llm_response = lambda t: asyncio.create_task(send_llm_response(t))
+
+            try:
+                # Wait for some audio to arrive
+                await asyncio.sleep(0.1)
+
+                # Stream response
+                await websocket.send(json.dumps({"type": "response_start"}))
+
+                async for audio_chunk in self.pipeline.stream(audio_stream()):
+                    await websocket.send(audio_chunk.data)
+
+                await websocket.send(json.dumps({"type": "response_end"}))
+
+            finally:
+                # Restore original callbacks
+                self.pipeline.on_transcript = original_on_transcript
+                self.pipeline.on_llm_response = original_on_llm_response
+
+        # Run receive and send concurrently
+        receive_task = asyncio.create_task(receive_audio())
+        send_task = asyncio.create_task(send_response())
+
+        try:
+            await asyncio.gather(receive_task, send_task)
+        except Exception as e:
+            receive_task.cancel()
+            send_task.cancel()
+            raise
+
+    async def broadcast(self, message: str):
+        """Broadcast a message to all connected clients."""
+        if self._clients:
+            await asyncio.gather(*[ws.send(message) for ws in self._clients.values()])
+
+    @property
+    def client_count(self) -> int:
+        """Return number of connected clients."""
+        return len(self._clients)
+
+
+async def run_server(
+    pipeline: Pipeline,
+    host: str = "0.0.0.0",
+    port: int = 8765,
+):
+    """
+    Convenience function to run the WebSocket server.
+
+    Args:
+        pipeline: Configured Pipeline instance
+        host: Host to bind to
+        port: Port to listen on
+    """
+    server = WebSocketServer(pipeline, host, port)
+    await server.start()
+
+    try:
+        await asyncio.Future()  # Run forever
+    finally:
+        await server.stop()
+
+
+async def run_server_from_config(
+    config: Optional["AudioEngineConfig"] = None,
+    host: Optional[str] = None,
+    port: Optional[int] = None,
+    system_prompt: Optional[str] = None,
+):
+    """
+    Create and run WebSocket server from AudioEngineConfig.
+
+    Approach:
+    1. Load config from environment (or use provided config)
+    2. Create Pipeline with providers from config
+    3. Initialize and run WebSocket server
+
+    Rationale: Single entry point to run full audio pipeline server.
+
+    Args:
+        config: AudioEngineConfig instance (loads from env if None)
+        host: Host to bind to (default: from config)
+        port: Port to listen on (default: from config)
+        system_prompt: Optional system prompt override
+    """
+    from core.config import AudioEngineConfig
+
+    if config is None:
+        config = AudioEngineConfig.from_env()
+
+    pipeline = config.create_pipeline(system_prompt=system_prompt)
+
+    host = host or config.streaming.host
+    port = port or config.streaming.port
+
+    logger.info(
+        f"Starting audio engine server with providers: "
+        f"ASR={config.asr.provider}, "
+        f"LLM={config.llm.provider}, "
+        f"TTS={config.tts.provider}"
+    )
+
+    await run_server(pipeline, host, port)

audio_engine/tts/__init__.py

@@ -0,0 +1,35 @@
+"""TTS (Text-to-Speech) providers."""
+
+from ..core.config import TTSConfig
+
+from .base import BaseTTS
+from .cartesia import CartesiaTTS
+
+__all__ = ["BaseTTS", "CartesiaTTS", "get_tts_from_config"]
+
+
+def get_tts_from_config(config: TTSConfig) -> BaseTTS:
+    """
+    Instantiate TTS provider from config.
+
+    Args:
+        config: TTSConfig object with provider name and settings
+
+    Returns:
+        Initialized BaseTTS provider instance
+
+    Raises:
+        ValueError: If provider name is not recognized
+    """
+    provider_name = config.provider.lower()
+
+    if provider_name == "cartesia":
+        return CartesiaTTS(
+            api_key=config.api_key,
+            voice_id=config.voice_id,  # None will use DEFAULT_VOICE_ID in CartesiaTTS
+            model=config.model or "sonic-3",
+            speed=config.speed,
+            **config.extra,
+        )
+    else:
+        raise ValueError(f"Unknown TTS provider: {config.provider}. " f"Supported: cartesia")

audio_engine/tts/base.py

@@ -0,0 +1,153 @@
+"""Abstract base class for TTS (Text-to-Speech) providers."""
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterator, Optional
+
+from ..core.types import AudioChunk, AudioFormat
+
+
+class BaseTTS(ABC):
+    """
+    Abstract base class for Text-to-Speech providers.
+
+    All TTS implementations must inherit from this class and implement
+    the required methods for both batch and streaming audio synthesis.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        voice_id: Optional[str] = None,
+        model: Optional[str] = None,
+        speed: float = 1.0,
+        output_format: AudioFormat = AudioFormat.PCM_24K,
+        **kwargs,
+    ):
+        """
+        Initialize the TTS provider.
+
+        Args:
+            api_key: API key for the provider
+            voice_id: Voice identifier to use
+            model: Model identifier (if applicable)
+            speed: Speech speed multiplier (1.0 = normal)
+            output_format: Desired audio output format
+            **kwargs: Additional provider-specific configuration
+        """
+        self.api_key = api_key
+        self.voice_id = voice_id
+        self.model = model
+        self.speed = speed
+        self.output_format = output_format
+        self.config = kwargs
+
+    @abstractmethod
+    async def synthesize(self, text: str) -> bytes:
+        """
+        Synthesize complete audio from text.
+
+        Args:
+            text: Text to convert to speech
+
+        Returns:
+            Complete audio as bytes
+        """
+        pass
+
+    @abstractmethod
+    async def synthesize_stream(self, text: str) -> AsyncIterator[AudioChunk]:
+        """
+        Synthesize streaming audio from text.
+
+        Args:
+            text: Text to convert to speech
+
+        Yields:
+            AudioChunk objects with audio data
+        """
+        pass
+
+    async def synthesize_stream_text(
+        self, text_stream: AsyncIterator[str]
+    ) -> AsyncIterator[AudioChunk]:
+        """
+        Synthesize streaming audio from streaming text input.
+
+        This enables sentence-by-sentence TTS as the LLM generates text.
+        Default implementation buffers until punctuation. Override for
+        providers with native text streaming support.
+
+        Args:
+            text_stream: Async iterator yielding text chunks
+
+        Yields:
+            AudioChunk objects with audio data
+        """
+        buffer = ""
+        sentence_enders = ".!?;"
+
+        async for text_chunk in text_stream:
+            buffer += text_chunk
+
+            # Check if we have a complete sentence
+            for ender in sentence_enders:
+                if ender in buffer:
+                    # Split at the sentence boundary
+                    parts = buffer.split(ender, 1)
+                    sentence = parts[0] + ender
+
+                    if sentence.strip():
+                        async for audio_chunk in self.synthesize_stream(sentence.strip()):
+                            yield audio_chunk
+
+                    buffer = parts[1] if len(parts) > 1 else ""
+                    break
+
+        # Handle remaining text
+        if buffer.strip():
+            async for audio_chunk in self.synthesize_stream(buffer.strip()):
+                yield audio_chunk
+
+    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 TTS service.
+        Override in subclasses if needed.
+        """
+        pass
+
+    async def disconnect(self):
+        """
+        Close connection to the TTS service.
+        Override in subclasses if needed.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the name of this TTS provider."""
+        pass
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Whether this provider supports streaming audio output."""
+        return True
+
+    @property
+    def sample_rate(self) -> int:
+        """Return the sample rate for this provider's output."""
+        format_rates = {
+            AudioFormat.PCM_16K: 16000,
+            AudioFormat.PCM_24K: 24000,
+            AudioFormat.PCM_44K: 44100,
+        }
+        return format_rates.get(self.output_format, 24000)