| 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 |
|
|
| |
| 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", ""), |
| ) |
|
|
| |
| 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() |
|
|
| |
| 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, |
| ) |
|
|
| |
| self.cleanup_manager.reset_context(self.context) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| self.logging_middleware = LoggingMiddleware(self.session_id) |
| self.metrics_middleware = MetricsMiddleware(self.session_id) |
|
|
| |
| self.event_bus.subscribe("*", self.logging_middleware) |
| self.event_bus.subscribe("*", self.metrics_middleware) |
|
|
| |
| 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() |
| |
| 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 = [] |
|
|
| |
| 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}") |
|
|
| |
| new_lang = data.get("language") |
| if new_lang: |
| self.session_config.language = new_lang.strip().lower() |
| changes.append(f"language={new_lang}") |
|
|
| |
| if "voice_interruption" in data: |
| self.session_config.allow_interruptions = bool(data["voice_interruption"]) |
| changes.append(f"allow_interruptions={self.session_config.allow_interruptions}") |
|
|
| |
| 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: |
| |
| 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) |
|
|
