orbitalsai 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

orbitalsai/streaming/async_client.py

@@ -0,0 +1,507 @@
+"""
+OrbitalsAI Async Streaming Client
+
+Asynchronous WebSocket client for real-time streaming transcription.
+"""
+
+import asyncio
+import json
+import logging
+from typing import Any, Dict, Optional
+
+try:
+    import websockets
+    from websockets.exceptions import (
+        ConnectionClosed,
+        ConnectionClosedOK,
+        ConnectionClosedError,
+        InvalidStatusCode,
+    )
+except ImportError:
+    raise ImportError(
+        "websockets is required for streaming. "
+        "Install it with: pip install websockets>=11.0.0"
+    )
+
+from .config import StreamingConfig
+from .events import StreamingEventHandlers
+from .exceptions import (
+    StreamingError,
+    ConnectionError,
+    AuthenticationError,
+    AudioFormatError,
+    ReconnectionFailedError,
+    SessionClosedError,
+    ProtocolError,
+    exception_from_close_code,
+)
+from .protocol import (
+    MessageType,
+    MessageParser,
+    create_websocket_url,
+    should_retry,
+    get_close_reason,
+)
+from .connection import ConnectionManager, ConnectionState
+from .audio import AudioBuffer
+
+logger = logging.getLogger("orbitalsai.streaming")
+
+
+class AsyncStreamingClient:
+    """
+    Asynchronous WebSocket client for streaming transcription.
+    
+    Provides real-time audio streaming and transcription via WebSocket.
+    Supports automatic reconnection, event callbacks, and graceful shutdown.
+    
+    Example:
+        async with AsyncStreamingClient(api_key="your_key") as client:
+            await client.connect(PrintingEventHandlers())
+            
+            # Stream audio file
+            with open("audio.pcm", "rb") as f:
+                while chunk := f.read(16000):
+                    await client.send_audio(chunk)
+            
+            await client.flush()
+    
+    Example with custom handlers:
+        class MyHandlers(StreamingEventHandlers):
+            def on_transcript_final(self, text, metadata):
+                print(f"Transcription: {text}")
+        
+        client = AsyncStreamingClient(api_key="your_key")
+        await client.connect(MyHandlers())
+        # ... send audio ...
+        await client.disconnect()
+    """
+    
+    def __init__(
+        self,
+        api_key: str,
+        config: Optional[StreamingConfig] = None,
+        base_url: str = "wss://api.orbitalsai.com"
+    ):
+        """
+        Initialize the async streaming client.
+        
+        Args:
+            api_key: OrbitalsAI API key or JWT token
+            config: Streaming configuration (optional)
+            base_url: WebSocket base URL (default: wss://api.orbitalsai.com)
+        """
+        self.api_key = api_key
+        self.config = config or StreamingConfig()
+        self.base_url = base_url
+        
+        self._ws: Optional[websockets.WebSocketClientProtocol] = None
+        self._handlers: Optional[StreamingEventHandlers] = None
+        self._receiver_task: Optional[asyncio.Task] = None
+        self._keepalive_task: Optional[asyncio.Task] = None
+        self._connected = False
+        self._session_id: Optional[str] = None
+        self._audio_buffer: Optional[AudioBuffer] = None
+        
+        # Connection management
+        self._connection_manager = ConnectionManager(
+            max_retries=self.config.max_retries,
+            base_delay=self.config.retry_delay,
+            max_delay=self.config.max_retry_delay,
+            connection_timeout=self.config.connection_timeout,
+        )
+        
+        # Reconnection state
+        self._should_reconnect = True
+        self._reconnect_lock = asyncio.Lock()
+    
+    @property
+    def is_connected(self) -> bool:
+        """Check if WebSocket is connected."""
+        return self._connected and self._ws is not None
+    
+    @property
+    def session_id(self) -> Optional[str]:
+        """Get current session ID."""
+        return self._session_id
+    
+    async def connect(self, handlers: StreamingEventHandlers) -> None:
+        """
+        Establish WebSocket connection and start receiver loop.
+        
+        Args:
+            handlers: Event handlers for callbacks
+            
+        Raises:
+            ConnectionError: If connection fails
+            AuthenticationError: If authentication fails
+        """
+        self._handlers = handlers
+        self._should_reconnect = True
+        
+        await self._connect_internal()
+    
+    async def _connect_internal(self) -> None:
+        """Internal connection method."""
+        self._connection_manager.mark_connecting()
+        
+        # Build WebSocket URL with API key
+        ws_url = create_websocket_url(self.base_url, self.api_key)
+        
+        logger.info(f"Connecting to {self.base_url}...")
+        
+        try:
+            # Connect with timeout
+            self._ws = await asyncio.wait_for(
+                websockets.connect(
+                    ws_url,
+                    ping_interval=self.config.keepalive_interval,
+                    ping_timeout=20,
+                    close_timeout=10,
+                    max_size=10 * 1024 * 1024,  # 10MB max message
+                ),
+                timeout=self.config.connection_timeout
+            )
+            
+            # Wait for ready message
+            ready_msg = await asyncio.wait_for(
+                self._ws.recv(),
+                timeout=self.config.connection_timeout
+            )
+            
+            # Parse ready message
+            try:
+                message = MessageParser.parse(ready_msg)
+            except ValueError as e:
+                raise ProtocolError(f"Invalid ready message: {e}")
+            
+            if message.type == MessageType.ERROR:
+                error_msg = message.error_message or "Connection error"
+                raise ConnectionError(error_msg)
+            
+            if message.type != MessageType.READY:
+                raise ProtocolError(f"Expected 'ready' message, got '{message.type.value}'")
+            
+            # Extract session info
+            self._session_id = message.session_id
+            self._connected = True
+            self._connection_manager.mark_connected()
+            
+            # Initialize audio buffer
+            self._audio_buffer = AudioBuffer(
+                chunk_size=self.config.chunk_size,
+                sample_rate=self.config.sample_rate
+            )
+            
+            logger.info(f"Connected: session_id={self._session_id}")
+            
+            # Call on_open handler
+            self._safe_call_handler(
+                "on_open",
+                {
+                    "session_id": message.session_id,
+                    "language": message.language,
+                    "supported_languages": message.supported_languages,
+                }
+            )
+            
+            # Start receiver loop
+            self._receiver_task = asyncio.create_task(self._receiver_loop())
+            
+        except asyncio.TimeoutError:
+            raise ConnectionError("Connection timeout")
+        except InvalidStatusCode as e:
+            if e.status_code == 401:
+                raise AuthenticationError("Invalid API key")
+            raise ConnectionError(f"Connection failed: HTTP {e.status_code}")
+        except ConnectionClosedError as e:
+            self._connection_manager.record_close(e.code, e.reason)
+            raise exception_from_close_code(e.code, e.reason)
+        except Exception as e:
+            if not isinstance(e, StreamingError):
+                raise ConnectionError(f"Connection failed: {e}")
+            raise
+    
+    async def send_audio(self, audio_data: bytes) -> None:
+        """
+        Send PCM16 audio chunk.
+        
+        Audio should be PCM16 mono little-endian format. The audio will be
+        buffered and sent in optimal chunk sizes.
+        
+        Args:
+            audio_data: Raw PCM16 mono little-endian bytes
+            
+        Raises:
+            SessionClosedError: If session is closed
+            AudioFormatError: If audio format is invalid
+        """
+        if not self.is_connected:
+            raise SessionClosedError("Session is not connected")
+        
+        # Validate audio format
+        if len(audio_data) % 2 != 0:
+            raise AudioFormatError(
+                f"Audio data must have even length (PCM16), got {len(audio_data)}"
+            )
+        
+        try:
+            # Send directly without additional buffering
+            # The server handles its own buffering
+            await self._ws.send(audio_data)
+        except ConnectionClosed as e:
+            self._handle_connection_closed(e)
+            raise SessionClosedError("Connection closed while sending audio")
+    
+    async def configure(
+        self,
+        language: Optional[str] = None,
+        sample_rate: Optional[int] = None
+    ) -> None:
+        """
+        Update session configuration dynamically.
+        
+        Args:
+            language: New transcription language (optional)
+            sample_rate: New sample rate in Hz (optional)
+            
+        Raises:
+            SessionClosedError: If session is closed
+        """
+        if not self.is_connected:
+            raise SessionClosedError("Session is not connected")
+        
+        config_msg = MessageParser.encode_config(
+            language=language,
+            sample_rate=sample_rate
+        )
+        
+        try:
+            await self._ws.send(config_msg)
+        except ConnectionClosed as e:
+            self._handle_connection_closed(e)
+            raise SessionClosedError("Connection closed while configuring")
+    
+    async def flush(self) -> None:
+        """
+        Force transcription of remaining audio buffer.
+        
+        Sends a flush command to process any buffered audio on the server.
+        The server will emit final transcripts for any remaining audio.
+        
+        Raises:
+            SessionClosedError: If session is closed
+        """
+        if not self.is_connected:
+            raise SessionClosedError("Session is not connected")
+        
+        # Flush local buffer first
+        if self._audio_buffer:
+            remaining = self._audio_buffer.flush()
+            if remaining:
+                try:
+                    await self._ws.send(remaining)
+                except ConnectionClosed:
+                    pass  # Will be handled by server flush
+        
+        try:
+            flush_msg = MessageParser.encode_flush()
+            await self._ws.send(flush_msg)
+        except ConnectionClosed as e:
+            self._handle_connection_closed(e)
+            raise SessionClosedError("Connection closed while flushing")
+    
+    async def disconnect(self) -> None:
+        """
+        Close connection gracefully.
+        
+        Stops the receiver loop and closes the WebSocket connection.
+        """
+        self._should_reconnect = False
+        self._connected = False
+        
+        # Cancel receiver task
+        if self._receiver_task and not self._receiver_task.done():
+            self._receiver_task.cancel()
+            try:
+                await self._receiver_task
+            except asyncio.CancelledError:
+                pass
+        
+        # Cancel keepalive task
+        if self._keepalive_task and not self._keepalive_task.done():
+            self._keepalive_task.cancel()
+            try:
+                await self._keepalive_task
+            except asyncio.CancelledError:
+                pass
+        
+        # Close WebSocket
+        if self._ws:
+            try:
+                await self._ws.close(code=1000, reason="Client disconnect")
+            except Exception as e:
+                logger.debug(f"Error closing WebSocket: {e}")
+            self._ws = None
+        
+        self._connection_manager.mark_closed()
+        logger.info("Disconnected")
+    
+    async def _receiver_loop(self) -> None:
+        """Background task to receive and dispatch messages."""
+        try:
+            async for message in self._ws:
+                await self._handle_message(message)
+        except ConnectionClosedOK:
+            logger.info("Connection closed normally")
+            self._safe_call_handler("on_close", 1000, "Normal closure")
+        except ConnectionClosedError as e:
+            logger.warning(f"Connection closed: code={e.code}, reason={e.reason}")
+            self._handle_connection_closed(e)
+        except asyncio.CancelledError:
+            logger.debug("Receiver loop cancelled")
+            raise
+        except Exception as e:
+            logger.exception(f"Receiver loop error: {e}")
+            self._safe_call_handler("on_error", e)
+        finally:
+            self._connected = False
+    
+    async def _handle_message(self, raw_message: str) -> None:
+        """
+        Handle a received WebSocket message.
+        
+        Args:
+            raw_message: Raw message string from WebSocket
+        """
+        try:
+            message = MessageParser.parse(raw_message)
+        except ValueError as e:
+            logger.warning(f"Invalid message: {e}")
+            self._safe_call_handler("on_error", ProtocolError(str(e)))
+            return
+        
+        # Dispatch based on message type
+        msg_type = message.type
+        
+        if msg_type == MessageType.PARTIAL:
+            self._safe_call_handler("on_transcript_partial", message.text or "")
+        
+        elif msg_type == MessageType.FINAL:
+            metadata = {
+                "cost": message.cost or 0,
+                "audio_seconds": message.audio_seconds or 0,
+                "remaining_percent": message.remaining_percent or 100,
+                "capped": message.capped,
+            }
+            self._safe_call_handler("on_transcript_final", message.text or "", metadata)
+        
+        elif msg_type == MessageType.SPEECH_START:
+            self._safe_call_handler("on_speech_start")
+        
+        elif msg_type == MessageType.SPEECH_END:
+            self._safe_call_handler("on_speech_end")
+        
+        elif msg_type == MessageType.LANGUAGE_SET:
+            self._safe_call_handler("on_language_detected", message.language or "")
+        
+        elif msg_type == MessageType.SAMPLE_RATE_SET:
+            self._safe_call_handler("on_sample_rate_changed", message.sample_rate or 16000)
+        
+        elif msg_type == MessageType.FLUSHED:
+            self._safe_call_handler("on_flushed")
+        
+        elif msg_type == MessageType.CREDITS_WARNING:
+            self._safe_call_handler("on_credits_warning", message.remaining_percent or 0)
+        
+        elif msg_type == MessageType.CREDITS_CRITICAL:
+            self._safe_call_handler("on_credits_critical", message.remaining_percent or 0)
+        
+        elif msg_type == MessageType.CREDITS_EXHAUSTED:
+            self._safe_call_handler("on_credits_exhausted")
+        
+        elif msg_type == MessageType.ERROR:
+            error = StreamingError(message.error_message or "Unknown error")
+            self._safe_call_handler("on_error", error)
+    
+    def _safe_call_handler(self, handler_name: str, *args) -> None:
+        """
+        Safely call an event handler.
+        
+        Wraps handler calls in try/except to prevent user code from
+        crashing the SDK.
+        
+        Args:
+            handler_name: Name of handler method
+            *args: Arguments to pass to handler
+        """
+        if not self._handlers:
+            return
+        
+        handler = getattr(self._handlers, handler_name, None)
+        if not handler:
+            return
+        
+        try:
+            handler(*args)
+        except Exception as e:
+            logger.warning(f"Handler {handler_name} raised exception: {e}")
+    
+    def _handle_connection_closed(self, error: ConnectionClosed) -> None:
+        """
+        Handle connection closure.
+        
+        Args:
+            error: The connection closed error
+        """
+        self._connected = False
+        self._connection_manager.record_close(error.code, error.reason)
+        
+        # Call on_close handler
+        self._safe_call_handler("on_close", error.code, error.reason)
+        
+        # Check for non-retryable errors
+        if not should_retry(error.code):
+            return
+        
+        # Attempt reconnection if enabled
+        if self._should_reconnect:
+            asyncio.create_task(self._attempt_reconnect())
+    
+    async def _attempt_reconnect(self) -> None:
+        """Attempt to reconnect after connection loss."""
+        async with self._reconnect_lock:
+            if not self._should_reconnect:
+                return
+            
+            while self._should_reconnect and self._connection_manager.should_reconnect():
+                self._connection_manager.increment_retry()
+                logger.info(
+                    f"Reconnection attempt {self._connection_manager.retry_count}/"
+                    f"{self.config.max_retries}"
+                )
+                
+                await self._connection_manager.wait_for_retry()
+                
+                try:
+                    await self._connect_internal()
+                    logger.info("Reconnected successfully")
+                    return
+                except Exception as e:
+                    logger.warning(f"Reconnection failed: {e}")
+            
+            # Failed to reconnect
+            if self._should_reconnect:
+                error = ReconnectionFailedError(
+                    f"Failed to reconnect after {self._connection_manager.retry_count} attempts",
+                    attempts=self._connection_manager.retry_count
+                )
+                self._safe_call_handler("on_error", error)
+    
+    async def __aenter__(self) -> 'AsyncStreamingClient':
+        """Async context manager entry."""
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Async context manager exit."""
+        await self.disconnect()

orbitalsai/streaming/audio/__init__.py

@@ -0,0 +1,33 @@
+"""
+OrbitalsAI Streaming Audio Utilities
+
+Audio processing utilities for streaming transcription.
+"""
+
+from .buffer import AudioBuffer
+from .converter import AudioConverter
+from .formats import (
+    AudioFormat,
+    PCM16_MONO,
+    get_format_for_file,
+    SUPPORTED_AUDIO_EXTENSIONS,
+)
+from .source import (
+    AudioSource,
+    FileAudioSource,
+    RawPCMFileSource,
+    MicrophoneSource,
+)
+
+__all__ = [
+    "AudioBuffer",
+    "AudioConverter",
+    "AudioFormat",
+    "PCM16_MONO",
+    "get_format_for_file",
+    "SUPPORTED_AUDIO_EXTENSIONS",
+    "AudioSource",
+    "FileAudioSource",
+    "RawPCMFileSource",
+    "MicrophoneSource",
+]

orbitalsai/streaming/audio/buffer.py

@@ -0,0 +1,171 @@
+"""
+OrbitalsAI Streaming Audio Buffer
+
+Audio chunking buffer for streaming transcription.
+"""
+
+from typing import List, Optional
+import logging
+
+logger = logging.getLogger("orbitalsai.streaming")
+
+
+class AudioBuffer:
+    """
+    Manages audio chunking for streaming.
+    
+    Accumulates audio data and returns complete chunks of the specified size.
+    Any remaining data is kept in the buffer until more audio is added or
+    the buffer is flushed.
+    
+    Example:
+        buffer = AudioBuffer(chunk_size=8000, sample_rate=16000)
+        
+        # Add audio and get complete chunks
+        chunks = buffer.add(audio_bytes)
+        for chunk in chunks:
+            await client.send_audio(chunk)
+        
+        # Flush remaining audio at the end
+        remaining = buffer.flush()
+        if remaining:
+            await client.send_audio(remaining)
+    """
+    
+    def __init__(self, chunk_size: int, sample_rate: int):
+        """
+        Initialize audio buffer.
+        
+        Args:
+            chunk_size: Number of samples per chunk
+            sample_rate: Audio sample rate in Hz
+        """
+        self.chunk_size = chunk_size
+        self.sample_rate = sample_rate
+        self._buffer = bytearray()
+        self._bytes_per_chunk = chunk_size * 2  # 2 bytes per int16 sample
+        self._total_samples_processed = 0
+    
+    def add(self, audio_data: bytes) -> List[bytes]:
+        """
+        Add audio data and return complete chunks.
+        
+        Adds audio to the internal buffer and extracts any complete chunks.
+        Remaining data stays in the buffer for the next call.
+        
+        Args:
+            audio_data: Raw PCM16 mono little-endian bytes
+            
+        Returns:
+            List of complete chunk bytes (each bytes_per_chunk size)
+        """
+        if not audio_data:
+            return []
+        
+        # Ensure even length (PCM16 = 2 bytes per sample)
+        if len(audio_data) % 2 != 0:
+            logger.warning(
+                f"Audio data has odd length ({len(audio_data)}), "
+                f"truncating last byte"
+            )
+            audio_data = audio_data[:-1]
+        
+        self._buffer.extend(audio_data)
+        
+        chunks = []
+        while len(self._buffer) >= self._bytes_per_chunk:
+            chunk = bytes(self._buffer[:self._bytes_per_chunk])
+            self._buffer = self._buffer[self._bytes_per_chunk:]
+            chunks.append(chunk)
+            self._total_samples_processed += self.chunk_size
+        
+        return chunks
+    
+    def flush(self) -> Optional[bytes]:
+        """
+        Flush any remaining audio in the buffer.
+        
+        Returns the remaining audio data (may be smaller than chunk_size).
+        After flushing, the buffer is empty.
+        
+        Returns:
+            Remaining audio bytes, or None if buffer is empty
+        """
+        if not self._buffer:
+            return None
+        
+        remaining = bytes(self._buffer)
+        samples = len(remaining) // 2
+        self._total_samples_processed += samples
+        self._buffer = bytearray()
+        
+        return remaining
+    
+    def clear(self) -> None:
+        """Clear the buffer without returning data."""
+        self._buffer = bytearray()
+    
+    @property
+    def buffered_duration_ms(self) -> float:
+        """
+        Get duration of buffered audio in milliseconds.
+        
+        Returns:
+            Buffered audio duration in milliseconds
+        """
+        samples = len(self._buffer) // 2
+        return (samples / self.sample_rate) * 1000
+    
+    @property
+    def buffered_samples(self) -> int:
+        """
+        Get number of buffered samples.
+        
+        Returns:
+            Number of samples in buffer
+        """
+        return len(self._buffer) // 2
+    
+    @property
+    def buffered_bytes(self) -> int:
+        """
+        Get number of buffered bytes.
+        
+        Returns:
+            Number of bytes in buffer
+        """
+        return len(self._buffer)
+    
+    @property
+    def is_empty(self) -> bool:
+        """
+        Check if buffer is empty.
+        
+        Returns:
+            True if buffer is empty, False otherwise
+        """
+        return len(self._buffer) == 0
+    
+    @property
+    def total_duration_ms(self) -> float:
+        """
+        Get total duration of audio processed (including flushed).
+        
+        Returns:
+            Total processed duration in milliseconds
+        """
+        return (self._total_samples_processed / self.sample_rate) * 1000
+    
+    @property
+    def total_samples(self) -> int:
+        """
+        Get total number of samples processed.
+        
+        Returns:
+            Total number of samples processed
+        """
+        return self._total_samples_processed
+    
+    def reset_stats(self) -> None:
+        """Reset statistics (total_samples_processed)."""
+        self._total_samples_processed = 0