vad_processor.py

"""
Enhanced VAD (Voice Activity Detection) processor
Provides multiple VAD options including pipecat's advanced VAD
"""

import asyncio
import logging
import numpy as np
from typing import Optional, Callable, Dict, Any
from enum import Enum

logger = logging.getLogger(__name__)

class VADType(Enum):
    DEEPGRAM = "deepgram"  # Use Deepgram's built-in VAD
    PIPECAT_SILERO = "pipecat_silero"  # Use pipecat's Silero VAD
    PIPECAT_WEBRTC = "pipecat_webrtc"  # Use pipecat's WebRTC VAD
    CUSTOM = "custom"  # Custom VAD implementation

class VADProcessor:
    """
    Voice Activity Detection processor with multiple VAD backends
    """
    
    def __init__(self, vad_type: VADType = VADType.DEEPGRAM, **kwargs):
        self.vad_type = vad_type
        self.vad_engine = None
        self.is_speaking = False
        self.speech_start_callback: Optional[Callable] = None
        self.speech_end_callback: Optional[Callable] = None
        
        # VAD configuration
        self.config = {
            "sample_rate": kwargs.get("sample_rate", 16000),
            "chunk_size": kwargs.get("chunk_size", 1600),  # 100ms at 16kHz
            "silence_duration_ms": kwargs.get("silence_duration_ms", 1000),  # 1s silence to end speech
            "speech_threshold": kwargs.get("speech_threshold", 0.5),  # Speech probability threshold
            "min_speech_duration_ms": kwargs.get("min_speech_duration_ms", 250),  # Min speech duration
        }
        
        logger.info(f"🎤 Initializing VAD processor: {vad_type.value}")
        
    async def initialize(self):
        """Initialize the VAD engine"""
        try:
            if self.vad_type == VADType.PIPECAT_SILERO:
                await self._init_pipecat_silero()
            elif self.vad_type == VADType.PIPECAT_WEBRTC:
                await self._init_pipecat_webrtc()
            elif self.vad_type == VADType.CUSTOM:
                await self._init_custom_vad()
            else:
                # DEEPGRAM - no initialization needed, handled by Deepgram service
                logger.info("🎙️ Using Deepgram built-in VAD")
                
        except Exception as e:
            logger.error(f"❌ Failed to initialize VAD: {e}")
            logger.info("🔄 Falling back to Deepgram VAD")
            self.vad_type = VADType.DEEPGRAM
    
    async def _init_pipecat_silero(self):
        """Initialize pipecat Silero VAD"""
        try:
            from pipecat.vad.silero import SileroVADAnalyzer
            
            self.vad_engine = SileroVADAnalyzer(
                sample_rate=self.config["sample_rate"],
                min_volume=0.6,
                min_silence_duration=self.config["silence_duration_ms"] / 1000.0,
                min_speech_duration=self.config["min_speech_duration_ms"] / 1000.0,
            )
            logger.info("✅ Pipecat Silero VAD initialized")
            
        except ImportError:
            logger.warning("⚠️ Pipecat Silero VAD not available, falling back to Deepgram")
            self.vad_type = VADType.DEEPGRAM
        except Exception as e:
            logger.error(f"❌ Silero VAD initialization failed: {e}")
            self.vad_type = VADType.DEEPGRAM
    
    async def _init_pipecat_webrtc(self):
        """Initialize pipecat WebRTC VAD"""
        try:
            from pipecat.vad.webrtc import WebRTCVADAnalyzer
            
            self.vad_engine = WebRTCVADAnalyzer(
                sample_rate=self.config["sample_rate"],
                aggressiveness=2,  # 0-3, higher = more aggressive
            )
            logger.info("✅ Pipecat WebRTC VAD initialized")
            
        except ImportError:
            logger.warning("⚠️ Pipecat WebRTC VAD not available, falling back to Deepgram")
            self.vad_type = VADType.DEEPGRAM
        except Exception as e:
            logger.error(f"❌ WebRTC VAD initialization failed: {e}")
            self.vad_type = VADType.DEEPGRAM
    
    async def _init_custom_vad(self):
        """Initialize custom VAD implementation"""
        logger.info("✅ Custom VAD initialized (energy-based)")
        # Custom VAD will use energy-based detection
        
    def set_callbacks(self, speech_start: Callable = None, speech_end: Callable = None):
        """Set callbacks for speech start/end events"""
        self.speech_start_callback = speech_start
        self.speech_end_callback = speech_end
    
    async def process_audio(self, audio_data: bytes) -> Dict[str, Any]:
        """
        Process audio data and return VAD results
        
        Args:
            audio_data: Raw audio bytes (16-bit linear PCM)
            
        Returns:
            Dict with VAD results: {
                "is_speech": bool,
                "confidence": float,
                "speech_start": bool,
                "speech_end": bool,
                "energy": float
            }
        """
        try:
            if self.vad_type == VADType.DEEPGRAM:
                # For Deepgram VAD, we rely on the service's vad_events
                return await self._process_deepgram_vad(audio_data)
            
            elif self.vad_type in [VADType.PIPECAT_SILERO, VADType.PIPECAT_WEBRTC]:
                return await self._process_pipecat_vad(audio_data)
            
            elif self.vad_type == VADType.CUSTOM:
                return await self._process_custom_vad(audio_data)
                
        except Exception as e:
            logger.error(f"VAD processing error: {e}")
            return {
                "is_speech": False,
                "confidence": 0.0,
                "speech_start": False,
                "speech_end": False,
                "energy": 0.0
            }
    
    async def _process_deepgram_vad(self, audio_data: bytes) -> Dict[str, Any]:
        """Process audio using Deepgram's built-in VAD"""
        # For Deepgram VAD, we rely on the WebSocket vad_events
        # This is just a placeholder - actual VAD happens in Deepgram service
        energy = self._calculate_energy(audio_data)
        
        return {
            "is_speech": energy > 0.01,  # Simple energy threshold
            "confidence": min(energy * 10, 1.0),
            "speech_start": False,  # Handled by Deepgram vad_events
            "speech_end": False,    # Handled by Deepgram vad_events  
            "energy": energy
        }
    
    async def _process_pipecat_vad(self, audio_data: bytes) -> Dict[str, Any]:
        """Process audio using pipecat VAD engines"""
        if not self.vad_engine:
            return await self._process_deepgram_vad(audio_data)
        
        try:
            # Convert bytes to numpy array
            audio_np = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
            
            # Process with pipecat VAD
            vad_result = await self.vad_engine.analyze_audio(audio_np)
            
            # Check for speech state changes
            was_speaking = self.is_speaking
            self.is_speaking = vad_result.get("is_speech", False)
            
            speech_start = not was_speaking and self.is_speaking
            speech_end = was_speaking and not self.is_speaking
            
            # Trigger callbacks
            if speech_start and self.speech_start_callback:
                await self.speech_start_callback()
            if speech_end and self.speech_end_callback:
                await self.speech_end_callback()
            
            return {
                "is_speech": self.is_speaking,
                "confidence": vad_result.get("confidence", 0.5),
                "speech_start": speech_start,
                "speech_end": speech_end,
                "energy": self._calculate_energy(audio_data)
            }
            
        except Exception as e:
            logger.error(f"Pipecat VAD error: {e}")
            return await self._process_deepgram_vad(audio_data)
    
    async def _process_custom_vad(self, audio_data: bytes) -> Dict[str, Any]:
        """Process audio using custom energy-based VAD"""
        energy = self._calculate_energy(audio_data)
        
        # Simple energy-based VAD with hysteresis
        speech_threshold_high = 0.02  # Threshold to start speech
        speech_threshold_low = 0.01   # Threshold to continue speech
        
        was_speaking = self.is_speaking
        
        if not self.is_speaking:
            # Not currently speaking - check if speech starts
            self.is_speaking = energy > speech_threshold_high
        else:
            # Currently speaking - check if speech continues
            self.is_speaking = energy > speech_threshold_low
        
        speech_start = not was_speaking and self.is_speaking
        speech_end = was_speaking and not self.is_speaking
        
        # Trigger callbacks
        if speech_start and self.speech_start_callback:
            await self.speech_start_callback()
        if speech_end and self.speech_end_callback:
            await self.speech_end_callback()
        
        return {
            "is_speech": self.is_speaking,
            "confidence": min(energy * 50, 1.0),  # Convert energy to confidence
            "speech_start": speech_start,
            "speech_end": speech_end,
            "energy": energy
        }
    
    def _calculate_energy(self, audio_data: bytes) -> float:
        """Calculate RMS energy of audio data"""
        try:
            if len(audio_data) == 0:
                return 0.0
                
            # Convert to numpy array
            audio_np = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32)
            
            # Calculate RMS energy
            rms = np.sqrt(np.mean(audio_np ** 2)) / 32768.0  # Normalize to [0, 1]
            
            return float(rms)
            
        except Exception as e:
            logger.error(f"Energy calculation error: {e}")
            return 0.0
    
    async def cleanup(self):
        """Cleanup VAD resources"""
        if self.vad_engine and hasattr(self.vad_engine, 'cleanup'):
            await self.vad_engine.cleanup()
        
        logger.info("🧹 VAD processor cleaned up")

# Factory function for easy VAD creation
async def create_vad_processor(vad_type: str = "deepgram", **kwargs) -> VADProcessor:
    """
    Create and initialize a VAD processor
    
    Args:
        vad_type: Type of VAD ("deepgram", "pipecat_silero", "pipecat_webrtc", "custom")
        **kwargs: Additional configuration options
        
    Returns:
        Initialized VADProcessor instance
    """
    try:
        vad_enum = VADType(vad_type.lower())
    except ValueError:
        logger.warning(f"Unknown VAD type '{vad_type}', using Deepgram")
        vad_enum = VADType.DEEPGRAM
    
    processor = VADProcessor(vad_enum, **kwargs)
    await processor.initialize()
    
    return processor