diff --git a/.env.example b/.env.example index 3e70778..0c0fde5 100644 --- a/.env.example +++ b/.env.example @@ -85,6 +85,11 @@ FAL_KEY= # --- Soniox --- SONIOX_API_KEY= +# --- AssemblyAI --- +ASSEMBLYAI_API_KEY= +# ASSEMBLYAI_BASE_URL=https://api.assemblyai.com +# ASSEMBLYAI_SPEAKER_LABELS=false + # --- Picovoice Leopard --- PICOVOICE_ACCESS_KEY= diff --git a/sapat/providers/assemblyai.py b/sapat/providers/assemblyai.py new file mode 100644 index 0000000..9cd2862 --- /dev/null +++ b/sapat/providers/assemblyai.py @@ -0,0 +1,152 @@ +# ABOUTME: AssemblyAI pre-recorded transcription provider +# ABOUTME: Uses upload -> transcript submission -> polling REST workflow + +import os +from typing import List, Optional + +import requests + +from sapat.providers import register +from sapat.providers.async_poll import AsyncPollProvider +from sapat.providers.base import ( + AudioFormat, + ProviderConfig, + TranscriptionResult, +) + + +def _env_bool(name: str) -> Optional[bool]: + value = os.getenv(name) + if value is None or value == "": + return None + return value.lower() in {"1", "true", "yes", "on"} + + +@register +class AssemblyAIProvider(AsyncPollProvider): + """AssemblyAI pre-recorded speech-to-text provider.""" + + name = "assemblyai" + config = ProviderConfig( + required_env_vars=["ASSEMBLYAI_API_KEY"], + max_file_size_mb=100.0, + preferred_format=AudioFormat.MP3, + default_model="universal-3-pro,universal-2", + ) + + poll_interval: float = float(os.getenv("ASSEMBLYAI_POLL_INTERVAL_SECONDS", "3")) + max_poll_time: float = float(os.getenv("ASSEMBLYAI_TIMEOUT_SECONDS", "600")) + + def __init__(self): + super().__init__() + self.api_key = os.getenv("ASSEMBLYAI_API_KEY", "") + self.base_url = os.getenv( + "ASSEMBLYAI_BASE_URL", "https://api.assemblyai.com" + ).rstrip("/") + + def _headers(self, content_type: Optional[str] = None) -> dict: + headers = {"Authorization": self.api_key} + if content_type: + headers["Content-Type"] = content_type + return headers + + def _speech_models(self, model: str) -> List[str]: + models = [part.strip() for part in model.split(",") if part.strip()] + return models or ["universal-3-pro", "universal-2"] + + def _upload(self, audio_file: str, model: str, language: str, **kwargs) -> str: + upload_url = self._upload_audio(audio_file) + return self._submit_transcript(upload_url, model, language, **kwargs) + + def _upload_audio(self, audio_file: str) -> str: + with open(audio_file, "rb") as f: + response = requests.post( + f"{self.base_url}/v2/upload", + headers=self._headers("application/octet-stream"), + data=f, + timeout=120, + ) + + if response.status_code != 200: + raise RuntimeError(f"AssemblyAI upload failed: {response.text}") + + upload_url = response.json().get("upload_url") + if not upload_url: + raise RuntimeError("AssemblyAI upload response did not include upload_url.") + return upload_url + + def _submit_transcript( + self, audio_url: str, model: str, language: str, **kwargs + ) -> str: + payload = { + "audio_url": audio_url, + "speech_models": self._speech_models(model), + } + + if language and language.lower() not in {"auto", "detect"}: + payload["language_code"] = language + else: + payload["language_detection"] = True + + prompt = kwargs.get("prompt") + if prompt: + payload["prompt"] = prompt + + for env_name, field_name in ( + ("ASSEMBLYAI_SPEAKER_LABELS", "speaker_labels"), + ("ASSEMBLYAI_PUNCTUATE", "punctuate"), + ("ASSEMBLYAI_FORMAT_TEXT", "format_text"), + ("ASSEMBLYAI_DISFLUENCIES", "disfluencies"), + ): + value = _env_bool(env_name) + if value is not None: + payload[field_name] = value + + response = requests.post( + f"{self.base_url}/v2/transcript", + headers=self._headers("application/json"), + json=payload, + timeout=60, + ) + + if response.status_code not in (200, 201): + raise RuntimeError( + f"AssemblyAI transcript submission failed: {response.text}" + ) + + transcript_id = response.json().get("id") + if not transcript_id: + raise RuntimeError("AssemblyAI transcript response did not include id.") + return transcript_id + + def _poll(self, job_id: str) -> str: + result = self._get_transcript(job_id) + status = result.get("status") + if status == "completed": + return "completed" + if status in {"error", "failed"}: + return "failed" + return "pending" + + def _fetch_result(self, job_id: str) -> TranscriptionResult: + result = self._get_transcript(job_id) + if result.get("status") != "completed": + raise RuntimeError(f"AssemblyAI transcript {job_id} is not completed.") + + return TranscriptionResult( + text=result.get("text", ""), + language=result.get("language_code"), + duration=result.get("audio_duration"), + segments=result.get("utterances") or result.get("words"), + raw_response=result, + ) + + def _get_transcript(self, transcript_id: str) -> dict: + response = requests.get( + f"{self.base_url}/v2/transcript/{transcript_id}", + headers=self._headers(), + timeout=30, + ) + if response.status_code != 200: + raise RuntimeError(f"AssemblyAI transcript fetch failed: {response.text}") + return response.json() diff --git a/sapat/providers/async_poll.py b/sapat/providers/async_poll.py index b2acdb1..a86cf0b 100644 --- a/sapat/providers/async_poll.py +++ b/sapat/providers/async_poll.py @@ -45,7 +45,12 @@ def transcribe( temperature: float = 0, **kwargs, ) -> TranscriptionResult: - job_id = self._upload(audio_file, model, language, **kwargs) + upload_kwargs = dict(kwargs) + if prompt: + upload_kwargs["prompt"] = prompt + upload_kwargs["temperature"] = temperature + + job_id = self._upload(audio_file, model, language, **upload_kwargs) elapsed = 0.0 while elapsed < self.max_poll_time: @@ -57,4 +62,6 @@ def transcribe( time.sleep(self.poll_interval) elapsed += self.poll_interval - raise TimeoutError(f"{self.name} job {job_id} timed out after {self.max_poll_time}s") + raise TimeoutError( + f"{self.name} job {job_id} timed out after {self.max_poll_time}s" + ) diff --git a/tests/providers/test_group_b.py b/tests/providers/test_group_b.py index 5dd344a..0a4e0f0 100644 --- a/tests/providers/test_group_b.py +++ b/tests/providers/test_group_b.py @@ -54,7 +54,9 @@ class TestSymblProvider: ) @patch("sapat.providers.symbl.requests.get") @patch("sapat.providers.symbl.requests.post") - def test_transcribe_submits_polls_and_fetches(self, mock_post, mock_get, audio_file): + def test_transcribe_submits_polls_and_fetches( + self, mock_post, mock_get, audio_file + ): # _upload: POST to /process/audio mock_post.return_value = FakeResponse( status_code=201, @@ -86,7 +88,10 @@ def test_transcribe_submits_polls_and_fetches(self, mock_post, mock_get, audio_f assert result.text == "First sentence.\nSecond sentence." mock_post.assert_called_once() assert "Authorization" in mock_post.call_args.kwargs["headers"] - assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer test-token" + assert ( + mock_post.call_args.kwargs["headers"]["Authorization"] + == "Bearer test-token" + ) @patch.dict( os.environ, @@ -164,7 +169,130 @@ def test_default_model(self): # =========================================================================== -# 2. Gladia +# 2. AssemblyAI +# =========================================================================== + + +class TestAssemblyAIProvider: + @patch.dict( + os.environ, + { + "ASSEMBLYAI_API_KEY": "test-key", + "ASSEMBLYAI_BASE_URL": "https://assembly.test", + "ASSEMBLYAI_SPEAKER_LABELS": "true", + }, + clear=False, + ) + @patch("sapat.providers.assemblyai.requests.get") + @patch("sapat.providers.assemblyai.requests.post") + def test_transcribe_uploads_submits_polls_and_fetches( + self, mock_post, mock_get, audio_file + ): + mock_post.side_effect = [ + FakeResponse( + status_code=200, + payload={"upload_url": "https://cdn.assembly.test/audio.mp3"}, + ), + FakeResponse(status_code=200, payload={"id": "transcript-123"}), + ] + mock_get.side_effect = [ + FakeResponse(status_code=200, payload={"status": "queued"}), + FakeResponse(status_code=200, payload={"status": "completed"}), + FakeResponse( + status_code=200, + payload={ + "status": "completed", + "text": "Hello from AssemblyAI.", + "language_code": "en", + "audio_duration": 12.5, + "utterances": [{"speaker": "A", "text": "Hello from AssemblyAI."}], + }, + ), + ] + + from sapat.providers.assemblyai import AssemblyAIProvider + + provider = AssemblyAIProvider() + provider.poll_interval = 0.01 + result = provider.transcribe( + audio_file, + model="universal-3-pro, universal-2", + language="en", + prompt="Product names: Sapat", + ) + + assert isinstance(result, TranscriptionResult) + assert result.text == "Hello from AssemblyAI." + assert result.language == "en" + assert result.duration == 12.5 + + upload_call, transcript_call = mock_post.call_args_list + assert upload_call.args[0] == "https://assembly.test/v2/upload" + assert upload_call.kwargs["headers"]["Authorization"] == "test-key" + assert ( + upload_call.kwargs["headers"]["Content-Type"] == "application/octet-stream" + ) + + assert transcript_call.args[0] == "https://assembly.test/v2/transcript" + payload = transcript_call.kwargs["json"] + assert payload["audio_url"] == "https://cdn.assembly.test/audio.mp3" + assert payload["speech_models"] == ["universal-3-pro", "universal-2"] + assert payload["language_code"] == "en" + assert payload["prompt"] == "Product names: Sapat" + assert payload["speaker_labels"] is True + + @patch.dict(os.environ, {"ASSEMBLYAI_API_KEY": "test-key"}, clear=False) + def test_default_model(self): + from sapat.providers.assemblyai import AssemblyAIProvider + + assert AssemblyAIProvider.config.default_model == "universal-3-pro,universal-2" + + @patch.dict(os.environ, {"ASSEMBLYAI_API_KEY": "test-key"}, clear=False) + def test_available_with_key(self): + from sapat.providers.assemblyai import AssemblyAIProvider + + assert AssemblyAIProvider.is_available() is True + + @patch.dict(os.environ, {}, clear=True) + def test_not_available_without_key(self): + from sapat.providers.assemblyai import AssemblyAIProvider + + assert AssemblyAIProvider.is_available() is False + + @patch.dict(os.environ, {"ASSEMBLYAI_API_KEY": "test-key"}, clear=False) + @patch("sapat.providers.assemblyai.requests.get") + @patch("sapat.providers.assemblyai.requests.post") + def test_failed_transcript_raises(self, mock_post, mock_get, audio_file): + mock_post.side_effect = [ + FakeResponse(status_code=200, payload={"upload_url": "https://cdn/a.mp3"}), + FakeResponse(status_code=200, payload={"id": "transcript-123"}), + ] + mock_get.return_value = FakeResponse( + status_code=200, + payload={"status": "error", "error": "bad audio"}, + ) + + from sapat.providers.assemblyai import AssemblyAIProvider + + provider = AssemblyAIProvider() + provider.poll_interval = 0.01 + with pytest.raises(RuntimeError, match="failed"): + provider.transcribe(audio_file, model="universal-3-pro") + + @patch.dict(os.environ, {"ASSEMBLYAI_API_KEY": "test-key"}, clear=False) + @patch("sapat.providers.assemblyai.requests.post") + def test_upload_response_must_include_url(self, mock_post, audio_file): + mock_post.return_value = FakeResponse(status_code=200, payload={}) + + from sapat.providers.assemblyai import AssemblyAIProvider + + provider = AssemblyAIProvider() + with pytest.raises(RuntimeError, match="upload_url"): + provider.transcribe(audio_file, model="universal-3-pro") + + +# =========================================================================== +# 3. Gladia # =========================================================================== @@ -172,7 +300,9 @@ class TestGladiaProvider: @patch.dict(os.environ, {"GLADIA_API_KEY": "test-key"}, clear=False) @patch("sapat.providers.gladia.requests.get") @patch("sapat.providers.gladia.requests.post") - def test_transcribe_uploads_creates_job_and_polls(self, mock_post, mock_get, audio_file): + def test_transcribe_uploads_creates_job_and_polls( + self, mock_post, mock_get, audio_file + ): # Step 1: upload audio -> audio_url # Step 2: create transcription job -> result_url upload_response = FakeResponse( @@ -263,7 +393,9 @@ class TestSpeechmaticsProvider: @patch.dict(os.environ, {"SPEECHMATICS_API_KEY": "test-key"}, clear=False) @patch("sapat.providers.speechmatics.requests.get") @patch("sapat.providers.speechmatics.requests.post") - def test_transcribe_creates_job_and_fetches_text(self, mock_post, mock_get, audio_file): + def test_transcribe_creates_job_and_fetches_text( + self, mock_post, mock_get, audio_file + ): create_response = FakeResponse(status_code=201, payload={"id": "job-123"}) mock_post.return_value = create_response @@ -311,7 +443,9 @@ def test_default_model(self): @patch("sapat.providers.speechmatics.requests.get") @patch("sapat.providers.speechmatics.requests.post") def test_rejected_job_raises(self, mock_post, mock_get, audio_file): - mock_post.return_value = FakeResponse(status_code=201, payload={"id": "job-123"}) + mock_post.return_value = FakeResponse( + status_code=201, payload={"id": "job-123"} + ) mock_get.return_value = FakeResponse( status_code=200, payload={"job": {"status": "rejected"}} ) @@ -346,6 +480,7 @@ def test_transcribe_converts_and_posts(self, mock_run, mock_post, audio_file): # Simulate ffmpeg creating the output file def fake_run(cmd, **kwargs): Path(cmd[-1]).write_bytes(b"fake oggopus") + mock_run.side_effect = fake_run mock_post.return_value = FakeResponse( @@ -382,6 +517,7 @@ def fake_run(cmd, **kwargs): def test_iam_token_auth_with_folder_id(self, mock_run, mock_post, audio_file): def fake_run(cmd, **kwargs): Path(cmd[-1]).write_bytes(b"fake oggopus") + mock_run.side_effect = fake_run mock_post.return_value = FakeResponse( @@ -432,9 +568,11 @@ def test_default_model(self): def _fake_oci(object_client, speech_client): """Build a fake oci module for testing Oracle provider.""" + def record_factory(name): def factory(**kwargs): return types.SimpleNamespace(_model_name=name, **kwargs) + return factory models = types.SimpleNamespace( @@ -472,13 +610,13 @@ class TestOracleProvider: def test_transcribe_uploads_polls_and_reads_output(self, audio_file): object_client = Mock() speech_client = Mock() - speech_client.create_transcription_job.return_value.data = types.SimpleNamespace( - id="job1" + speech_client.create_transcription_job.return_value.data = ( + types.SimpleNamespace(id="job1") ) - speech_client.list_transcription_tasks.return_value.data = types.SimpleNamespace( - items=[ - types.SimpleNamespace(id="task1", lifecycle_state="SUCCEEDED") - ] + speech_client.list_transcription_tasks.return_value.data = ( + types.SimpleNamespace( + items=[types.SimpleNamespace(id="task1", lifecycle_state="SUCCEEDED")] + ) ) speech_client.get_transcription_task.return_value.data = types.SimpleNamespace( output_location=types.SimpleNamespace( @@ -547,7 +685,9 @@ def test_oracle_model_requires_locale(self): def test_extract_text_from_json(self): from sapat.providers.oracle import OracleProvider - raw = json.dumps({"transcriptions": [{"transcription": "hello"}, {"transcription": "world"}]}) + raw = json.dumps( + {"transcriptions": [{"transcription": "hello"}, {"transcription": "world"}]} + ) assert OracleProvider._extract_text(raw) == "hello\nworld" def test_extract_text_from_plain_text(self): @@ -599,7 +739,9 @@ def test_transcribe_sends_correct_model_and_input(self, audio_file): ) def test_translate_flag_passed_through(self, audio_file): mock_replicate = MagicMock() - mock_replicate.Client.return_value.run.return_value = {"text": "translated text"} + mock_replicate.Client.return_value.run.return_value = { + "text": "translated text" + } from sapat.providers.replicate import ReplicateProvider diff --git a/tests/test_registry.py b/tests/test_registry.py index 3e9be8b..be5af90 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -24,7 +24,9 @@ class FakeProvider(TranscriptionProvider): name = "fake_test" config = ProviderConfig(required_env_vars=[]) - def transcribe(self, audio_file, model, language="en", prompt=None, temperature=0, **kwargs): + def transcribe( + self, audio_file, model, language="en", prompt=None, temperature=0, **kwargs + ): return TranscriptionResult(text="fake") @@ -32,6 +34,7 @@ def transcribe(self, audio_file, model, language="en", prompt=None, temperature= def reset_registry(): """Reset the registry before each test.""" import sapat.providers as reg + reg._registry.clear() reg._discovered = False yield @@ -92,21 +95,33 @@ def transcribe(self, **kw): class TestAutoDiscovery: def test_discovers_azure_when_env_set(self): - with patch.dict(os.environ, { - "AZURE_OPENAI_API_KEY": "test", - "AZURE_OPENAI_ENDPOINT": "https://test.openai.azure.com", - "AZURE_OPENAI_STT_API_VERSION": "2024-02-01", - }): + with patch.dict( + os.environ, + { + "AZURE_OPENAI_API_KEY": "test", + "AZURE_OPENAI_ENDPOINT": "https://test.openai.azure.com", + "AZURE_OPENAI_STT_API_VERSION": "2024-02-01", + }, + ): from sapat.providers.azure import AzureProvider + register(AzureProvider) assert "azure" in _registry def test_discovers_groq_when_env_set(self): with patch.dict(os.environ, {"GROQ_API_KEY": "test"}): from sapat.providers.groq import GroqProvider + register(GroqProvider) assert "groq" in _registry + def test_discovers_assemblyai_when_env_set(self): + with patch.dict(os.environ, {"ASSEMBLYAI_API_KEY": "test"}): + from sapat.providers.assemblyai import AssemblyAIProvider + + register(AssemblyAIProvider) + assert "assemblyai" in _registry + def test_no_discovery_without_env(self): with patch.dict(os.environ, {}, clear=True): _discover_providers()