gladia_handler.py

"""
Gladia real-time transcription handler for Gradio integration
"""

import asyncio
import json
import os
from typing import Optional, Callable
import httpx
from websockets.asyncio.client import ClientConnection, connect
from dotenv import load_dotenv


class GladiaHandler:
    """Handle Gladia real-time transcription with WebSocket"""
    
    def __init__(self, language: str = "de", partial: bool = True):
        load_dotenv()
        self.api_key = os.getenv("GLADIA_API_KEY")
        if not self.api_key:
            raise ValueError("GLADIA_API_KEY not found in environment")
        
        self.language = language
        self.partial = partial
        self.session_id: Optional[str] = None
        self.websocket: Optional[ClientConnection] = None
        self.receiver_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        # Callback for transcript updates
        self.on_transcript: Optional[Callable[[str, bool], None]] = None
    
    async def initialize(self):
        """Initialize session and connect to WebSocket"""
        if self.is_running:
            print("⚠️ Gladia session already running")
            return
        
        config = {
            "model": "solaria-1",
            "encoding": "wav/pcm",
            "sample_rate": 16000,
            "bit_depth": 16,
            "channels": 1,
            "language_config": {
                "languages": [self.language] if self.language != "auto" else [],
                "code_switching": self.language == "auto",
            },
            "messages_config": {
                "receive_partial_transcripts": self.partial,
                "receive_final_transcripts": True,
            },
        }
        
        try:
            # Initialize session
            response = httpx.post(
                "https://api.gladia.io/v2/live",
                headers={"x-gladia-key": self.api_key},
                json=config,
                timeout=30,
            )
            
            if response.status_code not in (200, 201):
                raise Exception(f"Failed to initialize: {response.status_code} - {response.text}")
            
            session_data = response.json()
            self.session_id = session_data["id"]
            websocket_url = session_data["url"]
            
            print(f"✅ Gladia session initialized: {self.session_id}")
            
            # Connect to WebSocket
            self.websocket = await connect(websocket_url)
            self.is_running = True
            
            # Start receiver task
            self.receiver_task = asyncio.create_task(self._receive_messages())
            
            print("✅ Gladia WebSocket connected")
            
        except Exception as e:
            print(f"❌ Gladia initialization failed: {e}")
            raise
    
    async def send_audio(self, audio_data: bytes):
        """Send audio chunk to Gladia"""
        if not self.is_running or not self.websocket:
            raise Exception("Gladia session not initialized")
        
        try:
            await self.websocket.send(audio_data)
        except Exception as e:
            print(f"❌ Failed to send audio: {e}")
            raise
    
    async def _receive_messages(self):
        """Receive and process messages from Gladia WebSocket"""
        try:
            async for message in self.websocket:
                try:
                    msg = json.loads(message)
                    self._handle_message(msg)
                except json.JSONDecodeError as e:
                    print(f"⚠️ JSON decode error: {e}")
        except Exception as e:
            if self.is_running:
                print(f"❌ Receiver error: {e}")
    
    def _handle_message(self, message: dict):
        """Handle incoming message from Gladia"""
        msg_type = message.get("type")
        
        if msg_type == "transcript":
            data = message.get("data", {})
            is_final = data.get("is_final", False)
            utterance = data.get("utterance", {})
            text = utterance.get("text", "").strip()
            
            if text and self.on_transcript:
                # Call the callback with transcript and is_final flag
                self.on_transcript(text, is_final)
        
        elif msg_type == "speech_start":
            print("🎤 Speech detected")
        
        elif msg_type == "speech_end":
            print("🎤 Speech ended")
        
        elif msg_type == "post_final_transcript":
            print("✅ Final transcript received")
    
    async def stop(self):
        """Stop the Gladia session"""
        if not self.is_running:
            return
        
        self.is_running = False
        
        try:
            # Send stop recording message
            if self.websocket:
                await self.websocket.send(json.dumps({"type": "stop_recording"}))
                await asyncio.sleep(0.5)  # Give it time to process
                
                # Cancel receiver task
                if self.receiver_task:
                    self.receiver_task.cancel()
                    try:
                        await self.receiver_task
                    except asyncio.CancelledError:
                        pass
                
                # Close WebSocket
                await self.websocket.close()
                
            print(f"✅ Gladia session stopped: {self.session_id}")
            
        except Exception as e:
            print(f"⚠️ Error stopping Gladia: {e}")
        
        finally:
            self.websocket = None
            self.receiver_task = None
            self.session_id = None
    
    def set_transcript_callback(self, callback: Callable[[str, bool], None]):
        """Set callback function for transcript updates"""
        self.on_transcript = callback


# Convenience function for Gradio integration
async def create_gladia_handler(language: str = "de", partial: bool = True) -> GladiaHandler:
    """Create and initialize a Gladia handler"""
    handler = GladiaHandler(language=language, partial=partial)
    await handler.initialize()
    return handler