feat: complete production hardening implementation, specialty canonicalization, google-genai migration, voice reliability, sqlite concurrency, and test suite consolidation
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
WebSocket API Package for Real-Time Patient Onboarding.
|
||||
"""
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
WebSocket Connection Manager (Section 18, 19, 28).
|
||||
Manages active WebSocket instances and per-session processing locks.
|
||||
Ensures session-level concurrency control (messages for same session processed sequentially,
|
||||
different sessions process concurrently without global blocking).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Set, Optional
|
||||
from fastapi import WebSocket, status
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
def __init__(self):
|
||||
# Maps session_id to set of active WebSocket instances
|
||||
self.active_connections: Dict[str, Set[WebSocket]] = {}
|
||||
# Per-session processing lock to enforce sequential message processing per session (Section 18 & 19)
|
||||
self.session_locks: Dict[str, asyncio.Lock] = {}
|
||||
# Stores processed message_ids and assistant responses for deduplication per session (Section 17)
|
||||
self.dedup_store: Dict[str, Dict[str, dict]] = {}
|
||||
|
||||
def get_session_lock(self, session_id: str) -> asyncio.Lock:
|
||||
"""Returns or creates a per-session lock."""
|
||||
if session_id not in self.session_locks:
|
||||
self.session_locks[session_id] = asyncio.Lock()
|
||||
return self.session_locks[session_id]
|
||||
|
||||
async def connect(self, session_id: str, websocket: WebSocket):
|
||||
"""Accept and register connection."""
|
||||
if session_id not in self.active_connections:
|
||||
self.active_connections[session_id] = set()
|
||||
self.active_connections[session_id].add(websocket)
|
||||
logger.info("websocket_connected", extra={"session_id": session_id})
|
||||
|
||||
def disconnect(self, session_id: str, websocket: WebSocket):
|
||||
"""Clean up disconnected socket reference."""
|
||||
if session_id in self.active_connections:
|
||||
self.active_connections[session_id].discard(websocket)
|
||||
if not self.active_connections[session_id]:
|
||||
del self.active_connections[session_id]
|
||||
# Clean up session lock if no connections remain
|
||||
if session_id in self.session_locks:
|
||||
del self.session_locks[session_id]
|
||||
logger.info("websocket_disconnected", extra={"session_id": session_id})
|
||||
|
||||
async def send_json(self, websocket: WebSocket, data: dict):
|
||||
"""Send JSON payload to specific socket."""
|
||||
try:
|
||||
await websocket.send_json(data)
|
||||
except Exception as e:
|
||||
logger.warning("websocket_send_failed", extra={"error": str(e)})
|
||||
|
||||
async def send_to_session(self, session_id: str, data: dict):
|
||||
"""Broadcast payload to all sockets connected to a session."""
|
||||
if session_id in self.active_connections:
|
||||
for ws in list(self.active_connections[session_id]):
|
||||
await self.send_json(ws, data)
|
||||
|
||||
def cache_processed_message(self, session_id: str, message_id: str, response_data: dict):
|
||||
"""Cache response data for duplicate message_id protection (Section 17)."""
|
||||
if session_id not in self.dedup_store:
|
||||
self.dedup_store[session_id] = {}
|
||||
self.dedup_store[session_id][message_id] = response_data
|
||||
|
||||
def get_cached_message_response(self, session_id: str, message_id: str) -> Optional[dict]:
|
||||
"""Return cached response for message_id if previously processed."""
|
||||
return self.dedup_store.get(session_id, {}).get(message_id)
|
||||
|
||||
|
||||
manager = ConnectionManager()
|
||||
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
WebSocket Route Handler (Section 2, 3, 4, 15, 16, 22, 23, 24, 25).
|
||||
Endpoint: WS /api/v1/ws/onboarding/{session_id}
|
||||
Handles connection lifecycle, authentication, session validation, input sanitization, rate limiting, and message dispatch.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Optional
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.database.database import get_db
|
||||
from app.core.security import get_auth_context, verify_session_ownership, AuthContext
|
||||
from app.services.onboarding_service import OnboardingService
|
||||
from app.services.websocket_service import WebSocketService
|
||||
from app.api.websocket.manager import manager
|
||||
from app.api.websocket.schemas import (
|
||||
WSEventType,
|
||||
WSErrorCode,
|
||||
WSConnectionAck,
|
||||
WSPatientMessage,
|
||||
WSErrorEvent,
|
||||
)
|
||||
|
||||
from app.database.database import SessionLocal
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/api/v1/ws/onboarding",
|
||||
tags=["WebSocket Onboarding"],
|
||||
)
|
||||
|
||||
# In-memory per-connection rate limiter (Section 26)
|
||||
# Maps session_id -> list of message timestamps
|
||||
WS_RATE_LIMIT_STORE = {}
|
||||
|
||||
|
||||
def check_ws_rate_limit(session_id: str, limit: int = 30, window_sec: int = 60) -> bool:
|
||||
import time
|
||||
now = time.time()
|
||||
timestamps = WS_RATE_LIMIT_STORE.get(session_id, [])
|
||||
timestamps = [t for t in timestamps if now - t < window_sec]
|
||||
if len(timestamps) >= limit:
|
||||
return False
|
||||
timestamps.append(now)
|
||||
WS_RATE_LIMIT_STORE[session_id] = timestamps
|
||||
return True
|
||||
|
||||
|
||||
@router.websocket("/{session_id}")
|
||||
async def websocket_onboarding_endpoint(
|
||||
websocket: WebSocket,
|
||||
session_id: str,
|
||||
token: Optional[str] = Query(None),
|
||||
x_api_key: Optional[str] = Query(None, alias="api_key"),
|
||||
x_patient_id: Optional[str] = Query(None, alias="patient_id"),
|
||||
):
|
||||
"""
|
||||
Real-Time Patient Onboarding WebSocket Endpoint.
|
||||
WS /api/v1/ws/onboarding/{session_id}
|
||||
"""
|
||||
db = SessionLocal()
|
||||
|
||||
|
||||
try:
|
||||
await _handle_websocket_lifecycle(
|
||||
websocket=websocket,
|
||||
session_id=session_id,
|
||||
db=db,
|
||||
)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
async def _handle_websocket_lifecycle(
|
||||
websocket: WebSocket,
|
||||
session_id: str,
|
||||
db: Session,
|
||||
):
|
||||
onboarding_service = OnboardingService()
|
||||
ws_service = WebSocketService()
|
||||
|
||||
await websocket.accept()
|
||||
|
||||
# Extract query params directly from WebSocket request URL (Section 25)
|
||||
query_token = websocket.query_params.get("token")
|
||||
query_api_key = websocket.query_params.get("api_key")
|
||||
query_patient_id = websocket.query_params.get("patient_id")
|
||||
|
||||
auth_token = query_token or query_api_key or websocket.headers.get("x-api-key")
|
||||
if not auth_token:
|
||||
# Allow test connection token if in header or subprotocol
|
||||
sec_protocol = websocket.headers.get("sec-websocket-protocol")
|
||||
if sec_protocol:
|
||||
auth_token = sec_protocol
|
||||
|
||||
try:
|
||||
auth_context = get_auth_context(
|
||||
api_key=auth_token,
|
||||
credentials=None,
|
||||
x_patient_id=query_patient_id or websocket.headers.get("x-patient-id")
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("websocket_auth_failed", extra={"session_id": session_id, "error": str(e)})
|
||||
await websocket.send_json(
|
||||
WSErrorEvent(code=WSErrorCode.UNAUTHORIZED.value, message="Authentication failed.").model_dump()
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
return
|
||||
|
||||
# 2. Validate session exists & active (Section 24)
|
||||
session = onboarding_service.get_session(db, session_id)
|
||||
if not session:
|
||||
logger.warning("websocket_session_not_found", extra={"session_id": session_id})
|
||||
await websocket.send_json(
|
||||
WSErrorEvent(code=WSErrorCode.SESSION_NOT_FOUND.value, message="Session not found.").model_dump()
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
return
|
||||
|
||||
# 3. Verify session ownership (Section 25)
|
||||
try:
|
||||
verify_session_ownership(session.patient_id, auth_context)
|
||||
except Exception as e:
|
||||
logger.warning("websocket_unauthorized_session_access", extra={"session_id": session_id, "error": str(e)})
|
||||
await websocket.send_json(
|
||||
WSErrorEvent(code=WSErrorCode.UNAUTHORIZED.value, message="Unauthorized session access.").model_dump()
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
return
|
||||
|
||||
if session.status != "in_progress":
|
||||
logger.warning("websocket_session_expired", extra={"session_id": session_id})
|
||||
await websocket.send_json(
|
||||
WSErrorEvent(code=WSErrorCode.SESSION_EXPIRED.value, message="Session expired or completed.").model_dump()
|
||||
)
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
return
|
||||
|
||||
# 4. Register connection & send ACK (Section 3 & 4)
|
||||
await manager.connect(session_id, websocket)
|
||||
|
||||
ack = WSConnectionAck(session_id=session_id, status="connected").model_dump()
|
||||
await manager.send_json(websocket, ack)
|
||||
|
||||
try:
|
||||
while True:
|
||||
# 5. Receive text message
|
||||
data_text = await websocket.receive_text()
|
||||
|
||||
# 6. Check Rate Limiting (Section 26)
|
||||
if not check_ws_rate_limit(session_id, limit=30, window_sec=60):
|
||||
await manager.send_json(
|
||||
websocket,
|
||||
WSErrorEvent(
|
||||
code=WSErrorCode.RATE_LIMITED.value,
|
||||
message="Please wait a moment before sending another message."
|
||||
).model_dump()
|
||||
)
|
||||
continue
|
||||
|
||||
# 7. Validate JSON structure & size (Section 15 & 16)
|
||||
if len(data_text) > 5000:
|
||||
await manager.send_json(
|
||||
websocket,
|
||||
WSErrorEvent(
|
||||
code=WSErrorCode.MESSAGE_TOO_LARGE.value,
|
||||
message="Message exceeds maximum allowed length of 5000 characters."
|
||||
).model_dump()
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
payload = json.loads(data_text)
|
||||
except Exception:
|
||||
await manager.send_json(
|
||||
websocket,
|
||||
WSErrorEvent(
|
||||
code=WSErrorCode.INVALID_MESSAGE.value,
|
||||
message="Please send a valid JSON message."
|
||||
).model_dump()
|
||||
)
|
||||
continue
|
||||
|
||||
if not isinstance(payload, dict) or payload.get("type") != WSEventType.PATIENT_MESSAGE.value:
|
||||
await manager.send_json(
|
||||
websocket,
|
||||
WSErrorEvent(
|
||||
code=WSErrorCode.INVALID_MESSAGE.value,
|
||||
message="Unsupported message type. Expected 'patient.message'."
|
||||
).model_dump()
|
||||
)
|
||||
continue
|
||||
|
||||
message_id = str(payload.get("message_id") or "").strip()
|
||||
text = str(payload.get("text") or "").strip()
|
||||
|
||||
if not message_id or not text:
|
||||
await manager.send_json(
|
||||
websocket,
|
||||
WSErrorEvent(
|
||||
code=WSErrorCode.INVALID_MESSAGE.value,
|
||||
message="Message must include non-empty 'message_id' and 'text'."
|
||||
).model_dump()
|
||||
)
|
||||
continue
|
||||
|
||||
# 8. Session-level lock for sequential processing per session (Section 18 & 19)
|
||||
async with manager.get_session_lock(session_id):
|
||||
await ws_service.handle_patient_message(
|
||||
websocket=websocket,
|
||||
db=db,
|
||||
session_id=session_id,
|
||||
message_id=message_id,
|
||||
text=text,
|
||||
)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
manager.disconnect(session_id, websocket)
|
||||
logger.info("websocket_client_disconnected", extra={"session_id": session_id})
|
||||
except Exception as e:
|
||||
manager.disconnect(session_id, websocket)
|
||||
logger.error("websocket_unexpected_error", extra={"session_id": session_id, "error": str(e)})
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
WebSocket Protocol Schemas and Event Constants (Section 5, 6, 15).
|
||||
Centralized definitions for event types, error codes, and message models.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Optional, Any, Dict
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class WSEventType(str, Enum):
|
||||
CONNECTION_ACK = "connection.ack"
|
||||
PATIENT_MESSAGE = "patient.message"
|
||||
ASSISTANT_MESSAGE = "assistant.message"
|
||||
PROCESSING_STARTED = "processing.started"
|
||||
PROCESSING_COMPLETED = "processing.completed"
|
||||
TYPING_STARTED = "typing.started"
|
||||
TYPING_STOPPED = "typing.stopped"
|
||||
STATE_UPDATED = "state.updated"
|
||||
SAFETY_ALERT = "safety.alert"
|
||||
ONBOARDING_COMPLETED = "onboarding.completed"
|
||||
ERROR = "error"
|
||||
SESSION_ENDED = "session.ended"
|
||||
|
||||
|
||||
class WSErrorCode(str, Enum):
|
||||
INVALID_MESSAGE = "INVALID_MESSAGE"
|
||||
MESSAGE_TOO_LARGE = "MESSAGE_TOO_LARGE"
|
||||
UNAUTHORIZED = "UNAUTHORIZED"
|
||||
SESSION_NOT_FOUND = "SESSION_NOT_FOUND"
|
||||
SESSION_EXPIRED = "SESSION_EXPIRED"
|
||||
PROCESSING_ERROR = "PROCESSING_ERROR"
|
||||
AI_TIMEOUT = "AI_TIMEOUT"
|
||||
AI_UNAVAILABLE = "AI_UNAVAILABLE"
|
||||
RATE_LIMITED = "RATE_LIMITED"
|
||||
INVALID_STATE = "INVALID_STATE"
|
||||
INTERNAL_ERROR = "INTERNAL_ERROR"
|
||||
|
||||
|
||||
class WSConnectionAck(BaseModel):
|
||||
type: str = WSEventType.CONNECTION_ACK.value
|
||||
session_id: str
|
||||
status: str = "connected"
|
||||
|
||||
|
||||
class WSPatientMessage(BaseModel):
|
||||
type: str = WSEventType.PATIENT_MESSAGE.value
|
||||
message_id: str = Field(..., description="Unique client message ID")
|
||||
text: str = Field(..., max_length=5000, description="Patient text message")
|
||||
|
||||
|
||||
class WSAssistantMessage(BaseModel):
|
||||
type: str = WSEventType.ASSISTANT_MESSAGE.value
|
||||
message_id: str
|
||||
reply: str
|
||||
language: str = "en-IN"
|
||||
|
||||
|
||||
class WSTypingEvent(BaseModel):
|
||||
type: str # typing.started | typing.stopped
|
||||
|
||||
|
||||
class WSProcessingStartedEvent(BaseModel):
|
||||
type: str = WSEventType.PROCESSING_STARTED.value
|
||||
|
||||
|
||||
class WSProcessingCompletedEvent(BaseModel):
|
||||
type: str = WSEventType.PROCESSING_COMPLETED.value
|
||||
processing_time_ms: int
|
||||
|
||||
|
||||
class WSSafetyAlertEvent(BaseModel):
|
||||
type: str = WSEventType.SAFETY_ALERT.value
|
||||
level: str = "emergency" # emergency | crisis
|
||||
message: Optional[str] = None
|
||||
|
||||
|
||||
class WSStateUpdatedEvent(BaseModel):
|
||||
type: str = WSEventType.STATE_UPDATED.value
|
||||
session_id: str
|
||||
|
||||
|
||||
class WSOnboardingCompletedEvent(BaseModel):
|
||||
type: str = WSEventType.ONBOARDING_COMPLETED.value
|
||||
session_id: str
|
||||
|
||||
|
||||
class WSErrorEvent(BaseModel):
|
||||
type: str = WSEventType.ERROR.value
|
||||
code: str
|
||||
message: str
|
||||
Reference in New Issue
Block a user