File size: 6,551 Bytes
6b62834 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | """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":
# Voice input: transcribe → agent → TTS → audio
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
# Send transcribed text back so client can display it
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
# Synthesize TTS audio from the final answer
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:
# Text chat
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 # Fallback: treat as text
|