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

audio_engine/tests/test_traditional_pipeline/test_cartesia_asr.py

@@ -0,0 +1,474 @@
+"""Tests for CartesiaASR (Speech-to-Text) provider."""
+
+import pytest
+import json
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+from asr.cartesia import CartesiaASR
+from core.types import AudioChunk, TranscriptChunk
+
+
+class TestCartesiaASRInitialization:
+    """Tests for CartesiaASR initialization."""
+
+    def test_init_with_api_key(self):
+        """Test initialization with API key."""
+        asr = CartesiaASR(api_key="sk_test123")
+        assert asr.api_key == "sk_test123"
+        assert asr.model == "ink-whisper"
+        assert asr.language == "en"
+        assert asr.encoding == "pcm_s16le"
+        assert asr.sample_rate == 16000
+
+    def test_init_with_custom_model(self):
+        """Test initialization with custom model."""
+        asr = CartesiaASR(api_key="sk_test", model="custom-model")
+        assert asr.model == "custom-model"
+
+    def test_init_with_custom_language(self):
+        """Test initialization with custom language."""
+        asr = CartesiaASR(api_key="sk_test", language="es")
+        assert asr.language == "es"
+
+    def test_init_with_vad_params(self):
+        """Test initialization with VAD parameters."""
+        asr = CartesiaASR(
+            api_key="sk_test",
+            min_volume=0.5,
+            max_silence_duration_secs=15.0,
+        )
+        assert asr.min_volume == 0.5
+        assert asr.max_silence_duration_secs == 15.0
+
+    def test_name_property(self):
+        """Test that name property returns 'cartesia'."""
+        asr = CartesiaASR(api_key="sk_test")
+        assert asr.name == "cartesia"
+
+
+class TestCartesiaASRConnection:
+    """Tests for CartesiaASR connection management."""
+
+    @pytest.mark.asyncio
+    async def test_connect_with_valid_api_key(self):
+        """Test successful connection with valid API key."""
+        asr = CartesiaASR(api_key="sk_test123")
+
+        with patch("websockets.connect", new_callable=AsyncMock) as mock_connect:
+            mock_ws = AsyncMock()
+            mock_connect.return_value = mock_ws
+
+            await asr.connect()
+
+            # Verify WebSocket was created with correct URL parameters
+            mock_connect.assert_called_once()
+            call_args = mock_connect.call_args[0][0]
+            assert "wss://api.cartesia.ai/stt/websocket" in call_args
+            assert "model=ink-whisper" in call_args
+            assert "language=en" in call_args
+            assert "encoding=pcm_s16le" in call_args
+            assert "sample_rate=16000" in call_args
+            assert "api_key=sk_test123" in call_args
+
+    @pytest.mark.asyncio
+    async def test_connect_without_api_key(self):
+        """Test connection failure without API key."""
+        asr = CartesiaASR(api_key=None)
+
+        with pytest.raises(ValueError, match="API key not provided"):
+            await asr.connect()
+
+    @pytest.mark.asyncio
+    async def test_connect_idempotent(self):
+        """Test that connect() is idempotent (doesn't reconnect if already connected)."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        with patch("websockets.connect", new_callable=AsyncMock) as mock_connect:
+            mock_connect.return_value = mock_ws
+            asr.websocket = mock_ws
+
+            await asr.connect()
+
+            # Should not call websockets.connect again
+            mock_connect.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_disconnect(self):
+        """Test WebSocket disconnect."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        # Create proper async mocks
+        mock_ws = MagicMock()
+        mock_ws.close = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Create a proper task mock
+        async def dummy_task():
+            await asyncio.sleep(1)
+
+        task = asyncio.create_task(dummy_task())
+        asr._receive_task = task
+
+        await asr.disconnect()
+
+        # Verify close was called
+        mock_ws.close.assert_called_once()
+        # Verify task was cancelled
+        assert task.cancelled()
+
+
+class TestCartesiaASRTranscribe:
+    """Tests for batch transcription."""
+
+    @pytest.mark.asyncio
+    async def test_transcribe_simple(self):
+        """Test batch transcription with simple audio."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop that puts responses in queue
+        async def mock_receive():
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "Hello world", "is_final": True}
+            )
+            await asr._response_queue.put({"type": "done"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        audio = b"\x00\x01" * 16000  # 1 second of audio at 16kHz
+
+        result = await asr.transcribe(audio)
+
+        assert result == "Hello world"
+        mock_ws.send.assert_called()
+
+    @pytest.mark.asyncio
+    async def test_transcribe_multiple_chunks(self):
+        """Test batch transcription with multiple response chunks."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop that puts multiple responses
+        async def mock_receive():
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "Hello ", "is_final": False}
+            )
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "world", "is_final": True}
+            )
+            await asr._response_queue.put({"type": "done"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        audio = b"\x00\x01" * 16000
+
+        result = await asr.transcribe(audio)
+
+        assert result == "Hello world"
+
+    @pytest.mark.asyncio
+    async def test_transcribe_error_response(self):
+        """Test batch transcription with error response from server."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop that returns error
+        async def mock_receive():
+            await asr._response_queue.put(
+                {"type": "error", "error": "Invalid audio format"}
+            )
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        audio = b"\x00\x01" * 16000
+
+        with pytest.raises(RuntimeError, match="Cartesia error"):
+            await asr.transcribe(audio)
+
+    @pytest.mark.asyncio
+    async def test_transcribe_sends_audio_in_chunks(self):
+        """Test that transcribe sends audio in 100ms chunks."""
+        asr = CartesiaASR(api_key="sk_test", sample_rate=16000)
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop
+        async def mock_receive():
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "test", "is_final": True}
+            )
+            await asr._response_queue.put({"type": "done"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        # 1 second of audio = 10 chunks of 100ms each
+        audio = b"\x00\x01" * 16000
+
+        await asr.transcribe(audio)
+
+        # Should have sent audio chunks + 'done' command
+        send_calls = [call[0][0] for call in mock_ws.send.call_args_list]
+        assert send_calls[-1] == "done"  # Last call should be 'done' command
+        assert len(send_calls) >= 2  # At least 1 audio chunk + done
+
+    @pytest.mark.asyncio
+    async def test_transcribe_timeout(self):
+        """Test transcribe timeout when server doesn't respond."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop that never sends response
+        async def mock_receive():
+            await asyncio.sleep(10)
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        audio = b"\x00\x01" * 16000
+
+        # Should timeout and return empty string
+        result = await asr.transcribe(audio)
+        assert result == ""
+
+
+class TestCartesiaASRTranscribeStream:
+    """Tests for streaming transcription."""
+
+    @pytest.mark.asyncio
+    async def test_transcribe_stream_simple(self):
+        """Test streaming transcription with single audio chunk."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop
+        async def mock_receive():
+            await asyncio.sleep(0.1)
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "Hello", "is_final": True}
+            )
+            await asr._response_queue.put({"type": "done"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        # Create audio stream
+        async def audio_stream():
+            yield AudioChunk(
+                data=b"\x00\x01" * 8000,
+                sample_rate=16000,
+                channels=1,
+                format="pcm_s16le",
+                is_final=True,
+            )
+
+        results = []
+        async for chunk in asr.transcribe_stream(audio_stream()):
+            results.append(chunk.text)
+
+        # Filter out empty results
+        non_empty = [r for r in results if r]
+        assert "Hello" in non_empty
+
+    @pytest.mark.asyncio
+    async def test_transcribe_stream_multiple_chunks(self):
+        """Test streaming transcription with multiple audio chunks."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop
+        async def mock_receive():
+            await asyncio.sleep(0.05)
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "Hello ", "is_final": False}
+            )
+            await asyncio.sleep(0.05)
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "world", "is_final": True}
+            )
+            await asr._response_queue.put({"type": "done"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        # Create audio stream with multiple chunks
+        async def audio_stream():
+            yield AudioChunk(
+                data=b"\x00\x01" * 4000,
+                sample_rate=16000,
+                channels=1,
+                format="pcm_s16le",
+                is_final=False,
+            )
+            yield AudioChunk(
+                data=b"\x00\x01" * 4000,
+                sample_rate=16000,
+                channels=1,
+                format="pcm_s16le",
+                is_final=True,
+            )
+
+        results = []
+        async for chunk in asr.transcribe_stream(audio_stream()):
+            results.append(chunk.text)
+
+        # Should receive partial and final results
+        non_empty = [r for r in results if r]
+        assert len(non_empty) >= 1
+
+    @pytest.mark.asyncio
+    async def test_transcribe_stream_sends_done_on_final(self):
+        """Test that 'done' command is sent when final audio chunk received."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop
+        async def mock_receive():
+            await asyncio.sleep(0.1)
+            await asr._response_queue.put(
+                {"type": "transcript", "text": "test", "is_final": True}
+            )
+            await asr._response_queue.put({"type": "done"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        # Create audio stream with final flag
+        async def audio_stream():
+            yield AudioChunk(
+                data=b"\x00\x01" * 8000,
+                sample_rate=16000,
+                channels=1,
+                format="pcm_s16le",
+                is_final=True,
+            )
+
+        async for _ in asr.transcribe_stream(audio_stream()):
+            pass
+
+        # Verify 'done' was sent
+        send_calls = [call[0][0] for call in mock_ws.send.call_args_list]
+        assert "done" in send_calls
+
+    @pytest.mark.asyncio
+    async def test_transcribe_stream_error(self):
+        """Test streaming transcription with error response."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        asr.websocket = mock_ws
+
+        # Mock receive loop with error
+        async def mock_receive():
+            await asyncio.sleep(0.1)
+            await asr._response_queue.put({"type": "error", "error": "Connection lost"})
+
+        asr._receive_task = asyncio.create_task(mock_receive())
+
+        async def audio_stream():
+            yield AudioChunk(
+                data=b"\x00\x01" * 8000,
+                sample_rate=16000,
+                channels=1,
+                format="pcm_s16le",
+                is_final=True,
+            )
+
+        # Should handle error gracefully
+        results = []
+        async for chunk in asr.transcribe_stream(audio_stream()):
+            results.append(chunk.text)
+
+
+class TestCartesiaASRReceiveLoop:
+    """Tests for background receive loop."""
+
+    @pytest.mark.asyncio
+    async def test_receive_loop_parses_json(self):
+        """Test that receive loop correctly parses JSON messages."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        messages = [
+            json.dumps({"type": "transcript", "text": "test", "is_final": True}),
+            json.dumps({"type": "done"}),
+        ]
+        mock_ws.__aiter__.return_value = messages
+
+        asr.websocket = mock_ws
+
+        # Run receive loop
+        await asr._receive_loop()
+
+        # Check that messages were queued
+        assert asr._response_queue.qsize() == 2
+
+    @pytest.mark.asyncio
+    async def test_receive_loop_handles_invalid_json(self):
+        """Test that receive loop handles invalid JSON gracefully."""
+        asr = CartesiaASR(api_key="sk_test")
+
+        mock_ws = AsyncMock()
+        messages = [
+            "invalid json {",
+            json.dumps({"type": "transcript", "text": "test", "is_final": True}),
+        ]
+        mock_ws.__aiter__.return_value = messages
+
+        asr.websocket = mock_ws
+
+        # Run receive loop - should not raise
+        await asr._receive_loop()
+
+        # Should still queue valid message
+        assert asr._response_queue.qsize() == 1
+
+
+class TestCartesiaASRFactory:
+    """Tests for factory integration."""
+
+    def test_factory_creates_cartesia_asr(self):
+        """Test that factory creates CartesiaASR from config."""
+        from core.config import ASRConfig
+        from asr import get_asr_from_config
+
+        config = ASRConfig(
+            provider="cartesia",
+            api_key="sk_test123",
+            model="ink-whisper",
+            language="en",
+        )
+
+        asr = get_asr_from_config(config)
+
+        assert isinstance(asr, CartesiaASR)
+        assert asr.api_key == "sk_test123"
+        assert asr.model == "ink-whisper"
+        assert asr.language == "en"
+
+    def test_factory_with_extra_params(self):
+        """Test factory passes extra parameters to CartesiaASR."""
+        from core.config import ASRConfig
+        from asr import get_asr_from_config
+
+        config = ASRConfig(
+            provider="cartesia",
+            api_key="sk_test",
+            extra={"min_volume": 0.5, "max_silence_duration_secs": 20.0},
+        )
+
+        asr = get_asr_from_config(config)
+
+        assert asr.min_volume == 0.5
+        assert asr.max_silence_duration_secs == 20.0

audio_engine/tests/test_traditional_pipeline/test_config_env.py

@@ -0,0 +1,97 @@
+"""Tests for AudioEngineConfig environment variable loading."""
+
+import os
+import pytest
+from core.config import AudioEngineConfig
+
+
+class TestConfigEnvironment:
+    """Test AudioEngineConfig environment variable loading."""
+
+    def test_from_env_uses_defaults(self):
+        """Does from_env() use defaults when no env vars set?"""
+        # Ensure env is clean
+        for key in [
+            "ASR_PROVIDER",
+            "ASR_API_KEY",
+            "LLM_PROVIDER",
+            "LLM_API_KEY",
+            "TTS_PROVIDER",
+            "TTS_API_KEY",
+            "DEEPGRAM_API_KEY",
+            "GROQ_API_KEY",
+            "CARTESIA_API_KEY",
+        ]:
+            os.environ.pop(key, None)
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.asr.provider == "cartesia"
+        assert config.llm.provider == "groq"
+        assert config.tts.provider == "cartesia"
+
+    def test_from_env_reads_asr_provider(self):
+        """Does from_env() read ASR_PROVIDER env var?"""
+        os.environ["ASR_PROVIDER"] = "openai"
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.asr.provider == "openai"
+        os.environ.pop("ASR_PROVIDER", None)
+
+    def test_from_env_reads_llm_provider(self):
+        """Does from_env() read LLM_PROVIDER env var?"""
+        os.environ["LLM_PROVIDER"] = "anthropic"
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.llm.provider == "anthropic"
+        os.environ.pop("LLM_PROVIDER", None)
+
+    def test_from_env_reads_tts_provider(self):
+        """Does from_env() read TTS_PROVIDER env var?"""
+        os.environ["TTS_PROVIDER"] = "elevenlabs"
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.tts.provider == "elevenlabs"
+        os.environ.pop("TTS_PROVIDER", None)
+
+    def test_from_env_reads_api_keys(self):
+        """Does from_env() read API key env vars?"""
+        os.environ["ASR_API_KEY"] = "test-asr-key"
+        os.environ["LLM_API_KEY"] = "test-llm-key"
+        os.environ["TTS_API_KEY"] = "test-tts-key"
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.asr.api_key == "test-asr-key"
+        assert config.llm.api_key == "test-llm-key"
+        assert config.tts.api_key == "test-tts-key"
+
+        os.environ.pop("ASR_API_KEY", None)
+        os.environ.pop("LLM_API_KEY", None)
+        os.environ.pop("TTS_API_KEY", None)
+
+    def test_from_env_fallback_to_provider_specific_keys(self):
+        """Does from_env() fallback to provider-specific keys (DEEPGRAM_API_KEY)?"""
+        os.environ["DEEPGRAM_API_KEY"] = "fallback-deepgram-key"
+        os.environ.pop("ASR_API_KEY", None)
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.asr.api_key == "fallback-deepgram-key"
+
+        os.environ.pop("DEEPGRAM_API_KEY", None)
+
+    def test_from_env_api_key_priority(self):
+        """Does ASR_API_KEY take priority over DEEPGRAM_API_KEY?"""
+        os.environ["ASR_API_KEY"] = "primary-key"
+        os.environ["DEEPGRAM_API_KEY"] = "fallback-key"
+
+        config = AudioEngineConfig.from_env()
+
+        assert config.asr.api_key == "primary-key"
+
+        os.environ.pop("ASR_API_KEY", None)
+        os.environ.pop("DEEPGRAM_API_KEY", None)

audio_engine/tests/test_traditional_pipeline/test_conversation_context.py

@@ -0,0 +1,115 @@
+"""Tests for ConversationContext."""
+
+import pytest
+from core.types import ConversationContext, ConversationMessage
+
+
+class TestConversationContext:
+    """Test ConversationContext message management."""
+
+    def test_add_user_message(self):
+        """Can we add a user message to context?"""
+        context = ConversationContext()
+
+        context.add_message("user", "Hello")
+
+        assert len(context.messages) == 1
+        assert context.messages[0].role == "user"
+        assert context.messages[0].content == "Hello"
+
+    def test_add_assistant_message(self):
+        """Can we add an assistant message to context?"""
+        context = ConversationContext()
+
+        context.add_message("assistant", "Hi there")
+
+        assert len(context.messages) == 1
+        assert context.messages[0].role == "assistant"
+        assert context.messages[0].content == "Hi there"
+
+    def test_message_history_order(self):
+        """Are messages stored in correct order?"""
+        context = ConversationContext()
+
+        context.add_message("user", "Message 1")
+        context.add_message("assistant", "Message 2")
+        context.add_message("user", "Message 3")
+
+        assert len(context.messages) == 3
+        assert context.messages[0].content == "Message 1"
+        assert context.messages[1].content == "Message 2"
+        assert context.messages[2].content == "Message 3"
+
+    def test_system_prompt_set_on_context(self):
+        """Can we set system_prompt on context?"""
+        prompt = "You are a helpful assistant"
+        context = ConversationContext(system_prompt=prompt)
+
+        assert context.system_prompt == prompt
+
+    def test_get_messages_for_llm_with_system_prompt(self):
+        """Does get_messages_for_llm() include system prompt?"""
+        context = ConversationContext(system_prompt="You are helpful")
+        context.add_message("user", "Hello")
+        context.add_message("assistant", "Hi")
+
+        messages = context.get_messages_for_llm()
+
+        assert len(messages) == 3
+        assert messages[0]["role"] == "system"
+        assert messages[0]["content"] == "You are helpful"
+
+    def test_get_messages_for_llm_without_system_prompt(self):
+        """Does get_messages_for_llm() work without system prompt?"""
+        context = ConversationContext()
+        context.add_message("user", "Hello")
+        context.add_message("assistant", "Hi")
+
+        messages = context.get_messages_for_llm()
+
+        assert len(messages) == 2
+        assert messages[0]["role"] == "user"
+        assert messages[1]["role"] == "assistant"
+
+    def test_get_messages_for_llm_format(self):
+        """Does get_messages_for_llm() return correct format?"""
+        context = ConversationContext()
+        context.add_message("user", "Test message")
+
+        messages = context.get_messages_for_llm()
+
+        assert isinstance(messages, list)
+        assert isinstance(messages[0], dict)
+        assert "role" in messages[0]
+        assert "content" in messages[0]
+
+    def test_max_history_trim(self):
+        """Does context trim history when max_history exceeded?"""
+        context = ConversationContext(max_history=3)
+
+        context.add_message("user", "Message 1")
+        context.add_message("user", "Message 2")
+        context.add_message("user", "Message 3")
+        context.add_message("user", "Message 4")
+
+        assert len(context.messages) == 3
+        assert context.messages[0].content == "Message 2"
+        assert context.messages[2].content == "Message 4"
+
+    def test_clear_messages(self):
+        """Can we clear the message history?"""
+        context = ConversationContext()
+        context.add_message("user", "Hello")
+        context.add_message("assistant", "Hi")
+
+        context.messages = []
+
+        assert len(context.messages) == 0
+
+    def test_message_with_timestamp(self):
+        """Can we add messages with timestamps?"""
+        context = ConversationContext()
+
+        context.add_message("user", "Hello", timestamp_ms=1000)
+
+        assert context.messages[0].timestamp_ms == 1000