import asyncio import base64 import json import time from typing import Dict, Any, Optional from uuid import uuid4 from fastapi import WebSocket, WebSocketDisconnect from sqlalchemy.orm import Session from app.database.database import SessionLocal from app.database.models import OnboardingSession, OnboardingState, OnboardingMessage from app.services.onboarding_service import OnboardingService from app.services.voice_service import VoiceService from app.ai.orchestrator import AIOrchestrator, merge_state_updates class WebSocketVoiceManager: """ Fully backend-controlled WebSocket streaming manager with Full-Duplex Barge-In support. 1. Automatic session initialization if session_id is 'auto' or 'new'. 2. Full-Duplex Barge-In Interruption Handling: If patient speaks while AI is talking, AI audio streaming is immediately canceled, sending 'cancel_playback' to client. 3. Automatic backend silence detection (processes turn after 1s silence). """ def __init__(self): self.active_connections: Dict[str, WebSocket] = {} self.audio_buffers: Dict[str, bytearray] = {} self.last_audio_times: Dict[str, float] = {} self.silence_tasks: Dict[str, asyncio.Task] = {} self.is_processing: Dict[str, bool] = {} self.is_ai_speaking: Dict[str, bool] = {} self.tts_tasks: Dict[str, asyncio.Task] = {} self.voice_service = VoiceService() self.onboarding_service = OnboardingService() self.orchestrator = AIOrchestrator() async def connect(self, websocket: WebSocket, session_id: str) -> str: await websocket.accept() db = SessionLocal() real_session_id = session_id try: # 1. Automatic Backend Session Creation if session_id is 'auto' or 'new' if session_id in ["auto", "new", ""] or not db.query(OnboardingSession).filter(OnboardingSession.session_id == session_id).first(): real_session_id = f"VOICE-AUTO-{uuid4().hex[:8].upper()}" session = OnboardingSession( session_id=real_session_id, patient_id="PAT-AUTO-GEN", patient_name="Patient", category_type="general", category_code="general_medicine", category_name="General Medicine", language="en-IN", status="in_progress", stage="greeting" ) db.add(session) db.commit() db.refresh(session) self.onboarding_service.get_state(db, real_session_id) self.active_connections[real_session_id] = websocket self.audio_buffers[real_session_id] = bytearray() self.last_audio_times[real_session_id] = 0.0 self.is_processing[real_session_id] = False self.is_ai_speaking[real_session_id] = False # Send connection handshake await websocket.send_json({ "type": "connected", "session_id": real_session_id, "timestamp": time.time(), "message": "Connected to Full-Duplex Voice Engine. Start speaking at any time." }) # Send initial greeting audio automatically greeting_text = self.voice_service.get_initial_greeting("Patient", "en-IN") audio_id, audio_bytes = await self.voice_service.text_to_speech(greeting_text, "en-IN") b64_greeting = base64.b64encode(audio_bytes).decode("utf-8") if audio_bytes else "" await websocket.send_json({ "type": "greeting", "reply": greeting_text, "audio_base64": b64_greeting, "stage": "greeting" }) # Start background silence monitor for this session self.silence_tasks[real_session_id] = asyncio.create_task( self._silence_monitor(websocket, real_session_id) ) return real_session_id finally: db.close() def disconnect(self, session_id: str): self._cancel_tts_task(session_id) if session_id in self.silence_tasks: self.silence_tasks[session_id].cancel() del self.silence_tasks[session_id] if session_id in self.active_connections: del self.active_connections[session_id] if session_id in self.audio_buffers: del self.audio_buffers[session_id] if session_id in self.last_audio_times: del self.last_audio_times[session_id] if session_id in self.is_processing: del self.is_processing[session_id] if session_id in self.is_ai_speaking: del self.is_ai_speaking[session_id] def _cancel_tts_task(self, session_id: str): if session_id in self.tts_tasks and not self.tts_tasks[session_id].done(): self.tts_tasks[session_id].cancel() del self.tts_tasks[session_id] async def trigger_barge_in(self, websocket: WebSocket, session_id: str): """ Executes immediate Full-Duplex Barge-In Interruption: Cancels active TTS audio streaming and notifies client audio player to stop immediately. """ if self.is_ai_speaking.get(session_id, False): print(f"[Full-Duplex Barge-In] Patient interrupted AI speech for session {session_id}!") self._cancel_tts_task(session_id) self.is_ai_speaking[session_id] = False try: await websocket.send_json({ "type": "cancel_playback", "reason": "barge_in", "timestamp": time.time(), "message": "AI playback canceled due to patient voice interruption." }) except Exception: pass async def _silence_monitor(self, websocket: WebSocket, session_id: str): """ Backend Automatic Silence Detector: Monitors incoming audio stream. If patient stops talking for >= 1.0 second, the backend automatically processes the clinical turn. """ try: while True: await asyncio.sleep(0.2) buf_len = len(self.audio_buffers.get(session_id, b"")) last_time = self.last_audio_times.get(session_id, 0.0) processing = self.is_processing.get(session_id, False) # Silence threshold: 1.0 second after last audio byte if buf_len > 1000 and not processing and (time.time() - last_time) >= 1.0: print(f"[Backend VAD] 1.0s Silence detected for {session_id}. Auto-processing turn...") await self.process_audio_turn(websocket, session_id) except asyncio.CancelledError: pass except Exception as e: print(f"[Silence Monitor Error]: {e}") async def handle_websocket_session(self, websocket: WebSocket, session_id: str): real_sid = await self.connect(websocket, session_id) try: while True: message = await websocket.receive() # 1. Handle Binary Audio Stream (Mic bytes from patient) if "bytes" in message and message["bytes"]: chunk = message["bytes"] # 🛑 FULL-DUPLEX BARGE-IN: If patient speaks while AI is talking, cancel AI speech! if self.is_ai_speaking.get(real_sid, False): await self.trigger_barge_in(websocket, real_sid) self.audio_buffers[real_sid].extend(chunk) self.last_audio_times[real_sid] = time.time() # 2. Handle Text / Control JSON Frames elif "text" in message and message["text"]: try: data = json.loads(message["text"]) msg_type = data.get("type", "") if msg_type == "ping": await websocket.send_json({"type": "pong", "timestamp": time.time()}) elif msg_type in ["end_of_speech", "finish_turn", "process"]: await self.process_audio_turn(websocket, real_sid) elif msg_type == "text_turn": if self.is_ai_speaking.get(real_sid, False): await self.trigger_barge_in(websocket, real_sid) user_text = data.get("text", "").strip() if user_text: await self.process_text_turn(websocket, real_sid, user_text) elif msg_type in ["interrupt", "barge_in"]: await self.trigger_barge_in(websocket, real_sid) except json.JSONDecodeError: pass except WebSocketDisconnect: self.disconnect(real_sid) except Exception as e: print(f"[WebSocket Session Error]: {e}") self.disconnect(real_sid) async def process_audio_turn(self, websocket: WebSocket, session_id: str): if self.is_processing.get(session_id, False): return self.is_processing[session_id] = True buffer = bytes(self.audio_buffers.get(session_id, b"")) self.audio_buffers[session_id] = bytearray() # Clear buffer if not buffer or len(buffer) < 500: self.is_processing[session_id] = False return db = SessionLocal() try: sess = db.query(OnboardingSession).filter(OnboardingSession.session_id == session_id).first() language = sess.language if sess else "en-IN" finally: db.close() await websocket.send_json({"type": "processing", "status": "transcribing"}) transcript = await self.voice_service.speech_to_text( audio_bytes=buffer, filename="stream.webm", language_code=language ) if not transcript: await websocket.send_json({"type": "warning", "message": "Unclear audio."}) self.is_processing[session_id] = False return from app.ai.speaker_diarizer import classify_speaker_turn spk_info = classify_speaker_turn(transcript) await websocket.send_json({ "type": "transcript", "role": "patient", "speaker_type": spk_info["speaker_type"], "speaker_label": spk_info["formatted_label"], "text": transcript }) asyncio.create_task(self._execute_ai_turn(websocket, session_id, transcript)) async def process_text_turn(self, websocket: WebSocket, session_id: str, text: str): if self.is_processing.get(session_id, False): return from app.ai.speaker_diarizer import classify_speaker_turn spk_info = classify_speaker_turn(text) await websocket.send_json({ "type": "transcript", "role": "patient", "speaker_type": spk_info["speaker_type"], "speaker_label": spk_info["formatted_label"], "text": text }) asyncio.create_task(self._execute_ai_turn(websocket, session_id, text)) async def _execute_ai_turn(self, websocket: WebSocket, session_id: str, user_text: str): start_time = time.time() db = SessionLocal() try: sess = db.query(OnboardingSession).filter(OnboardingSession.session_id == session_id).first() if not sess: await websocket.send_json({"type": "error", "message": f"Session {session_id} not found."}) return state_obj = db.query(OnboardingState).filter(OnboardingState.session_id == session_id).first() current_state = json.loads(state_obj.state_json) if state_obj and state_obj.state_json else {} history_objs = ( db.query(OnboardingMessage) .filter(OnboardingMessage.session_id == session_id) .order_by(OnboardingMessage.created_at.asc()) .all() ) conv_history = [{"role": m.role, "message": m.message, "stage": m.stage} for m in history_objs] await websocket.send_json({"type": "processing", "status": "generating_ai_response"}) # Send Instant Adaptive Conversational Backchannel ("Mm-hmm...", "I see...") from app.voice.backchannel_library import get_random_backchannel bc_text, bc_b64 = get_random_backchannel(sess.language or "en-IN") await websocket.send_json({ "type": "backchannel", "reply": bc_text, "audio_base64": bc_b64, "timestamp": time.time() }) # AI Orchestrator Call res = await self.orchestrator.process(sess, state_obj, conv_history, user_text) reply_text = res.get("reply", "") next_stage = res.get("next_stage", "symptom_collection") completed = res.get("completed", False) state_updates = res.get("state_updates", {}) # 🔒 HIPAA / ABDM / GDPR PII/PHI Anonymization from app.services.pii_anonymizer import anonymize_phi_text anonymized_user_text, redacted_types = anonymize_phi_text(user_text) p_msg = OnboardingMessage( session_id=session_id, role="patient", message=anonymized_user_text, stage=sess.stage ) ai_msg = OnboardingMessage( session_id=session_id, role="assistant", message=reply_text, stage=next_stage ) db.add(p_msg) db.add(ai_msg) sess.stage = next_stage if completed: sess.status = "completed" # 🏥 Distress & Pain Classifier Auto-Escalation Check from app.ai.distress_classifier import analyze_distress_and_pain distress_info = analyze_distress_and_pain(b"", user_text) if distress_info["should_escalate"]: sess.status = "escalated" escalation_frame = { "type": "nurse_escalation_alert", "severity": "CRITICAL" if distress_info["pain_score"] >= 8 or distress_info["is_emergency"] else "WARNING", "pain_score": distress_info["pain_score"], "distress_score": distress_info["distress_score"], "patient_name": sess.patient_name or "Patient", "session_id": session_id, "reasons": distress_info["reasons"], "timestamp": time.time(), "message": f"CRITICAL ALERT: Live Nurse Escalation Triggered! Pain Score: {distress_info['pain_score']}/10." } await websocket.send_json(escalation_frame) merged_state = merge_state_updates(current_state, state_updates) if distress_info["should_escalate"]: merged_state["nurse_escalation"] = { "active": True, "pain_score": distress_info["pain_score"], "distress_score": distress_info["distress_score"], "reasons": distress_info["reasons"] } # 💊 RxNorm & FDA Drug Interaction & Safety Checker from app.ai.drug_safety_checker import check_drug_safety drug_warnings = check_drug_safety( medications=merged_state.get("medications", []), allergies=merged_state.get("allergies", []), medical_history=merged_state.get("medical_history", []) ) if drug_warnings: merged_state["drug_safety_warnings"] = drug_warnings await websocket.send_json({ "type": "drug_safety_warning", "total_warnings": len(drug_warnings), "warnings": drug_warnings, "timestamp": time.time() }) if state_obj: state_obj.state_json = json.dumps(merged_state) else: db.add(OnboardingState(session_id=session_id, state_json=json.dumps(merged_state))) db.commit() # ⚡ 1. Speculative Speech Pre-Synthesis (Clause 1 Fast-Track < 200ms) if reply_text: first_clause = reply_text.split('.')[0].split('?')[0].split(',')[0].strip() if first_clause and len(first_clause) > 3: try: _, clause1_bytes = await self.voice_service.text_to_speech( first_clause, language_code=sess.language or "en-IN", patient_age=sess.age ) if clause1_bytes: clause1_b64 = base64.b64encode(clause1_bytes).decode("utf-8") await websocket.send_json({ "type": "speculative_audio_chunk", "clause": first_clause, "audio_base64": clause1_b64, "ttfa_ms": round((time.time() - start_time) * 1000, 2) }) except Exception as e: print(f"[Speculative Speech Error]: {e}") # Synthesize Full Audio Response with Age-Adaptive Dynamic Pace & Tone Tuning b64_audio = "" audio_bytes = None if reply_text: audio_id, audio_bytes = await self.voice_service.text_to_speech( reply_text, language_code=sess.language or "en-IN", patient_age=sess.age ) if audio_bytes: b64_audio = base64.b64encode(audio_bytes).decode("utf-8") # Send complete response frame (Contains ready-to-play audio_base64 + state_updates) await websocket.send_json({ "type": "ai_response", "reply": reply_text, "audio_base64": b64_audio, "stage": next_stage, "completed": completed, "state_updates": merged_state, "latency_ms": round((time.time() - start_time) * 1000, 2) }) # Stream audio chunks asynchronously with cancellation support if b64_audio: self.is_ai_speaking[session_id] = True self.tts_tasks[session_id] = asyncio.create_task( self._stream_tts_chunks(websocket, session_id, b64_audio) ) try: await self.tts_tasks[session_id] except asyncio.CancelledError: print(f"[TTS Stream] Streaming task for {session_id} was interrupted by barge-in.") finally: self.is_processing[session_id] = False db.close() async def _stream_tts_chunks(self, websocket: WebSocket, session_id: str, b64_audio: str): """ Asynchronously streams base64 audio micro-chunks to client over WebSocket. Can be canceled instantaneously if patient interrupts (barge-in). """ try: await websocket.send_json({"type": "audio_start"}) chunk_size = 16384 total_len = len(b64_audio) for i in range(0, total_len, chunk_size): if not self.is_ai_speaking.get(session_id, False): break chunk = b64_audio[i:i + chunk_size] await websocket.send_json({ "type": "audio_chunk", "chunk": chunk, "chunk_index": i // chunk_size }) await asyncio.sleep(0.005) if self.is_ai_speaking.get(session_id, False): await websocket.send_json({"type": "audio_end"}) except asyncio.CancelledError: pass finally: self.is_ai_speaking[session_id] = False