aionboarding/ai-patient-onboarding/app/services/websocket_service.py

"""
WebSocket Transport Service (Section 7, 8, 10, 11, 12, 13, 14, 17).
Thin bridge between WebSocket events and core onboarding business services.
Delegates all clinical logic, state management, and AI orchestration to OnboardingService/AIOnboardingService.
"""

import time
import logging
from typing import Dict, Any
from fastapi import WebSocket
from sqlalchemy.orm import Session

from app.api.websocket.manager import manager
from app.api.websocket.schemas import (
    WSEventType,
    WSErrorCode,
    WSAssistantMessage,
    WSTypingEvent,
    WSProcessingStartedEvent,
    WSProcessingCompletedEvent,
    WSSafetyAlertEvent,
    WSStateUpdatedEvent,
    WSOnboardingCompletedEvent,
    WSErrorEvent,
)
from app.services.onboarding_service import OnboardingService
from app.services.ai_onboarding_service import AIOnboardingService
from app.database.models import OnboardingMessage

logger = logging.getLogger(__name__)


class WebSocketService:
    def __init__(self):
        self.onboarding_service = OnboardingService()
        self.ai_service = AIOnboardingService()

    async def handle_patient_message(
        self,
        websocket: WebSocket,
        db: Session,
        session_id: str,
        message_id: str,
        text: str,
    ):
        """
        Process patient message through transport pipeline (Section 7).
        Enforces typing events, safety alerts, state updates, and deduplication.
        """
        t0 = time.perf_counter()

        # 1. Deduplication check (Section 17)
        cached_resp = manager.get_cached_message_response(session_id, message_id)
        if cached_resp:
            logger.info("websocket_duplicate_message_received", extra={"session_id": session_id, "message_id": message_id})
            await manager.send_json(websocket, cached_resp)
            return

        # 2. Emit processing.started
        await manager.send_json(websocket, WSProcessingStartedEvent().model_dump())

        # 3. Emit typing.started
        await manager.send_json(websocket, WSTypingEvent(type=WSEventType.TYPING_STARTED.value).model_dump())

        try:
            # Fetch session & verify active status
            session = self.onboarding_service.get_session(db, session_id)
            if not session or session.status != "in_progress":
                await manager.send_json(
                    websocket,
                    WSErrorEvent(
                        code=WSErrorCode.SESSION_EXPIRED.value,
                        message="Onboarding session is expired or not active."
                    ).model_dump()
                )
                return

            # Save patient message
            self.onboarding_service.save_message(
                db=db,
                session_id=session_id,
                role="patient",
                message=text,
                stage=session.stage,
            )

            state = self.onboarding_service.get_state(db, session_id)

            messages = (
                db.query(OnboardingMessage)
                .filter(OnboardingMessage.session_id == session_id)
                .order_by(OnboardingMessage.created_at.asc())
                .all()
            )
            conversation_history = [
                {"role": m.role, "message": m.message, "stage": m.stage}
                for m in messages
            ]

            # 4. Delegate to AI Orchestrator
            ai_result = await self.ai_service.process_message(
                session=session,
                state=state,
                conversation_history=conversation_history,
                user_message=text,
            )

            # 5. Check Emergency / Crisis priority (Section 12)
            is_emergency = ai_result.get("is_emergency") or ai_result.get("next_stage") == "emergency_referral"
            is_crisis = ai_result.get("is_crisis") or ai_result.get("next_stage") == "crisis_handling"
            if is_emergency or is_crisis:
                level = "crisis" if is_crisis else "emergency"
                await manager.send_json(
                    websocket,
                    WSSafetyAlertEvent(
                        level=level,
                        message="Safety event detected. Normal intake flow suspended."
                    ).model_dump()
                )

            reply_text = ai_result.get("reply", "Got it. When did it start?")
            next_stage = ai_result.get("next_stage", session.stage)
            session.stage = next_stage

            # Save assistant message
            self.onboarding_service.save_message(
                db=db,
                session_id=session_id,
                role="assistant",
                message=reply_text,
                stage=next_stage,
            )

            # Update state in DB if modified
            updated_state_dict = ai_result.get("updated_state") or (ai_result.get("state_updates") if isinstance(ai_result.get("state_updates"), dict) else None)
            if updated_state_dict and isinstance(updated_state_dict, dict):
                self.onboarding_service.update_state(db, state, updated_state_dict)
            else:
                db.commit()

            # 6. Emit state.updated (Section 13)
            await manager.send_json(
                websocket,
                WSStateUpdatedEvent(session_id=session_id).model_dump()
            )

            # 7. Construct & emit assistant.message (Section 5)
            assistant_resp = WSAssistantMessage(
                message_id=f"resp_{message_id}",
                reply=reply_text,
                language=session.language or "en-IN"
            ).model_dump()

            await manager.send_json(websocket, assistant_resp)

            # Cache response for duplicate message protection
            manager.cache_processed_message(session_id, message_id, assistant_resp)

            # 8. Check onboarding completion (Section 14)
            if ai_result.get("completed") or session.status == "completed":
                session.status = "completed"
                db.commit()
                await manager.send_json(
                    websocket,
                    WSOnboardingCompletedEvent(session_id=session_id).model_dump()
                )

        except Exception as e:
            logger.error("websocket_message_processing_failed", extra={"session_id": session_id, "error": str(e)})
            await manager.send_json(
                websocket,
                WSErrorEvent(
                    code=WSErrorCode.PROCESSING_ERROR.value,
                    message="An error occurred while processing your message."
                ).model_dump()
            )
        finally:
            # Always ensure typing state is stopped (Section 10)
            await manager.send_json(
                websocket,
                WSTypingEvent(type=WSEventType.TYPING_STOPPED.value).model_dump()
            )
            # Emit processing.completed (Section 11)
            latency_ms = int((time.perf_counter() - t0) * 1000)
            await manager.send_json(
                websocket,
                WSProcessingCompletedEvent(processing_time_ms=latency_ms).model_dump()
            )