| """WebSocket handler for bidirectional streaming communication. |
| |
| Supports: |
| - Real-time chat with streaming responses |
| - Voice input streaming (audio chunks → STT → agent → TTS → audio chunks) |
| - Progress events during agent execution |
| """ |
|
|
| import asyncio |
| import json |
| import uuid |
| from typing import Optional |
|
|
| from fastapi import APIRouter, WebSocket, WebSocketDisconnect |
|
|
| router = APIRouter() |
|
|
|
|
| class ConnectionManager: |
| """Manages active WebSocket connections.""" |
|
|
| def __init__(self): |
| self._connections: dict[str, WebSocket] = {} |
|
|
| async def connect(self, websocket: WebSocket, session_id: str) -> None: |
| await websocket.accept() |
| self._connections[session_id] = websocket |
|
|
| def disconnect(self, session_id: str) -> None: |
| self._connections.pop(session_id, None) |
|
|
| def get(self, session_id: str) -> Optional[WebSocket]: |
| return self._connections.get(session_id) |
|
|
| @property |
| def active_count(self) -> int: |
| return len(self._connections) |
|
|
|
|
| manager = ConnectionManager() |
|
|
|
|
| @router.websocket("/ws/{session_id}") |
| async def websocket_endpoint(websocket: WebSocket, session_id: str): |
| """Main WebSocket endpoint for real-time agent interaction. |
| |
| Message format (client → server): |
| { |
| "type": "chat" | "voice" | "control", |
| "payload": { |
| "message": "...", // For chat type |
| "audio": "<base64>", // For voice type |
| "mode": "auto", // Agent mode |
| "action": "cancel" // For control type |
| } |
| } |
| |
| Message format (server → client): |
| { |
| "type": "text_delta" | "thought" | "tool_call" | "audio" | "error" | "done", |
| "data": {...} |
| } |
| """ |
| await manager.connect(websocket, session_id) |
|
|
| try: |
| from agentic_rag.runtime.unified_context import UnifiedContext |
| from agentic_rag.orchestration.l2_capabilities.chat import ChatCapability |
|
|
| context = UnifiedContext.create(session_id=session_id) |
| chat = ChatCapability() |
|
|
| while True: |
| raw = await websocket.receive_text() |
| msg = json.loads(raw) |
| msg_type = msg.get("type", "chat") |
| payload = msg.get("payload", {}) |
|
|
| if msg_type == "control": |
| action = payload.get("action", "") |
| if action == "cancel": |
| await websocket.send_json({"type": "done", "data": {"cancelled": True}}) |
| continue |
|
|
| elif msg_type == "voice": |
| |
| audio_b64 = payload.get("audio", "") |
| text = await _process_voice_input(audio_b64) |
| if not text.strip(): |
| await websocket.send_json({"type": "error", "data": {"error": "No speech detected"}}) |
| continue |
|
|
| |
| await websocket.send_json({"type": "transcript", "data": {"text": text}}) |
|
|
| full_answer = "" |
| async for event in chat.chat_stream( |
| message=text, |
| session_id=session_id, |
| context=context, |
| ): |
| await _send_event(websocket, event) |
| if event.event_type.value == "text_delta": |
| full_answer += event.data.get("content", "") or event.data.get("delta", "") |
| elif event.event_type.value == "done": |
| ans = event.data.get("final_answer", "") |
| if ans: |
| full_answer = ans |
|
|
| |
| if full_answer.strip(): |
| tts = _create_tts() |
| audio_bytes = await tts.synthesize(full_answer) |
| if audio_bytes: |
| import base64 |
| audio_b64 = base64.b64encode(audio_bytes).decode() |
| await websocket.send_json({ |
| "type": "audio", |
| "data": {"audio": audio_b64, "format": tts.response_format}, |
| }) |
|
|
| await websocket.send_json({"type": "done", "data": {}}) |
|
|
| else: |
| |
| message = payload.get("message", "") |
| mode = payload.get("mode") |
|
|
| async for event in chat.chat_stream( |
| message=message, |
| session_id=session_id, |
| mode=mode, |
| context=context, |
| ): |
| await _send_event(websocket, event) |
|
|
| except WebSocketDisconnect: |
| pass |
| finally: |
| manager.disconnect(session_id) |
|
|
|
|
| async def _send_event(websocket: WebSocket, event) -> None: |
| """Send an AgentEvent to the WebSocket client.""" |
| try: |
| await websocket.send_json({ |
| "type": event.event_type.value, |
| "data": event.model_dump_json() if hasattr(event, 'model_dump_json') else event.data, |
| }) |
| except Exception: |
| pass |
|
|
|
|
| def _create_stt(): |
| """Create STT service from settings.""" |
| from agentic_rag.config.settings import get_settings |
| from agentic_rag.core.voice.stt import STTService |
| s = get_settings().voice |
| return STTService( |
| provider=s.stt_provider, |
| model=s.stt_model, |
| api_base=s.stt_api_base, |
| api_key=s.stt_api_key, |
| language=s.stt_language, |
| sample_rate=s.sample_rate, |
| ) |
|
|
|
|
| def _create_tts(): |
| """Create TTS service from settings.""" |
| from agentic_rag.config.settings import get_settings |
| from agentic_rag.core.voice.tts import TTSService |
| s = get_settings().voice |
| return TTSService( |
| provider=s.tts_provider, |
| model=s.tts_model, |
| api_base=s.tts_api_base, |
| api_key=s.tts_api_key, |
| task_type=s.tts_task_type, |
| instructions=s.tts_instructions, |
| language=s.tts_language, |
| speaker=s.tts_speaker, |
| voice=s.tts_voice, |
| speed=s.tts_speed, |
| response_format=s.tts_response_format, |
| ) |
|
|
|
|
| async def _process_voice_input(audio_b64: str) -> str: |
| """Process voice input: base64 → STT → text.""" |
| import base64 |
| try: |
| audio_bytes = base64.b64decode(audio_b64) |
| stt = _create_stt() |
| return await stt.transcribe_bytes(audio_bytes) |
| except Exception: |
| return audio_b64 |
|
|
