runtime/session.py

import uuid
import logging
import asyncio
from typing import Any
from fastapi import WebSocket, WebSocketDisconnect
from models.event import Event, EventType
from models.session_state import SessionConfig
from models.conversation_profile import ConversationProfile
from runtime.event_bus import EventBus
from runtime.task_supervisor import TaskSupervisor
from runtime.cleanup_manager import CleanupManager
from runtime.sentence_accumulator import SentenceAccumulator
from runtime.error_handler import ErrorHandler
from runtime.turn_manager import TurnManager
from runtime.runtime_context import RuntimeContext
from runtime.audio_buffer_manager import AudioBufferManager
from runtime.transcript_stabilizer import TranscriptStabilizer
from queues.queue_registry import QueueRegistry
from queues import queue_names
from services.vad_service import VADService
from services.asr_service import ASRService
from services.llm_service import LLMService
from services.tts_service import TTSService
from services.tool_service import ToolService
from services.memory_service import MemoryService
from middleware.logging_middleware import LoggingMiddleware
from middleware.metrics_middleware import MetricsMiddleware

logger = logging.getLogger("s2s_server")

class ClientSession:
    def __init__(
        self,
        websocket: WebSocket,
        initial_config: dict,
        tts_model: Any = None,
        nv_client: Any = None,
        groq_client: Any = None,
        get_voice_state: Any = None,
        run_transcription: Any = None,
        clean_and_check_text: Any = None,
        available_voices: Any = None,
        persona_system_prompts: Any = None,
    ):
        self.session_id = str(uuid.uuid4())[:8]
        self.websocket = websocket

        # Config & Profiles
        self.session_config = SessionConfig(
            voice_name=initial_config.get("voice_name", "cosette"),
            language=initial_config.get("language", "en"),
            allow_interruptions=initial_config.get("voice_interruption", False),
        )
        self.conversation_profile = ConversationProfile(
            persona=initial_config.get("persona", "assistant"),
            system_prompt=initial_config.get("system_prompt", ""),
        )

        # Supervisors and Audio Buffer
        self.event_bus = EventBus(self.session_id)
        self.task_supervisor = TaskSupervisor(self.session_id)
        self.cleanup_manager = CleanupManager()
        self.queue_registry = QueueRegistry(self.session_id)
        self.audio_buffer = AudioBufferManager()

        # Context Setup
        self.context = RuntimeContext(
            session_id=self.session_id,
            websocket=self.websocket,
            event_bus=self.event_bus,
            queue_registry=self.queue_registry,
            task_supervisor=self.task_supervisor,
            cleanup_manager=self.cleanup_manager,
            session_config=self.session_config,
            conversation_profile=self.conversation_profile,
            logger=logger,
            tts_model=tts_model,
            nv_client=nv_client,
            groq_client=groq_client,
            get_voice_state=get_voice_state,
            run_transcription=run_transcription,
            clean_and_check_text=clean_and_check_text,
            available_voices=available_voices,
            persona_system_prompts=persona_system_prompts,
            audio_buffer=self.audio_buffer,
        )

        # Link cleanup manager context
        self.cleanup_manager.reset_context(self.context)

        # Subsystems
        self.error_handler = ErrorHandler(self.context)
        self.context.error_handler = self.error_handler

        self.sentence_accumulator = SentenceAccumulator(self.context)
        self.context.sentence_accumulator = self.sentence_accumulator

        # Services
        self.vad = VADService(self.context)
        self.asr = ASRService(self.context)
        self.llm = LLMService(self.context)
        self.tts = TTSService(self.context)
        self.tools = ToolService(self.context)
        self.memory = MemoryService(self.context)

        # Link services to context
        self.context.vad_service = self.vad
        self.context.asr_service = self.asr
        self.context.llm_service = self.llm
        self.context.tts_service = self.tts
        self.context.memory_service = self.memory

        # Turn Orchestrator and Transcript Stabilizer
        self.turn_manager = TurnManager(self.context)
        self.context.turn_manager = self.turn_manager
        
        self.transcript_stabilizer = TranscriptStabilizer(self.context)
        self.context.transcript_stabilizer = self.transcript_stabilizer

        # Middlewares
        self.logging_middleware = LoggingMiddleware(self.session_id)
        self.metrics_middleware = MetricsMiddleware(self.session_id)

        # Register wildcards
        self.event_bus.subscribe("*", self.logging_middleware)
        self.event_bus.subscribe("*", self.metrics_middleware)

        # Internal event bindings
        self.event_bus.subscribe(EventType.LLM_TOKEN, self.on_llm_token)
        self.event_bus.subscribe(EventType.LLM_FINISHED, self.on_llm_finished)

    async def on_llm_token(self, event: Event):
        self.sentence_accumulator.handle_token(event.payload)

    async def on_llm_finished(self, event: Event):
        self.sentence_accumulator.flush()

    async def start(self):
        self.context.logger.info(f"[{self.session_id}] ClientSession: Starting all services...")
        await self.vad.start()
        await self.asr.start()
        await self.llm.start()
        await self.tts.start()
        await self.tools.start()
        await self.memory.start()
        self.context.logger.info(f"[{self.session_id}] ClientSession: Services running.")

    async def stop(self):
        self.context.logger.info(f"[{self.session_id}] ClientSession: Stopping all services...")
        await self.vad.stop()
        await self.asr.stop()
        await self.llm.stop()
        await self.tts.stop()
        await self.tools.stop()
        await self.memory.stop()
        self.task_supervisor.cancel_all()

    async def cleanup(self):
        self.context.logger.info(f"[{self.session_id}] ClientSession: Initiating session cleanup.")
        await self.event_bus.publish(EventType.CLIENT_DISCONNECTED)
        await self.stop()
        self.cleanup_manager.perform_cleanup()

    async def run(self):
        await self.start()
        # Receive websocket message stream
        try:
            while True:
                message = await self.websocket.receive()
                await self.handle_message(message)
        except WebSocketDisconnect:
            self.context.logger.info(f"[{self.session_id}] ClientSession: WebSocket connection closed.")
        except Exception as e:
            if isinstance(e, RuntimeError) and "once a disconnect message has been received" in str(e):
                self.context.logger.info(f"[{self.session_id}] ClientSession: WebSocket connection closed (RuntimeError).")
            else:
                self.context.logger.error(f"[{self.session_id}] ClientSession error in receive loop: {e}", exc_info=True)
                raise

    async def handle_message(self, message: dict):
        if "text" in message:
            import json
            try:
                data = json.loads(message["text"])
                msg_type = data.get("type")

                if msg_type == "interrupt":
                    await self.event_bus.publish(EventType.BARGE_IN)

                elif msg_type == "config":
                    changes = []

                    # 1. Update Voice Setting
                    new_voice = data.get("voice")
                    if new_voice:
                        if new_voice in self.context.available_voices:
                            self.session_config.voice_name = new_voice
                            changes.append(f"voice={new_voice}")

                    # 2. Update Language Setting
                    new_lang = data.get("language")
                    if new_lang:
                        self.session_config.language = new_lang.strip().lower()
                        changes.append(f"language={new_lang}")

                    # 3. Update Voice Interruption Toggle
                    if "voice_interruption" in data:
                        self.session_config.allow_interruptions = bool(data["voice_interruption"])
                        changes.append(f"allow_interruptions={self.session_config.allow_interruptions}")

                    # 4. Update Conversation Profile Persona
                    new_persona = data.get("persona")
                    if new_persona:
                        if new_persona in self.context.persona_system_prompts:
                            self.conversation_profile.persona = new_persona
                            self.conversation_profile.system_prompt = self.context.persona_system_prompts[new_persona]
                            changes.append(f"persona={new_persona}")
                        elif new_persona == "custom" and "system_prompt" in data:
                            self.conversation_profile.persona = "custom"
                            self.conversation_profile.system_prompt = data["system_prompt"]
                            changes.append("persona=custom")

                    if changes:
                        self.context.logger.info(
                            f"[{self.session_id}] Config modified: {', '.join(changes)}"
                        )
                        await self.websocket.send_json({
                            "type": "status",
                            "message": f"Config updated: {', '.join(changes)}"
                        })
            except Exception as e:
                self.context.logger.error(
                    f"[{self.session_id}] ClientSession failed to process config frame: {e}"
                )

        elif "bytes" in message:
            # Route raw audio chunks to input registry queue
            import numpy as np
            raw_bytes = message["bytes"]
            chunk = np.frombuffer(raw_bytes, dtype=np.float32)
            if len(chunk) > 0:
                self.queue_registry.get(queue_names.AUDIO).put_nowait(chunk)