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

orbitalsai/streaming/protocol.py

@@ -0,0 +1,245 @@
+"""
+OrbitalsAI Streaming Protocol
+
+Message encoding/decoding for the streaming WebSocket protocol.
+"""
+
+import json
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+logger = logging.getLogger("orbitalsai.streaming")
+
+
+class MessageType(Enum):
+    """WebSocket message types for the streaming protocol."""
+    
+    # Server → Client
+    READY = "ready"
+    PARTIAL = "partial"
+    FINAL = "final"
+    SPEECH_START = "speech_start"
+    SPEECH_END = "speech_end"
+    LANGUAGE_SET = "language_set"
+    SAMPLE_RATE_SET = "sample_rate_set"
+    CREDITS_WARNING = "credits_warning"
+    CREDITS_CRITICAL = "credits_critical"
+    CREDITS_EXHAUSTED = "credits_exhausted"
+    FLUSHED = "flushed"
+    ERROR = "error"
+    
+    # Client → Server
+    CONFIG = "config"
+    FLUSH = "flush"
+    AUDIO = "audio"  # Binary, not JSON
+
+
+@dataclass
+class Message:
+    """
+    Represents a parsed WebSocket message.
+    
+    Attributes:
+        type: Message type enum
+        data: Message payload data
+        raw: Original raw message (for debugging)
+    """
+    type: MessageType
+    data: Dict[str, Any] = field(default_factory=dict)
+    raw: Optional[str] = None
+    
+    @property
+    def text(self) -> Optional[str]:
+        """Get transcript text (for partial/final messages)."""
+        return self.data.get("text")
+    
+    @property
+    def session_id(self) -> Optional[str]:
+        """Get session ID (for ready message)."""
+        return self.data.get("session_id")
+    
+    @property
+    def language(self) -> Optional[str]:
+        """Get language (for ready/language_set messages)."""
+        return self.data.get("language")
+    
+    @property
+    def sample_rate(self) -> Optional[int]:
+        """Get sample rate (for sample_rate_set message)."""
+        return self.data.get("sample_rate")
+    
+    @property
+    def supported_languages(self) -> List[str]:
+        """Get supported languages (for ready message)."""
+        return self.data.get("supported_languages", [])
+    
+    @property
+    def cost(self) -> Optional[float]:
+        """Get cost (for final message)."""
+        return self.data.get("cost")
+    
+    @property
+    def audio_seconds(self) -> Optional[float]:
+        """Get audio duration (for final message)."""
+        return self.data.get("audio_seconds")
+    
+    @property
+    def remaining_percent(self) -> Optional[int]:
+        """Get remaining credits percent (for final/credits_* messages)."""
+        return self.data.get("remaining_percent")
+    
+    @property
+    def capped(self) -> bool:
+        """Check if billing was capped (for final message)."""
+        return self.data.get("capped", False)
+    
+    @property
+    def error_message(self) -> Optional[str]:
+        """Get error message (for error message)."""
+        return self.data.get("message")
+
+
+class MessageParser:
+    """
+    Parser for streaming protocol messages.
+    
+    Handles conversion between raw WebSocket messages and Message objects.
+    """
+    
+    @staticmethod
+    def parse(raw: str) -> Message:
+        """
+        Parse a raw JSON message from the server.
+        
+        Args:
+            raw: Raw JSON string from WebSocket
+            
+        Returns:
+            Parsed Message object
+            
+        Raises:
+            ValueError: If message is invalid or has unknown type
+        """
+        try:
+            data = json.loads(raw)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON message: {e}")
+        
+        msg_type_str = data.get("type")
+        if not msg_type_str:
+            raise ValueError("Message missing 'type' field")
+        
+        try:
+            msg_type = MessageType(msg_type_str)
+        except ValueError:
+            logger.warning(f"Unknown message type: {msg_type_str}")
+            raise ValueError(f"Unknown message type: {msg_type_str}")
+        
+        return Message(type=msg_type, data=data, raw=raw)
+    
+    @staticmethod
+    def encode_config(
+        language: Optional[str] = None,
+        sample_rate: Optional[int] = None
+    ) -> str:
+        """
+        Encode a config message to send to server.
+        
+        Args:
+            language: Optional language to set
+            sample_rate: Optional sample rate to set
+            
+        Returns:
+            JSON string to send
+        """
+        msg = {"type": "config"}
+        if language is not None:
+            msg["language"] = language
+        if sample_rate is not None:
+            msg["sample_rate"] = sample_rate
+        return json.dumps(msg)
+    
+    @staticmethod
+    def encode_flush() -> str:
+        """
+        Encode a flush message to send to server.
+        
+        Returns:
+            JSON string to send
+        """
+        return json.dumps({"type": "flush"})
+
+
+def create_websocket_url(base_url: str, api_key: str) -> str:
+    """
+    Create WebSocket URL with authentication.
+    
+    Args:
+        base_url: Base URL (e.g., "wss://api.orbitalsai.com")
+        api_key: API key or JWT token
+        
+    Returns:
+        Full WebSocket URL with token parameter
+    """
+    # Ensure we have the right protocol
+    if base_url.startswith("https://"):
+        base_url = "wss://" + base_url[8:]
+    elif base_url.startswith("http://"):
+        base_url = "ws://" + base_url[7:]
+    elif not base_url.startswith(("ws://", "wss://")):
+        base_url = "wss://" + base_url
+    
+    # Remove trailing slash
+    base_url = base_url.rstrip("/")
+    
+    # Build URL with endpoint and token
+    return f"{base_url}/api/v1/streaming/transcribe?token={api_key}"
+
+
+# Close codes and their meanings
+CLOSE_CODES = {
+    1000: ("Normal closure", False),           # No retry needed
+    1001: ("Going away", True),                # Server shutdown, retry
+    1006: ("Abnormal closure", True),          # Network issue, retry
+    4001: ("Authentication failed", False),    # Don't retry
+    4002: ("Service unavailable", True),       # Retry with backoff
+    4003: ("Server busy", True),               # Retry with longer backoff
+    4004: ("Insufficient balance", False),     # Don't retry
+    4005: ("Credits exhausted", False),        # Don't retry
+}
+
+
+def should_retry(close_code: int) -> bool:
+    """
+    Determine if connection should be retried based on close code.
+    
+    Args:
+        close_code: WebSocket close code
+        
+    Returns:
+        True if connection should be retried, False otherwise
+    """
+    if close_code in CLOSE_CODES:
+        _, should_retry = CLOSE_CODES[close_code]
+        return should_retry
+    
+    # Default: retry for server errors (5xx equivalent)
+    return close_code >= 1000 and close_code < 4000
+
+
+def get_close_reason(close_code: int) -> str:
+    """
+    Get human-readable reason for close code.
+    
+    Args:
+        close_code: WebSocket close code
+        
+    Returns:
+        Human-readable reason string
+    """
+    if close_code in CLOSE_CODES:
+        reason, _ = CLOSE_CODES[close_code]
+        return reason
+    return f"Unknown close code: {close_code}"