473 lines
20 KiB
Python
473 lines
20 KiB
Python
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
|