atom-audio-engine 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

audio_engine/streaming/websocket_server.py

@@ -0,0 +1,341 @@
+"""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/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for the audio engine."""

audio_engine/tests/test_personaplex/__init__.py

@@ -0,0 +1 @@
+"""Tests for the audio engine."""

audio_engine/tests/test_personaplex/test_personaplex.py

@@ -0,0 +1,10 @@
+"""PersonaPlex test suite.
+
+Organized by step for maintainability:
+- test_personaplex_config.py     (Step 1: Config)
+- test_personaplex_message.py    (Step 2: Message encoding/decoding)
+- test_personaplex_transcript.py (Step 3: Transcript save/load)
+- test_personaplex_session.py    (Step 4: Session management)
+- test_personaplex_client.py     (Step 5: Client connection)
+- test_personaplex_pipeline.py   (Step 6-7: Pipeline lifecycle + mock messages)
+"""

audio_engine/tests/test_personaplex/test_personaplex_client.py

@@ -0,0 +1,259 @@
+"""Tests for PersonaPlexClient WebSocket connection (Step 5)."""
+
+import pytest
+from unittest.mock import AsyncMock, patch
+import asyncio
+
+from pipelines.personaplex import PersonaPlexClient, PersonaPlexConfig
+
+
+class TestPersonaPlexClientInit:
+    """Test client initialization and URL building."""
+
+    def test_client_init_with_defaults(self):
+        """Can we create a client with default config?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        assert client.config == config
+        assert client.connection is None
+        assert not client._is_connected
+
+    def test_client_init_with_custom_config(self):
+        """Can we create a client with custom config?"""
+        config = PersonaPlexConfig(
+            server_url="wss://custom.example.com",
+            voice_prompt="NATM0.pt",
+        )
+        client = PersonaPlexClient(config)
+
+        assert client.config.server_url == "wss://custom.example.com"
+        assert client.config.voice_prompt == "NATM0.pt"
+
+    def test_url_building_includes_voice_prompt(self):
+        """Does URL building include voice_prompt parameter?"""
+        config = PersonaPlexConfig(voice_prompt="NATF0.pt")
+        client = PersonaPlexClient(config)
+
+        url = client._build_url("Test prompt")
+
+        assert "voice_prompt=NATF0.pt" in url
+        assert "text_prompt=" in url
+
+    def test_url_building_includes_temperatures(self):
+        """Does URL include temperature parameters?"""
+        config = PersonaPlexConfig(
+            text_temperature=0.5,
+            audio_temperature=0.9,
+        )
+        client = PersonaPlexClient(config)
+
+        url = client._build_url("Test prompt")
+
+        assert "text_temperature=0.5" in url
+        assert "audio_temperature=0.9" in url
+
+    def test_url_building_encodes_system_prompt(self):
+        """Is system prompt included in URL?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        url = client._build_url("Hello world!")
+
+        assert "text_prompt=" in url
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexClientConnection:
+    """Test client WebSocket connection lifecycle."""
+
+    async def test_connect_opens_websocket(self):
+        """Does connect establish a WebSocket?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        with patch("websockets.connect", new_callable=AsyncMock) as mock_connect:
+            mock_conn = AsyncMock()
+            mock_connect.return_value = mock_conn
+
+            await client.connect("Hello assistant")
+
+            assert client._is_connected
+            mock_connect.assert_called_once()
+
+    async def test_disconnect_closes_websocket(self):
+        """Does disconnect close the WebSocket?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        mock_conn = AsyncMock()
+        client.connection = mock_conn
+        client._is_connected = True
+
+        await client.disconnect()
+
+        mock_conn.close.assert_called_once()
+        assert not client._is_connected
+        assert client.connection is None
+
+    async def test_context_manager_exits_cleanly(self):
+        """Does context manager call disconnect on exit?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        # Manually set up connection
+        mock_conn = AsyncMock()
+        client.connection = mock_conn
+        client._is_connected = True
+
+        async with client:
+            # Still connected during the context
+            assert client._is_connected
+
+        # Should be disconnected after exiting
+        mock_conn.close.assert_called_once()
+        assert not client._is_connected
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexClientSendReceive:
+    """Test sending and receiving data."""
+
+    async def test_send_audio_creates_message(self):
+        """Does send_audio create a binary message?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        mock_conn = AsyncMock()
+        client.connection = mock_conn
+        client._is_connected = True
+
+        audio_data = b"fake_opus_data"
+        await client.send_audio(audio_data)
+
+        mock_conn.send.assert_called_once()
+        sent_data = mock_conn.send.call_args[0][0]
+
+        # Should start with 0x01 (audio type)
+        assert sent_data[0] == 0x01
+        assert sent_data[1:] == audio_data
+
+    async def test_receive_audio_parses_message(self):
+        """Does receive_audio return AudioChunk?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        # Create a raw audio message (0x01 + audio data)
+        audio_data = b"received_opus"
+        raw_message = bytes([0x01]) + audio_data
+
+        mock_conn = AsyncMock()
+        mock_conn.recv.return_value = raw_message
+        client.connection = mock_conn
+        client._is_connected = True
+
+        from pipelines.personaplex import AudioChunk
+
+        chunk = await client.receive_audio()
+
+        assert isinstance(chunk, AudioChunk)
+        assert chunk.data == audio_data
+        assert chunk.sample_rate == 48000
+
+    async def test_receive_text_parses_message(self):
+        """Does receive_text return TextChunk?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        # Create a raw text message (0x02 + UTF-8 text)
+        text = "Hello from server"
+        raw_message = bytes([0x02]) + text.encode("utf-8")
+
+        mock_conn = AsyncMock()
+        mock_conn.recv.return_value = raw_message
+        client.connection = mock_conn
+        client._is_connected = True
+
+        from pipelines.personaplex import TextChunk
+
+        chunk = await client.receive_text()
+
+        assert isinstance(chunk, TextChunk)
+        assert chunk.text == text
+
+    async def test_receive_error_returns_none(self):
+        """Does receive_audio handle error messages gracefully?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        # Create an error message (0x05 + error text)
+        error_text = "Connection failed"
+        raw_message = bytes([0x05]) + error_text.encode("utf-8")
+
+        mock_conn = AsyncMock()
+        mock_conn.recv.return_value = raw_message
+        client.connection = mock_conn
+        client._is_connected = True
+
+        result = await client.receive_audio()
+
+        # Should return None, not raise
+        assert result is None
+
+
+@pytest.mark.asyncio
+class TestPersonaPlexClientStreaming:
+    """Test streaming message handling."""
+
+    async def test_stream_messages_yields_raw_messages(self):
+        """Does stream_messages yield PersonaPlexMessage objects?"""
+        from pipelines.personaplex import PersonaPlexMessage, MessageType
+
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        # Create mock messages - raw bytes
+        audio_msg = bytes([0x01]) + b"audio_data"
+        text_msg = bytes([0x02]) + "Hello".encode("utf-8")
+
+        # Mock connection returns bytes
+        mock_conn = AsyncMock()
+        mock_conn.__aiter__.return_value = [audio_msg, text_msg]
+        client.connection = mock_conn
+        client._is_connected = True
+
+        messages = []
+        try:
+            async for msg in client.stream_messages():
+                messages.append(msg)
+                if len(messages) >= 2:
+                    break
+        except (StopAsyncIteration, asyncio.CancelledError):
+            pass
+
+        # Should get 2 PersonaPlexMessage objects
+        assert len(messages) == 2
+        assert messages[0].type == MessageType.AUDIO
+        assert messages[1].type == MessageType.TEXT
+
+    async def test_stream_messages_skips_errors(self):
+        """Does stream_messages skip error messages?"""
+        config = PersonaPlexConfig()
+        client = PersonaPlexClient(config)
+
+        # Error message type (0x05)
+        error_msg = bytes([0x05]) + b"Server error"
+
+        mock_conn = AsyncMock()
+        mock_conn.recv.return_value = error_msg
+        client.connection = mock_conn
+        client._is_connected = True
+
+        chunks = []
+        async for chunk in client.stream_messages():
+            chunks.append(chunk)
+            if len(chunks) > 1:  # Prevent infinite loop
+                break
+
+        # Error should be skipped (logged), not returned as chunk
+        assert len(chunks) == 0

audio_engine/tests/test_personaplex/test_personaplex_config.py

@@ -0,0 +1,71 @@
+"""Tests for PersonaPlexConfig (Step 1)."""
+
+import pytest
+from pathlib import Path
+from pipelines.personaplex import PersonaPlexConfig
+
+
+class TestPersonaPlexConfig:
+    """Test PersonaPlexConfig creation and validation."""
+
+    def test_default_creation(self):
+        """Can we create PersonaPlexConfig with defaults?"""
+        config = PersonaPlexConfig()
+        assert config.voice_prompt == "NATF0.pt"
+        assert config.text_temperature == 0.7
+        assert config.audio_temperature == 0.8
+        assert config.sample_rate == 48000
+
+    def test_custom_values(self):
+        """Can we create config with custom values?"""
+        config = PersonaPlexConfig(
+            voice_prompt="NATM1.pt",
+            text_temperature=0.5,
+            audio_temperature=0.9,
+        )
+        assert config.voice_prompt == "NATM1.pt"
+        assert config.text_temperature == 0.5
+        assert config.audio_temperature == 0.9
+
+    def test_reject_invalid_text_temp(self):
+        """Does it reject text_temperature > 2.0?"""
+        with pytest.raises(ValueError):
+            PersonaPlexConfig(text_temperature=2.5)
+
+    def test_reject_invalid_audio_temp(self):
+        """Does it reject audio_temperature < 0.0?"""
+        with pytest.raises(ValueError):
+            PersonaPlexConfig(audio_temperature=-0.1)
+
+    def test_reject_invalid_topk(self):
+        """Does it reject invalid top-K values?"""
+        with pytest.raises(ValueError):
+            PersonaPlexConfig(text_topk=0)
+
+    def test_reject_invalid_sample_rate(self):
+        """Does it reject invalid sample rates?"""
+        with pytest.raises(ValueError):
+            PersonaPlexConfig(sample_rate=22050)
+
+    def test_from_dict(self):
+        """Can we load config from a dictionary?"""
+        data = {
+            "voice_prompt": "NATF2.pt",
+            "text_temperature": 0.8,
+            "audio_temperature": 0.9,
+        }
+        config = PersonaPlexConfig.from_dict(data)
+        assert config.voice_prompt == "NATF2.pt"
+        assert config.text_temperature == 0.8
+        assert config.audio_temperature == 0.9
+
+    def test_transcript_dir_created(self):
+        """Does it create transcript directory on init?"""
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config = PersonaPlexConfig(
+                transcript_path=tmpdir + "/transcripts/",
+                save_transcripts=True,
+            )
+            assert Path(config.transcript_path).exists()