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)