videosdk-plugins-assemblyai 0.0.30__py3-none-any.whl → 0.0.32__py3-none-any.whl

videosdk/plugins/assemblyai/stt.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 import asyncio
-import io
+import json
 import os
-import wave
-from typing import Any
+from typing import Any, Optional
+from urllib.parse import urlencode
+import logging
 
-import aiohttp
 import numpy as np
-from videosdk.agents import STT, STTResponse, SpeechData, SpeechEventType, global_event_emitter
+import aiohttp
+from videosdk.agents import STT as BaseSTT, STTResponse, SpeechData, SpeechEventType, global_event_emitter
 
 try:
     from scipy import signal
@@ -16,134 +17,280 @@ try:
 except ImportError:
     SCIPY_AVAILABLE = False
 
-ASSEMBLYAI_API_URL = "https://api.assemblyai.com/v2"
+logger = logging.getLogger(__name__)
 
-class AssemblyAISTT(STT):
-    """
-    VideoSDK Agent Framework STT plugin for AssemblyAI.
-    """
+class AssemblyAISTT(BaseSTT):
 
     def __init__(
         self,
         *,
         api_key: str | None = None,
-        language_code: str = "en_us",
         input_sample_rate: int = 48000,
         target_sample_rate: int = 16000,
-        silence_threshold: float = 0.015,
-        silence_duration: float = 0.8,
+        format_turns: bool = True,
+        word_boost: list[str] | None = None,
+        end_of_turn_confidence_threshold: float = 0.5,
+        min_end_of_turn_silence_when_confident: int = 800,
+        max_turn_silence: int = 2000,
     ) -> None:
+        """Initialize the AssemblyAI STT plugin.
+
+        Args:
+            api_key (str | None, optional): AssemblyAI API key. Uses ASSEMBLYAI_API_KEY environment variable if not provided. Defaults to None.
+            input_sample_rate (int): The input sample rate to use for the STT plugin. Defaults to 48000.
+            target_sample_rate (int): The target sample rate to use for the STT plugin. Defaults to 16000.
+            format_turns (bool): Whether to format turns. Defaults to True.
+            word_boost (list[str] | None, optional): The word boost to use for the STT plugin. Defaults to None.
+            end_of_turn_confidence_threshold (float): The end of turn confidence threshold to use for the STT plugin. Defaults to 0.5.
+            min_end_of_turn_silence_when_confident (int): The minimum end of turn silence when confident to use for the STT plugin. Defaults to 800.
+            max_turn_silence (int): The maximum turn silence to use for the STT plugin. Defaults to 2000.
+        """
         super().__init__()
+        
         if not SCIPY_AVAILABLE:
             raise ImportError("scipy is not installed. Please install it with 'pip install scipy'")
 
         self.api_key = api_key or os.getenv("ASSEMBLYAI_API_KEY")
         if not self.api_key:
-            raise ValueError("AssemblyAI API key must be provided either through the 'api_key' parameter or the 'ASSEMBLYAI_API_KEY' environment variable.")
-
-        self.language_code = language_code
+            raise ValueError(
+                "AssemblyAI API key must be provided either through the 'api_key' parameter "
+                "or the 'ASSEMBLYAI_API_KEY' environment variable."
+            )
         self.input_sample_rate = input_sample_rate
         self.target_sample_rate = target_sample_rate
-        self.silence_threshold_bytes = int(silence_threshold * 32767)
-        self.silence_duration_frames = int(silence_duration * self.input_sample_rate)
-
-        self._session = aiohttp.ClientSession(headers={"Authorization": self.api_key})
-        self._audio_buffer = bytearray()
-        self._is_speaking = False
-        self._silence_frames = 0
-        self._lock = asyncio.Lock()
+        self.format_turns = format_turns
+        self.word_boost = word_boost or []
+        self.end_of_turn_confidence_threshold = end_of_turn_confidence_threshold
+        self.min_end_of_turn_silence_when_confident = min_end_of_turn_silence_when_confident
+        self.max_turn_silence = max_turn_silence
 
-    async def process_audio(self, audio_frames: bytes, **kwargs: Any) -> None:
-        async with self._lock:
-            is_silent_chunk = self._is_silent(audio_frames)
+        connection_params = {
+            "sample_rate": self.target_sample_rate,
+            "format_turns": self.format_turns,
+        }
+        
             
-            if not is_silent_chunk:
-                if not self._is_speaking:
-                    self._is_speaking = True
-                    global_event_emitter.emit("speech_started")
-                self._audio_buffer.extend(audio_frames)
-                self._silence_frames = 0
-            else:
-                if self._is_speaking:
-                    self._silence_frames += len(audio_frames) // 4 
-                    if self._silence_frames > self.silence_duration_frames:
-                        global_event_emitter.emit("speech_stopped")
-                        asyncio.create_task(self._transcribe_buffer())
-                        self._is_speaking = False
-                        self._silence_frames = 0
+        if self.end_of_turn_confidence_threshold != 0.7:
+            connection_params["end_of_turn_confidence_threshold"] = self.end_of_turn_confidence_threshold
+        if self.min_end_of_turn_silence_when_confident != 1500:
+            connection_params["min_end_of_turn_silence_when_confident"] = self.min_end_of_turn_silence_when_confident
+        if self.max_turn_silence != 3000:
+            connection_params["max_turn_silence"] = self.max_turn_silence
+        
+        if self.word_boost:
+            connection_params["word_boost"] = json.dumps(self.word_boost)
 
-    def _is_silent(self, audio_chunk: bytes) -> bool:
-        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
-        return np.max(np.abs(audio_data)) < self.silence_threshold_bytes
+        self.ws_url = f"wss://streaming.assemblyai.com/v3/ws?{urlencode(connection_params)}"
+        logger.info(f"[AssemblyAI] WebSocket URL: {self.ws_url}")
 
-    async def _transcribe_buffer(self):
-        async with self._lock:
-            if not self._audio_buffer:
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
+        self._ws_task: Optional[asyncio.Task] = None
+        
+        self._stream_buffer = bytearray()
+        self._target_chunk_size = int(0.1 * self.target_sample_rate * 2)  
+        self._min_chunk_size = int(0.05 * self.target_sample_rate * 2)   
+        
+        self._last_speech_event_time = 0.0
+        self._last_transcript = ""
+        self._is_speaking = False
+
+    async def process_audio(
+        self,
+        audio_frames: bytes,
+        **kwargs: Any
+    ) -> None:
+        """Process audio frames and send to AssemblyAI's Streaming API"""
+        
+        if not self._ws:
+            await self._connect_ws()
+            self._ws_task = asyncio.create_task(self._listen_for_responses())
+            
+        try:
+            resampled_audio = self._resample_audio(audio_frames)
+            if not resampled_audio:
                 return
-            audio_to_send = self._audio_buffer
-            self._audio_buffer = bytearray()
+                
+            self._stream_buffer.extend(resampled_audio)
+            
+            while len(self._stream_buffer) >= self._target_chunk_size:
+                chunk_to_send = bytes(self._stream_buffer[:self._target_chunk_size])
+                self._stream_buffer = self._stream_buffer[self._target_chunk_size:]
+                
+                await self._ws.send_bytes(chunk_to_send)
+                
+        except Exception as e:
+            logger.error(f"Error in process_audio: {str(e)}")
+            self.emit("error", str(e))
+            if self._ws:
+                await self._ws.close()
+                self._ws = None
+                if self._ws_task:
+                    self._ws_task.cancel()
+                    self._ws_task = None
+
+    async def _listen_for_responses(self) -> None:
+        """Background task to listen for WebSocket responses"""
+        if not self._ws:
+            return
+            
+        try:
+            async for msg in self._ws:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    data = msg.json()
+                    responses = self._handle_ws_message(data)
+                    for response in responses:
+                        if self._transcript_callback:
+                            await self._transcript_callback(response)
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    logger.error(f"WebSocket error: {self._ws.exception()}")
+                    self.emit("error", f"WebSocket error: {self._ws.exception()}")
+                    break
+        except Exception as e:
+            logger.error(f"Error in WebSocket listener: {str(e)}")
+            self.emit("error", f"Error in WebSocket listener: {str(e)}")
+        finally:
+            if self._ws:
+                await self._ws.close()
+                self._ws = None
+                
+    async def _connect_ws(self) -> None:
+        """Establish WebSocket connection with AssemblyAI's Streaming API"""
+        
+        if not self._session:
+            self._session = aiohttp.ClientSession()
+            
+        headers = {
+            "Authorization": self.api_key,
+        }
+        
+        try:
+            self._ws = await self._session.ws_connect(self.ws_url, headers=headers)
+            logger.info("[AssemblyAI] WebSocket connection opened")
+        except Exception as e:
+            logger.error(f"Error connecting to WebSocket: {str(e)}")
+            raise
+        
+    def _handle_ws_message(self, msg: dict) -> list[STTResponse]:
+        """Handle incoming WebSocket messages and generate STT responses"""
+        responses = []
         
         try:
-            resampled_audio_bytes = self._resample_audio(audio_to_send)
-            wav_audio = self._create_wav_in_memory(resampled_audio_bytes)
-
-            upload_url = f"{ASSEMBLYAI_API_URL}/upload"
-            async with self._session.post(upload_url, data=wav_audio) as response:
-                response.raise_for_status()
-                upload_data = await response.json()
-                audio_url = upload_data["upload_url"]
-
-            transcript_url = f"{ASSEMBLYAI_API_URL}/transcript"
-            payload = {"audio_url": audio_url, "language_code": self.language_code}
-            async with self._session.post(transcript_url, json=payload) as response:
-                response.raise_for_status()
-                transcript_data = await response.json()
-                transcript_id = transcript_data["id"]
-
-            poll_url = f"{ASSEMBLYAI_API_URL}/transcript/{transcript_id}"
-            while True:
-                await asyncio.sleep(1)
-                async with self._session.get(poll_url) as response:
-                    response.raise_for_status()
-                    result = await response.json()
-                    if result["status"] == "completed":
-                        if result.get("text") and self._transcript_callback:
-                            event = STTResponse(
-                                event_type=SpeechEventType.FINAL,
-                                data=SpeechData(text=result["text"], language=self.language_code, confidence=result.get("confidence", 1.0))
-                            )
-                            await self._transcript_callback(event)
-                        break
-                    elif result["status"] == "error":
-                        raise Exception(f"AssemblyAI transcription failed: {result.get('error')}")
+            msg_type = msg.get('type')
+            logger.info(f"[AssemblyAI] Message type: {msg_type}")
+
+            if msg_type == "Begin":
+                session_id = msg.get('id')
+                logger.info(f"[AssemblyAI] Session began: ID={session_id}")
+                
+            elif msg_type == "Turn":
+                transcript = msg.get('transcript', '')
+                formatted = msg.get('turn_is_formatted', False)
+                confidence = msg.get('confidence', 1.0)
+                
+                if transcript and transcript.strip():
+                    self._last_transcript = transcript.strip()
+                    
+                    event_type = SpeechEventType.FINAL if formatted else SpeechEventType.INTERIM
+                    
+                    response = STTResponse(
+                        event_type=event_type,
+                        data=SpeechData(
+                            text=transcript.strip(),
+                            confidence=confidence
+                        )
+                    )
+                    
+                    responses.append(response)
+                    
+                    if not self._is_speaking:
+                        self._is_speaking = True
+                        global_event_emitter.emit("speech_started")
+                        
+                    if formatted:
+                        self._is_speaking = False
+                        self._last_transcript = ""
+                        
+            elif msg_type == "Termination":
+                if self._last_transcript and self._is_speaking:
+                    final_response = STTResponse(
+                        event_type=SpeechEventType.FINAL,
+                        data=SpeechData(
+                            text=self._last_transcript,
+                            confidence=1.0
+                        )
+                    )
+                    responses.append(final_response)
+                    self._last_transcript = ""
+                    self._is_speaking = False
+                
+            elif msg_type == "Error":
+                error_msg = msg.get('error', 'Unknown error')
+                logger.error(f"AssemblyAI Error: {error_msg}")
 
         except Exception as e:
-            print(f"!!! ASSEMBLYAI PLUGIN FATAL ERROR: {e} ({type(e).__name__}) !!!")
-            self.emit("error", f"AssemblyAI transcription error: {e}")
+            logger.error(f"Error handling WebSocket message: {str(e)}")
+        
+        return responses
 
     def _resample_audio(self, audio_bytes: bytes) -> bytes:
-        raw_audio = np.frombuffer(audio_bytes, dtype=np.int16)
-        if raw_audio.size == 0: return b''
-        stereo_audio = raw_audio.reshape(-1, 2)
-        mono_audio = stereo_audio.astype(np.float32).mean(axis=1)
-        resampled_data = signal.resample(mono_audio, int(len(mono_audio) * self.target_sample_rate / self.input_sample_rate))
-        return resampled_data.astype(np.int16).tobytes()
-        
-    def _create_wav_in_memory(self, pcm_data: bytes) -> io.BytesIO:
-        """Creates a WAV file in memory from raw PCM data."""
-        wav_buffer = io.BytesIO()
-        with wave.open(wav_buffer, 'wb') as wf:
-            wf.setnchannels(1)  # Mono
-            wf.setsampwidth(2)  # 16-bit
-            wf.setframerate(self.target_sample_rate)
-            wf.writeframes(pcm_data)
-        wav_buffer.seek(0)
-        return wav_buffer
+        """Resample audio from input sample rate to target sample rate and convert to mono."""
+        try:
+            if not audio_bytes:
+                return b''
+
+            raw_audio = np.frombuffer(audio_bytes, dtype=np.int16)
+            if raw_audio.size == 0:
+                return b''
+
+            if raw_audio.size % 2 == 0: 
+                stereo_audio = raw_audio.reshape(-1, 2)
+                mono_audio = stereo_audio.astype(np.float32).mean(axis=1)
+            else:
+                mono_audio = raw_audio.astype(np.float32)
+
+            if self.input_sample_rate != self.target_sample_rate:
+                target_length = int(len(mono_audio) * self.target_sample_rate / self.input_sample_rate)
+                resampled_data = signal.resample(mono_audio, target_length)
+            else:
+                resampled_data = mono_audio
+
+            resampled_data = np.clip(resampled_data, -32767, 32767)
+            return resampled_data.astype(np.int16).tobytes()
+
+        except Exception as e:
+            logger.error(f"Error resampling audio: {e}")
+            return b''
 
     async def aclose(self) -> None:
-        if self._is_speaking and self._audio_buffer:
-            await self._transcribe_buffer()
-            await asyncio.sleep(1)
+        """Cleanup resources"""
+        
+        if len(self._stream_buffer) >= self._min_chunk_size and self._ws:
+            try:
+                final_chunk = bytes(self._stream_buffer)
+                await self._ws.send_bytes(final_chunk)
+            except Exception as e:
+                logger.error(f"Error sending final audio: {e}")
+        
+        if self._ws:
+            try:
+                await self._ws.send_str(json.dumps({"type": "Terminate"}))
+                await asyncio.sleep(0.5)  
+            except Exception as e:
+                logger.error(f"Error sending termination: {e}")
 
-        if self._session and not self._session.closed:
-            await self._session.close()
+        if self._ws_task:
+            self._ws_task.cancel()
+            try:
+                await self._ws_task
+            except asyncio.CancelledError:
+                pass
+            self._ws_task = None
+            
+        if self._ws:
+            await self._ws.close()
+            self._ws = None
+            
+        if self._session:
+            await self._session.close()
+            self._session = None

videosdk/plugins/assemblyai/version.py

@@ -1 +1 @@
-__version__ = "0.0.30"
+__version__ = "0.0.32"

{videosdk_plugins_assemblyai-0.0.30.dist-info → videosdk_plugins_assemblyai-0.0.32.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: videosdk-plugins-assemblyai
-Version: 0.0.30
+Version: 0.0.32
 Summary: VideoSDK Agent Framework plugin for AssemblyAI
 Author: videosdk
 License-Expression: Apache-2.0
@@ -12,7 +12,8 @@ Classifier: Topic :: Multimedia :: Sound/Audio
 Classifier: Topic :: Multimedia :: Video
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.11
-Requires-Dist: videosdk-agents>=0.0.30
+Requires-Dist: aiohttp
+Requires-Dist: videosdk-agents>=0.0.32
 Description-Content-Type: text/markdown
 
 # VideoSDK Assembly AI Plugin

videosdk_plugins_assemblyai-0.0.32.dist-info/RECORD

@@ -0,0 +1,6 @@
+videosdk/plugins/assemblyai/__init__.py,sha256=T4deawBZKrOiGrKFcfksus-wmb5rF5KY7_p6QBRd4QE,59
+videosdk/plugins/assemblyai/stt.py,sha256=kdUsBJtc0utn1bpFWU2Y7szv2PdGvx5HE75mHWQc1Ao,11873
+videosdk/plugins/assemblyai/version.py,sha256=WKDnjJM7gYpD9fIwhK2qAZICJAT2ndquQ6VcOar074Y,23
+videosdk_plugins_assemblyai-0.0.32.dist-info/METADATA,sha256=xPOmbeUdbO6dokBwmAlR605XYcM_X4t2-XFR7sa8ab4,804
+videosdk_plugins_assemblyai-0.0.32.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+videosdk_plugins_assemblyai-0.0.32.dist-info/RECORD,,

videosdk_plugins_assemblyai-0.0.30.dist-info/RECORD

@@ -1,6 +0,0 @@
-videosdk/plugins/assemblyai/__init__.py,sha256=T4deawBZKrOiGrKFcfksus-wmb5rF5KY7_p6QBRd4QE,59
-videosdk/plugins/assemblyai/stt.py,sha256=4qBQSn0gvQ4ET0ilFW6lry7NEOd_Sc1yqG6v4op6u7M,6247
-videosdk/plugins/assemblyai/version.py,sha256=8ZeepqkW4DvpVeNm92mx0tIzgvVevS4NKWkTXXHuXNY,23
-videosdk_plugins_assemblyai-0.0.30.dist-info/METADATA,sha256=1pk6-Le16OuZ_8ayfLQTQYg9nlP8HHn3bmlwiukI_A0,781
-videosdk_plugins_assemblyai-0.0.30.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-videosdk_plugins_assemblyai-0.0.30.dist-info/RECORD,,