onboarding.py

import asyncio
import time
import uuid
from typing import Any, Awaitable, Callable, Dict, List, Optional, cast

import logging

from utils.executors import llm_executor, run_blocking
from utils.llm.clients import get_llm
from utils.llm.usage_tracker import track_usage, Features
from utils.llm.conversation_processing import _word_count  # type: ignore[reportPrivateUsage]  # shared word count helper

logger = logging.getLogger(__name__)

# The app shows these three together as "Talk About" topics rather than one at
# a time, so the user speaks freely; the handler still walks them in order
# against the accumulated transcript (see _check_answer).
ONBOARDING_QUESTIONS: List[Dict[str, str]] = [
    {'question': "Where do you live?", 'category': 'location'},
    {'question': "What do you do for work?", 'category': 'work'},
    {'question': "What is your long-term goal?", 'category': 'long_term_goal'},
]


class OnboardingHandler:
    """Handles onboarding question flow via websocket"""

    # Special speaker ID for Omi question segments (use 99 to avoid conflicts with real speakers)
    OMI_SPEAKER_ID = 99

    def __init__(
        self,
        uid: str,
        send_message: Callable[[Dict[str, Any]], Awaitable[None]],
        stream_transcript: Optional[Callable[[List[Dict[str, Any]]], None]] = None,
        session_id: Optional[str] = None,
    ) -> None:
        self.uid = uid
        # Server-generated provenance.  Clients may request onboarding mode,
        # but they cannot choose or forge this session identity; consumers use
        # it instead of trusting request.source.
        # The admission/session identity is issued by the authenticated
        # backend.  Keep the UUID fallback for existing internal callers, but
        # never accept a client-provided value here.
        self.session_id = session_id if isinstance(session_id, str) and len(session_id) >= 16 else uuid.uuid4().hex
        self.send_message = send_message
        self.stream_transcript = stream_transcript  # Callback to inject segments into transcript stream
        self.questions: List[Dict[str, str]] = ONBOARDING_QUESTIONS.copy()
        self.current_question_index = 0
        self.answers: List[Dict[str, Any]] = []
        self.current_transcript = ''
        self.silence_timer: Optional[asyncio.Task[None]] = None
        self.is_checking_answer = False
        # Segments that arrive while an AI answer check is awaiting the LLM are
        # queued here and replayed once the check finishes, so speech covering
        # later topics is evaluated instead of dropped.
        self.pending_segments: List[Dict[str, Any]] = []
        self.completed = False
        self.started = False
        self.start_time: Optional[float] = None  # Track when onboarding started
        self.last_segment_end: float = 0.0  # Track end time for question segment timing

    @property
    def current_question(self) -> Optional[Dict[str, str]]:
        if self.current_question_index < len(self.questions):
            return self.questions[self.current_question_index]
        return None

    def _get_elapsed_time(self) -> float:
        """Get elapsed time since onboarding started."""
        if self.start_time is None:
            self.start_time = time.time()
        return time.time() - self.start_time

    def _create_question_segment(self) -> Optional[Dict[str, Any]]:
        """Create a transcript segment for the current question."""
        if not self.current_question:
            return None

        # Question segment starts after last segment ended, with small gap
        start_time = self.last_segment_end + 0.5 if self.last_segment_end > 0 else 0.0
        # Estimate end time based on question length (rough: 150 words per minute)
        words = len(self.current_question['question'].split())
        duration = max(1.0, words / 2.5)  # At least 1 second
        end_time = start_time + duration

        return {
            'id': str(uuid.uuid4()),
            'text': self.current_question['question'],
            'start': start_time,
            'end': end_time,
            'speaker': f'SPEAKER_{self.OMI_SPEAKER_ID}',  # Use consistent format with STT output
            'speaker_id': self.OMI_SPEAKER_ID,
            'is_user': False,
            'person_id': None,
        }

    def update_segment_timing(self, segments: List[Dict[str, Any]]) -> None:
        """Update timing tracking based on received segments."""
        for segment in segments:
            end_time = segment.get('end', 0)
            if end_time > self.last_segment_end:
                self.last_segment_end = end_time

    def on_segments_received(self, segments: List[Dict[str, Any]]) -> None:
        """Called when new transcript segments are received"""
        if self.completed:
            return
        if self.is_checking_answer:
            # An AI answer check can await up to three LLM calls; speech that
            # arrives during that window must not be lost. Queue it and replay
            # it when _check_answer finishes.
            self.pending_segments.extend(segments)
            return

        # Update timing tracking
        self.update_segment_timing(segments)

        # Accumulate transcript for current question (ignore Omi segments)
        new_text = ' '.join(s.get('text', '') for s in segments if s.get('speaker_id') != self.OMI_SPEAKER_ID).strip()
        if new_text:
            if self.current_transcript:
                self.current_transcript += ' ' + new_text
            else:
                self.current_transcript = new_text

        # Reset silence timer
        if self.silence_timer:
            self.silence_timer.cancel()

        # Start new silence timer (2 seconds)
        self.silence_timer = asyncio.create_task(self._silence_check())

    async def _silence_check(self) -> None:
        """Check answer after 2 seconds of silence"""
        await asyncio.sleep(2.0)

        if self.completed or self.is_checking_answer:
            return

        if not self.current_transcript.strip():
            return

        await self._check_answer()

    async def skip_current_question(self) -> None:
        """Skip the current question and move to the next one"""
        if self.completed or self.is_checking_answer:
            return

        # Cancel any pending silence timer
        if self.silence_timer:
            self.silence_timer.cancel()
            self.silence_timer = None

        # Record that this question was skipped
        if self.current_question:
            self.answers.append(
                {
                    'question': self.current_question['question'],
                    'answer': self.current_transcript.strip() if self.current_transcript.strip() else '[skipped]',
                    'category': self.current_question['category'],
                    'skipped': True,
                }
            )

        # Send event to app
        await self._send_event(
            'question_skipped',
            {
                'question_index': self.current_question_index,
            },
        )

        # Move to next question
        self.current_question_index += 1
        self.current_transcript = ''

        if self.current_question_index >= len(self.questions):
            await self._complete_onboarding()
        else:
            await self.send_current_question()

    async def start(self) -> None:
        """Start the question flow after the client has attached its listener.

        The client explicitly requests this transition so the first question cannot
        race the WebSocket subscription that consumes it. Repeated requests are safe
        and do not inject duplicate question segments.
        """
        if self.completed or self.started:
            return

        self.start_time = time.time()
        await self.send_current_question()
        self.started = True

    async def _check_answer(self) -> None:
        """Use AI to check if question was answered"""
        if self.is_checking_answer or not self.current_question:
            return

        self.is_checking_answer = True

        try:
            transcript = self.current_transcript.strip()

            # The topics are shown to the user all at once, so one stretch of
            # speech may cover several of them. Keep the transcript across
            # questions and advance through every question it answers.
            while self.current_question and not self.completed:
                question = self.current_question['question']
                answered = False
                if _word_count(transcript) >= 2:
                    answered = await self._ai_check_answer(question, transcript)
                if not answered:
                    break

                self.answers.append(
                    {
                        'question': question,
                        'answer': transcript,
                        'category': self.current_question['category'],
                    }
                )
                await self._send_event(
                    'question_answered',
                    {
                        'question_index': self.current_question_index,
                        'answered': True,
                    },
                )
                self.current_question_index += 1
                if self.current_question_index >= len(self.questions):
                    await self._complete_onboarding()
                else:
                    await self.send_current_question()

        finally:
            self.is_checking_answer = False
            # Replay what was said while the checks were awaiting the LLM so it
            # accumulates into the transcript and restarts the silence timer
            # for the next evaluation.
            pending, self.pending_segments = self.pending_segments, []
            if pending and not self.completed:
                self.on_segments_received(pending)

    async def _ai_check_answer(self, question: str, transcript: str) -> bool:
        """Use AI to determine if answer is valid"""
        try:
            prompt = f"""Determine if this transcript answers the question. Be lenient - any attempt to answer counts.
Question: "{question}"
Transcript: "{transcript}"

Reply with only "yes" or "no"."""

            with track_usage(self.uid, Features.ONBOARDING):
                response = await run_blocking(llm_executor, get_llm('onboarding').invoke, prompt)
            return 'yes' in cast(str, cast(Any, response).content).lower()
        except Exception as e:
            logger.error(f"AI check error: {e}")
            # Fallback: 2+ words is an answer
            return len(transcript.split()) >= 2

    async def _send_event(self, event_type: str, data: Dict[str, Any]) -> None:
        """Send message event to client"""
        event = {'type': event_type, **data}
        await self.send_message(event)

    async def send_current_question(self) -> None:
        """Send current question to client and inject as transcript segment"""
        if self.current_question:
            # Create and inject question segment into transcript stream
            question_segment = self._create_question_segment()
            if question_segment and self.stream_transcript:
                self.stream_transcript([question_segment])
                # Update last_segment_end to account for the question segment
                self.last_segment_end = question_segment['end']

            await self._send_event(
                'onboarding_question',
                {
                    'question': self.current_question['question'],
                    'question_index': self.current_question_index,
                    'total_questions': len(self.questions),
                    'question_segment_id': question_segment['id'] if question_segment else None,
                },
            )

    async def _complete_onboarding(self) -> None:
        """Signal completion when all questions answered.

        The conversation is already being created/updated by the normal
        transcription flow in transcribe.py. We just signal completion
        so the client can finalize the speech profile and trigger
        conversation processing.
        """
        self.completed = True

        await self._send_event(
            'onboarding_complete',
            {
                'answers_count': len(self.answers),
            },
        )

    def cleanup(self) -> None:
        """Cleanup resources"""
        if self.silence_timer:
            self.silence_timer.cancel()