server.py

#!/usr/bin/env python3
"""
Voice Control Server for OpenClaw
"""

import asyncio
import json
import logging
import numpy as np
import websockets
import time
from collections import deque
from openwakeword import Model as WakeWordModel
import whisper

# Configuration
WEBSOCKET_HOST = "0.0.0.0"
WEBSOCKET_PORT = 8765
SAMPLE_RATE = 16000
CHANNELS = 1

# Wake word detection - testing all models to find best match for "Andy"
WAKE_WORD_MODELS = [
    "alexa_v0.1",
    "hey_jarvis_v0.1",
    "hey_mycroft_v0.1",  # Might respond to "Hey Andy"
    "hey_rhasspy_v0.1",
]
WAKE_WORD_THRESHOLD = 0.05  # Low threshold for testing
SILENCE_THRESHOLD = 300  # Lower = more sensitive to silence
SILENCE_DURATION = 2.0   # Wait 2 seconds of silence before stopping
MAX_RECORDING_DURATION = 30

recent_logs = deque(maxlen=100)

# Keep last 5 seconds of audio for playback (16000 samples/sec * 5 = 80000 samples)
PLAYBACK_BUFFER_SECONDS = 5
audio_buffer = deque(maxlen=SAMPLE_RATE * PLAYBACK_BUFFER_SECONDS)

class LogHandler(logging.Handler):
    def emit(self, record):
        log_entry = self.format(record)
        recent_logs.append(log_entry)

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
logger.addHandler(LogHandler())

class VoiceControlServer:
    def __init__(self):
        logger.info("Initializing Voice Control Server...")
        
        logger.info(f"Loading wake word models: {WAKE_WORD_MODELS}")
        self.wake_word_model = WakeWordModel(
            wakeword_models=WAKE_WORD_MODELS,
            inference_framework='onnx'
        )
        logger.info(f"Loaded models: {list(self.wake_word_model.models.keys())}")
        
        logger.info("Loading Whisper model (medium) for better Hebrew...")
        self.whisper_model = whisper.load_model("medium")
        
        logger.info("Server initialized successfully!")
        self.max_score_seen = 0.0
        
    def calculate_rms(self, audio_chunk: np.ndarray) -> float:
        return float(np.sqrt(np.mean(audio_chunk.astype(np.float32) ** 2)))
    
    def detect_wake_word(self, audio_chunk: np.ndarray) -> tuple:
        prediction = self.wake_word_model.predict(audio_chunk)
        
        # Check all wake words, return highest score
        best_word = None
        best_score = 0.0
        for model_name, score in prediction.items():
            score = float(score)
            if score > best_score:
                best_score = score
                best_word = model_name
        
        # Track max score for debugging
        if best_score > self.max_score_seen:
            self.max_score_seen = best_score
            logger.info(f"📈 New max score: {best_score:.4f} for '{best_word}'")
        
        detected = best_score > WAKE_WORD_THRESHOLD
        return detected, best_score, best_word, prediction
    
    def transcribe(self, audio_data: np.ndarray) -> str:
        audio_float = audio_data.astype(np.float32) / 32768.0
        result = self.whisper_model.transcribe(
            audio_float,
            language="he",  # Hebrew
            fp16=False
        )
        return result["text"].strip()
    
    def check_for_andy(self, audio_data: np.ndarray) -> bool:
        """Quick check if 'Andy' was said in recent audio"""
        audio_float = audio_data.astype(np.float32) / 32768.0
        result = self.whisper_model.transcribe(
            audio_float,
            language="en",  # Andy is English
            fp16=False,
            condition_on_previous_text=False
        )
        text = result["text"].lower()
        # Check for Andy or similar sounds
        andy_variants = ["andy", "andi", "andie", "and he", "handy", "אנדי"]
        for variant in andy_variants:
            if variant in text:
                logger.info(f"🎯 Detected 'Andy' in: '{result['text']}'")
                return True
        return False
    
    async def handle_client(self, websocket):
        client_id = id(websocket)
        logger.info(f"Client {client_id} connected from {websocket.remote_address}")
        
        # Reset model state for new client
        self.wake_word_model.reset()
        self.max_score_seen = 0.0
        
        state = "listening"
        recording_buffer = []
        ptt_recording_buffer = []
        ptt_active = False
        silence_frames = 0
        chunk_count = 0
        last_score_log = time.time()
        
        try:
            async for message in websocket:
                if isinstance(message, bytes):
                    # Verify audio format
                    if len(message) != 3200:  # 1600 samples * 2 bytes
                        logger.warning(f"Unexpected audio chunk size: {len(message)} bytes")
                    
                    audio_chunk = np.frombuffer(message, dtype=np.int16)
                    chunk_count += 1
                    rms = self.calculate_rms(audio_chunk)
                    
                    # Store in playback buffer
                    audio_buffer.extend(audio_chunk.tolist())
                    
                    # Log audio stats occasionally to debug
                    if chunk_count == 10:
                        logger.info(f"📊 Audio sample stats: min={audio_chunk.min()}, max={audio_chunk.max()}, dtype={audio_chunk.dtype}")
                    
                    # Send audio level to client
                    if chunk_count % 5 == 0:
                        await websocket.send(json.dumps({
                            "type": "audio_level",
                            "level": min(100, int(rms / 100)),
                            "chunks": chunk_count
                        }))
                    
                    if ptt_active:
                        ptt_recording_buffer.append(audio_chunk)
                    elif state == "listening":
                        detected, score, best_word, all_scores = self.detect_wake_word(audio_chunk)
                        
                        # Log periodically
                        now = time.time()
                        if now - last_score_log > 2.0:
                            scores_str = ", ".join([f"{k.split('_')[0]}:{v:.4f}" for k, v in all_scores.items()])
                            logger.info(f"Scores: [{scores_str}] RMS:{rms:.0f} max_seen:{self.max_score_seen:.4f}")
                            last_score_log = now
                        
                        # Send wake word score to client
                        if chunk_count % 10 == 0:
                            await websocket.send(json.dumps({
                                "type": "wake_score",
                                "score": round(score, 4),
                                "threshold": WAKE_WORD_THRESHOLD,
                                "rms": round(rms, 1),
                                "max_seen": round(self.max_score_seen, 4),
                                "word": best_word.split('_')[0] if best_word else "none"
                            }))
                        
                        if detected:
                            logger.info(f"🎤 WAKE WORD DETECTED: '{best_word}' (score: {score:.3f})")
                            state = "recording"
                            recording_buffer = [audio_chunk]
                            silence_frames = 0
                            
                            await websocket.send(json.dumps({
                                "type": "status",
                                "status": "listening",
                                "message": f"I'm listening... ({best_word.split('_')[0]})"
                            }))
                    
                    elif state == "recording":
                        recording_buffer.append(audio_chunk)
                        
                        if rms < SILENCE_THRESHOLD:
                            silence_frames += 1
                        else:
                            silence_frames = 0
                        
                        total_duration = len(recording_buffer) * len(audio_chunk) / SAMPLE_RATE
                        silence_duration = silence_frames * len(audio_chunk) / SAMPLE_RATE
                        
                        if silence_duration >= SILENCE_DURATION or total_duration >= MAX_RECORDING_DURATION:
                            state = "processing"
                            logger.info(f"Recording complete ({total_duration:.1f}s)")
                            
                            await websocket.send(json.dumps({
                                "type": "status",
                                "status": "processing",
                                "message": "Processing..."
                            }))
                            
                            full_audio = np.concatenate(recording_buffer)
                            text = self.transcribe(full_audio)
                            logger.info(f"📝 Transcribed (Hebrew): {text}")
                            
                            if text:
                                await websocket.send(json.dumps({
                                    "type": "transcription",
                                    "text": text
                                }))
                                
                                response_text = f"I heard: {text}"
                                await websocket.send(json.dumps({
                                    "type": "response",
                                    "text": response_text
                                }))
                            
                            # Reset model state after processing
                            self.wake_word_model.reset()
                            
                            state = "listening"
                            recording_buffer = []
                            silence_frames = 0
                            
                            await websocket.send(json.dumps({
                                "type": "status",
                                "status": "idle",
                                "message": "מוכן - אמור 'Alexa' או 'Hey Andy'"
                            }))
                
                elif isinstance(message, str):
                    try:
                        cmd = json.loads(message)
                        if cmd.get("type") == "ping":
                            await websocket.send(json.dumps({"type": "pong"}))
                        elif cmd.get("type") == "ptt_start":
                            ptt_active = True
                            ptt_recording_buffer = []
                            await websocket.send(json.dumps({
                                "type": "status",
                                "status": "listening",
                                "message": "Push-to-talk active..."
                            }))
                        elif cmd.get("type") == "ptt_end":
                            if ptt_active and ptt_recording_buffer:
                                await websocket.send(json.dumps({
                                    "type": "status",
                                    "status": "processing",
                                    "message": "Processing..."
                                }))
                                full_audio = np.concatenate(ptt_recording_buffer)
                                text = self.transcribe(full_audio)
                                logger.info(f"📝 PTT transcribed (Hebrew): {text}")
                                if text:
                                    await websocket.send(json.dumps({
                                        "type": "transcription",
                                        "text": text
                                    }))
                                    await websocket.send(json.dumps({
                                        "type": "response",
                                        "text": f"I heard: {text}"
                                    }))
                            ptt_active = False
                            ptt_recording_buffer = []
                            await websocket.send(json.dumps({
                                "type": "status",
                                "status": "idle",
                                "message": "מוכן - לחץ על כפתור הדיבור הצף"
                            }))
                        elif cmd.get("type") == "get_logs":
                            await websocket.send(json.dumps({
                                "type": "logs",
                                "logs": list(recent_logs)
                            }))
                        elif cmd.get("type") == "playback":
                            # Send back last N seconds of audio
                            seconds = cmd.get("seconds", 3)
                            samples_needed = min(SAMPLE_RATE * seconds, len(audio_buffer))
                            if samples_needed > 0:
                                # Get last N samples
                                audio_data = list(audio_buffer)[-samples_needed:]
                                audio_array = np.array(audio_data, dtype='<i2')  # Force little-endian
                                logger.info(f"🔊 Sending playback: {len(audio_array)} samples ({len(audio_array)/SAMPLE_RATE:.1f}s), min={audio_array.min()}, max={audio_array.max()}")
                                await websocket.send(audio_array.tobytes())
                            else:
                                await websocket.send(json.dumps({
                                    "type": "error",
                                    "message": "No audio buffered yet"
                                }))
                    except json.JSONDecodeError:
                        pass
                        
        except websockets.exceptions.ConnectionClosed:
            logger.info(f"Client {client_id} disconnected")
        except Exception as e:
            logger.error(f"Error handling client {client_id}: {e}", exc_info=True)

    async def start(self):
        logger.info(f"Starting WebSocket server on ws://{WEBSOCKET_HOST}:{WEBSOCKET_PORT}")
        
        async with websockets.serve(
            self.handle_client,
            WEBSOCKET_HOST,
            WEBSOCKET_PORT,
            max_size=10 * 1024 * 1024,
            ping_interval=30,
            ping_timeout=10
        ):
            logger.info("Server is running. Say 'Alexa' or 'Hey Jarvis' to activate.")
            await asyncio.Future()

def main():
    server = VoiceControlServer()
    try:
        asyncio.run(server.start())
    except KeyboardInterrupt:
        logger.info("Server stopped.")

if __name__ == "__main__":
    main()