feat: complete real-time voice-ai engine upgrades and frontend integration guide
This commit is contained in:
@@ -0,0 +1,472 @@
|
||||
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
|
||||
Reference in New Issue
Block a user