549 lines
22 KiB
Python
549 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable
|
|
|
|
from services.ai_voice_runtime_service.audiosocket import (
|
|
AUDIO_SOCKET_PACKET_DTMF,
|
|
AUDIO_SOCKET_PACKET_HANGUP,
|
|
AUDIO_SOCKET_PACKET_PCM16,
|
|
AUDIO_SOCKET_PACKET_UUID,
|
|
EnergyVAD,
|
|
chunk_audio,
|
|
encode_audio_packet,
|
|
normalize_media_uuid,
|
|
pcm16le_to_wav_bytes,
|
|
read_packet,
|
|
resample_pcm16le,
|
|
)
|
|
from services.ai_voice_runtime_service.providers.asr import ASRProvider
|
|
from services.ai_voice_runtime_service.providers.tts import TTSProvider
|
|
from services.shared.models import VoiceAITurnDecisionOut
|
|
|
|
|
|
logger = logging.getLogger("uvicorn.error")
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class MediaRegistration:
|
|
voice_session_id: str
|
|
call_id: str
|
|
interaction_id: str
|
|
ai_session_id: str | None = None
|
|
language: str | None = None
|
|
media_uuid: str | None = None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class MediaActor:
|
|
registration: MediaRegistration
|
|
reader: asyncio.StreamReader
|
|
writer: asyncio.StreamWriter
|
|
vad: EnergyVAD
|
|
frame_ms: int
|
|
frame_bytes: int
|
|
turn_queue: asyncio.Queue[tuple[bytes, bool]] = field(default_factory=asyncio.Queue)
|
|
worker_task: asyncio.Task | None = None
|
|
keepalive_task: asyncio.Task | None = None
|
|
handoff_task: asyncio.Task | None = None
|
|
playback_interrupt: asyncio.Event = field(default_factory=asyncio.Event)
|
|
write_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
|
closed: bool = False
|
|
state: str = "greeting"
|
|
barge_in_pending: bool = False
|
|
last_media_touch_monotonic: float = 0.0
|
|
last_outbound_audio_monotonic: float = 0.0
|
|
first_pcm_logged: bool = False
|
|
keepalive_loop_logged: bool = False
|
|
|
|
|
|
class AudioSocketMediaRuntime:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
enabled: bool,
|
|
host: str,
|
|
port: int,
|
|
frame_ms: int,
|
|
idle_timeout_seconds: float,
|
|
registration_wait_timeout_seconds: float,
|
|
min_speech_ms: int,
|
|
trailing_silence_ms: int,
|
|
max_turn_ms: int,
|
|
asr_provider: ASRProvider,
|
|
tts_provider: TTSProvider,
|
|
load_registration_by_media_uuid: Callable[[str], MediaRegistration | None],
|
|
mark_media_connected: Callable[[str, str], None],
|
|
mark_media_ended: Callable[[str, str], None],
|
|
touch_media_frame: Callable[[str], None],
|
|
set_state: Callable[[str, str, str | None, dict[str, Any] | None], None],
|
|
get_pending_greeting: Callable[[str], str | None],
|
|
mark_reply_delivered: Callable[[str, str, bool], None],
|
|
process_turn: Callable[[str, str, str | None, bool, dict[str, Any] | None], VoiceAITurnDecisionOut],
|
|
request_handoff: Callable[[str, str, VoiceAITurnDecisionOut], None],
|
|
handle_media_error: Callable[[str, str, dict[str, Any] | None], None],
|
|
) -> None:
|
|
self._enabled = enabled
|
|
self._host = host
|
|
self._port = port
|
|
self._frame_ms = max(frame_ms, 20)
|
|
self._frame_bytes = int((8000 * self._frame_ms / 1000.0) * 2)
|
|
self._idle_timeout_seconds = max(idle_timeout_seconds, 5.0)
|
|
self._registration_wait_timeout_seconds = max(registration_wait_timeout_seconds, 0.0)
|
|
self._outbound_keepalive_interval_seconds = 0.75
|
|
self._min_speech_ms = max(min_speech_ms, self._frame_ms)
|
|
self._trailing_silence_ms = max(trailing_silence_ms, self._frame_ms)
|
|
self._max_turn_ms = max(max_turn_ms, self._frame_ms)
|
|
self._asr_provider = asr_provider
|
|
self._tts_provider = tts_provider
|
|
self._load_registration_by_media_uuid = load_registration_by_media_uuid
|
|
self._mark_media_connected = mark_media_connected
|
|
self._mark_media_ended = mark_media_ended
|
|
self._touch_media_frame = touch_media_frame
|
|
self._set_state = set_state
|
|
self._get_pending_greeting = get_pending_greeting
|
|
self._mark_reply_delivered = mark_reply_delivered
|
|
self._process_turn = process_turn
|
|
self._request_handoff = request_handoff
|
|
self._handle_media_error = handle_media_error
|
|
self._server: asyncio.base_events.Server | None = None
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
self._actors: dict[str, MediaActor] = {}
|
|
|
|
@property
|
|
def address(self) -> str:
|
|
return f"{self._host}:{self._port}"
|
|
|
|
async def start(self) -> None:
|
|
if not self._enabled or self._server is not None:
|
|
return
|
|
self._loop = asyncio.get_running_loop()
|
|
self._server = await asyncio.start_server(self._handle_connection, self._host, self._port)
|
|
|
|
async def _await_registration(
|
|
self,
|
|
media_uuid: str,
|
|
*,
|
|
writer: asyncio.StreamWriter | None = None,
|
|
frame_bytes: int | None = None,
|
|
) -> MediaRegistration | None:
|
|
deadline = time.monotonic() + self._registration_wait_timeout_seconds
|
|
last_keepalive_at = 0.0
|
|
keepalive_payload = b"\x00" * max(frame_bytes or 0, 0)
|
|
keepalive_logged = False
|
|
while True:
|
|
registration = await asyncio.to_thread(self._load_registration_by_media_uuid, media_uuid)
|
|
if registration is not None:
|
|
return registration
|
|
if time.monotonic() >= deadline:
|
|
logger.warning("audiosocket.registration_timeout media_uuid=%s", media_uuid)
|
|
return None
|
|
if writer is not None and keepalive_payload:
|
|
now = time.monotonic()
|
|
if now - last_keepalive_at >= self._outbound_keepalive_interval_seconds:
|
|
writer.write(encode_audio_packet(keepalive_payload))
|
|
await writer.drain()
|
|
if not keepalive_logged:
|
|
keepalive_logged = True
|
|
logger.warning(
|
|
"audiosocket.prereg_keepalive media_uuid=%s frame_bytes=%s",
|
|
media_uuid,
|
|
len(keepalive_payload),
|
|
)
|
|
last_keepalive_at = now
|
|
await asyncio.sleep(0.05)
|
|
|
|
async def stop(self) -> None:
|
|
actor_ids = list(self._actors.keys())
|
|
for session_id in actor_ids:
|
|
await self._close_session(session_id, reason="runtime_shutdown")
|
|
if self._server is not None:
|
|
self._server.close()
|
|
await self._server.wait_closed()
|
|
self._server = None
|
|
|
|
def close_session_sync(self, session_id: str, *, reason: str) -> None:
|
|
if not self._loop:
|
|
return
|
|
future = asyncio.run_coroutine_threadsafe(
|
|
self._close_session(session_id, reason=reason),
|
|
self._loop,
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
future.result(timeout=5.0)
|
|
|
|
async def _handle_connection(
|
|
self,
|
|
reader: asyncio.StreamReader,
|
|
writer: asyncio.StreamWriter,
|
|
) -> None:
|
|
actor: MediaActor | None = None
|
|
close_reason = "connection_closed"
|
|
close_as_error = False
|
|
peer = writer.get_extra_info("peername")
|
|
try:
|
|
logger.warning("audiosocket.accept peer=%s", peer)
|
|
packet_type, payload = await read_packet(reader, timeout=self._idle_timeout_seconds)
|
|
logger.warning(
|
|
"audiosocket.initial_packet peer=%s packet_type=%s payload_len=%s",
|
|
peer,
|
|
packet_type,
|
|
len(payload),
|
|
)
|
|
if packet_type != AUDIO_SOCKET_PACKET_UUID:
|
|
raise RuntimeError("AudioSocket UUID handshake is required")
|
|
media_uuid = normalize_media_uuid(payload)
|
|
logger.warning("audiosocket.handshake peer=%s media_uuid=%s", peer, media_uuid)
|
|
registration = await self._await_registration(
|
|
media_uuid,
|
|
writer=writer,
|
|
frame_bytes=self._frame_bytes,
|
|
)
|
|
if registration is None:
|
|
raise RuntimeError(f"Unknown AudioSocket media_uuid: {media_uuid}")
|
|
logger.warning(
|
|
"audiosocket.registered peer=%s session_id=%s media_uuid=%s",
|
|
peer,
|
|
registration.voice_session_id,
|
|
media_uuid,
|
|
)
|
|
|
|
actor = MediaActor(
|
|
registration=registration,
|
|
reader=reader,
|
|
writer=writer,
|
|
vad=EnergyVAD(
|
|
frame_ms=self._frame_ms,
|
|
min_speech_ms=self._min_speech_ms,
|
|
trailing_silence_ms=self._trailing_silence_ms,
|
|
max_turn_ms=self._max_turn_ms,
|
|
),
|
|
frame_ms=self._frame_ms,
|
|
frame_bytes=self._frame_bytes,
|
|
)
|
|
previous_actor = self._actors.get(registration.voice_session_id)
|
|
if previous_actor is not None:
|
|
await self._cleanup_actor(previous_actor, reason="media_replaced", error=False)
|
|
self._actors[registration.voice_session_id] = actor
|
|
actor.last_outbound_audio_monotonic = time.monotonic()
|
|
await asyncio.to_thread(self._mark_media_connected, registration.voice_session_id, media_uuid)
|
|
actor.keepalive_task = asyncio.create_task(self._keepalive_loop(actor))
|
|
actor.worker_task = asyncio.create_task(self._worker(actor))
|
|
|
|
while not actor.closed:
|
|
packet_type, payload = await read_packet(reader, timeout=self._idle_timeout_seconds)
|
|
if packet_type == AUDIO_SOCKET_PACKET_PCM16:
|
|
await self._handle_pcm(actor, payload)
|
|
continue
|
|
if packet_type == AUDIO_SOCKET_PACKET_DTMF:
|
|
logger.warning(
|
|
"audiosocket.dtmf session_id=%s payload_len=%s",
|
|
actor.registration.voice_session_id,
|
|
len(payload),
|
|
)
|
|
continue
|
|
if packet_type == AUDIO_SOCKET_PACKET_HANGUP:
|
|
close_reason = "audiosocket_hangup"
|
|
break
|
|
if packet_type == AUDIO_SOCKET_PACKET_UUID:
|
|
logger.warning(
|
|
"audiosocket.extra_uuid session_id=%s payload_len=%s",
|
|
actor.registration.voice_session_id,
|
|
len(payload),
|
|
)
|
|
continue
|
|
raise RuntimeError(f"Unsupported AudioSocket packet type: {packet_type}")
|
|
except asyncio.TimeoutError:
|
|
close_reason = "media_idle_timeout"
|
|
close_as_error = True
|
|
logger.warning("audiosocket.timeout peer=%s", peer)
|
|
except asyncio.IncompleteReadError:
|
|
close_reason = "connection_closed"
|
|
logger.warning("audiosocket.peer_closed peer=%s", peer)
|
|
except Exception as exc:
|
|
close_reason = str(exc)[:1000] or "media_runtime_error"
|
|
close_as_error = True
|
|
logger.warning("audiosocket.error peer=%s reason=%s", peer, close_reason)
|
|
finally:
|
|
if actor is not None:
|
|
await self._cleanup_actor(actor, reason=close_reason, error=close_as_error)
|
|
else:
|
|
writer.close()
|
|
with contextlib.suppress(Exception):
|
|
await writer.wait_closed()
|
|
|
|
async def _handle_pcm(self, actor: MediaActor, pcm_frame: bytes) -> None:
|
|
if actor.closed or not pcm_frame:
|
|
return
|
|
if not actor.first_pcm_logged:
|
|
actor.first_pcm_logged = True
|
|
logger.warning(
|
|
"audiosocket.first_pcm session_id=%s state=%s frame_bytes=%s",
|
|
actor.registration.voice_session_id,
|
|
actor.state,
|
|
len(pcm_frame),
|
|
)
|
|
now = time.monotonic()
|
|
if now - actor.last_media_touch_monotonic >= 1.0:
|
|
actor.last_media_touch_monotonic = now
|
|
await asyncio.to_thread(self._touch_media_frame, actor.registration.voice_session_id)
|
|
|
|
if actor.state not in {"listening", "speaking"}:
|
|
return
|
|
|
|
vad_result = actor.vad.feed(pcm_frame)
|
|
if actor.state == "speaking" and vad_result.speech_started:
|
|
actor.playback_interrupt.set()
|
|
actor.barge_in_pending = True
|
|
if vad_result.utterance_pcm:
|
|
await actor.turn_queue.put((vad_result.utterance_pcm, actor.barge_in_pending))
|
|
actor.barge_in_pending = False
|
|
|
|
async def _worker(self, actor: MediaActor) -> None:
|
|
greeting_text = await asyncio.to_thread(
|
|
self._get_pending_greeting,
|
|
actor.registration.voice_session_id,
|
|
)
|
|
if greeting_text:
|
|
logger.warning(
|
|
"audiosocket.greeting session_id=%s text_len=%s",
|
|
actor.registration.voice_session_id,
|
|
len(greeting_text),
|
|
)
|
|
await self._speak_text(actor, greeting_text, is_greeting=True)
|
|
if not actor.closed:
|
|
await self._set_actor_state(actor, "listening")
|
|
|
|
while not actor.closed:
|
|
pcm_bytes, barge_in = await actor.turn_queue.get()
|
|
if actor.closed:
|
|
break
|
|
await self._process_utterance(actor, pcm_bytes, barge_in)
|
|
|
|
async def _process_utterance(self, actor: MediaActor, pcm_bytes: bytes, barge_in: bool) -> None:
|
|
await self._set_actor_state(actor, "thinking")
|
|
wav_bytes = pcm16le_to_wav_bytes(pcm_bytes, sample_rate_hz=8000)
|
|
transcription = await asyncio.to_thread(
|
|
lambda: self._asr_provider.transcribe(
|
|
wav_bytes,
|
|
language_hint=actor.registration.language,
|
|
)
|
|
)
|
|
transcript_text = str(transcription.text or "").strip()
|
|
if not transcript_text:
|
|
await self._set_actor_state(actor, "listening")
|
|
return
|
|
|
|
metadata = {
|
|
"turn_duration_ms": int(len(pcm_bytes) / 16),
|
|
"media_uuid": actor.registration.media_uuid,
|
|
}
|
|
decision = await asyncio.to_thread(
|
|
self._process_turn,
|
|
actor.registration.voice_session_id,
|
|
transcript_text,
|
|
transcription.language or actor.registration.language,
|
|
barge_in,
|
|
metadata,
|
|
)
|
|
handoff_task: asyncio.Task | None = None
|
|
if decision.needs_handoff:
|
|
await self._set_actor_state(actor, "handoff_requested", decision.handoff_reason)
|
|
handoff_task = self._start_handoff_request(actor, transcript_text, decision)
|
|
if decision.reply_text:
|
|
await self._speak_text(actor, decision.reply_text, is_greeting=False)
|
|
if actor.closed:
|
|
return
|
|
if decision.needs_handoff:
|
|
if handoff_task is not None and handoff_task.done():
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await handoff_task
|
|
return
|
|
await self._set_actor_state(actor, "listening")
|
|
|
|
def _start_handoff_request(
|
|
self,
|
|
actor: MediaActor,
|
|
customer_request_text: str,
|
|
decision: VoiceAITurnDecisionOut,
|
|
) -> asyncio.Task:
|
|
existing = actor.handoff_task
|
|
if existing is not None and not existing.done():
|
|
return existing
|
|
|
|
async def _run() -> None:
|
|
try:
|
|
await asyncio.to_thread(
|
|
self._request_handoff,
|
|
actor.registration.voice_session_id,
|
|
customer_request_text,
|
|
decision,
|
|
)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"audiosocket.handoff_request_failed session_id=%s error=%s",
|
|
actor.registration.voice_session_id,
|
|
str(exc)[:500],
|
|
)
|
|
|
|
actor.handoff_task = asyncio.create_task(_run())
|
|
return actor.handoff_task
|
|
|
|
async def _speak_text(self, actor: MediaActor, text: str, *, is_greeting: bool) -> None:
|
|
if actor.closed or not text:
|
|
return
|
|
await self._set_actor_state(actor, "speaking")
|
|
synth_started_at = time.monotonic()
|
|
synthesis = await asyncio.to_thread(
|
|
lambda: self._tts_provider.synthesize(
|
|
text,
|
|
language=actor.registration.language,
|
|
)
|
|
)
|
|
if not synthesis.audio_bytes:
|
|
raise RuntimeError("TTS provider returned empty audio")
|
|
logger.warning(
|
|
"audiosocket.tts_ready session_id=%s greeting=%s synth_ms=%s audio_bytes=%s",
|
|
actor.registration.voice_session_id,
|
|
is_greeting,
|
|
int((time.monotonic() - synth_started_at) * 1000),
|
|
len(synthesis.audio_bytes),
|
|
)
|
|
|
|
pcm_8k = resample_pcm16le(
|
|
synthesis.audio_bytes,
|
|
input_rate_hz=synthesis.sample_rate_hz,
|
|
output_rate_hz=8000,
|
|
)
|
|
first_frame_sent = False
|
|
for frame in chunk_audio(pcm_8k, frame_bytes=actor.frame_bytes):
|
|
if actor.closed or actor.playback_interrupt.is_set():
|
|
break
|
|
await self._write_audio_packet(actor, frame)
|
|
if not first_frame_sent:
|
|
first_frame_sent = True
|
|
logger.warning(
|
|
"audiosocket.first_frame session_id=%s greeting=%s frame_bytes=%s",
|
|
actor.registration.voice_session_id,
|
|
is_greeting,
|
|
len(frame),
|
|
)
|
|
await asyncio.sleep(actor.frame_ms / 1000.0)
|
|
|
|
interrupted = actor.playback_interrupt.is_set()
|
|
actor.playback_interrupt.clear()
|
|
if not interrupted and not actor.closed:
|
|
await asyncio.to_thread(
|
|
self._mark_reply_delivered,
|
|
actor.registration.voice_session_id,
|
|
text,
|
|
is_greeting,
|
|
)
|
|
logger.warning(
|
|
"audiosocket.reply_delivered session_id=%s greeting=%s",
|
|
actor.registration.voice_session_id,
|
|
is_greeting,
|
|
)
|
|
|
|
async def _write_audio_packet(self, actor: MediaActor, pcm_frame: bytes) -> None:
|
|
if actor.closed:
|
|
return
|
|
async with actor.write_lock:
|
|
if actor.closed:
|
|
return
|
|
actor.writer.write(encode_audio_packet(pcm_frame))
|
|
await actor.writer.drain()
|
|
actor.last_outbound_audio_monotonic = time.monotonic()
|
|
|
|
async def _keepalive_loop(self, actor: MediaActor) -> None:
|
|
silence_frame = b"\x00" * actor.frame_bytes
|
|
while not actor.closed:
|
|
await asyncio.sleep(0.25)
|
|
if actor.closed:
|
|
break
|
|
if (time.monotonic() - actor.last_outbound_audio_monotonic) < self._outbound_keepalive_interval_seconds:
|
|
continue
|
|
if not actor.keepalive_loop_logged:
|
|
actor.keepalive_loop_logged = True
|
|
logger.warning(
|
|
"audiosocket.keepalive_loop session_id=%s state=%s frame_bytes=%s",
|
|
actor.registration.voice_session_id,
|
|
actor.state,
|
|
actor.frame_bytes,
|
|
)
|
|
await self._write_audio_packet(actor, silence_frame)
|
|
|
|
async def _set_actor_state(
|
|
self,
|
|
actor: MediaActor,
|
|
state: str,
|
|
handoff_reason: str | None = None,
|
|
) -> None:
|
|
if actor.closed and state not in {"error", "closed"}:
|
|
return
|
|
if actor.state == state and (handoff_reason or None) is None:
|
|
return
|
|
logger.warning(
|
|
"audiosocket.state session_id=%s from_state=%s to_state=%s handoff_reason=%s",
|
|
actor.registration.voice_session_id,
|
|
actor.state,
|
|
state,
|
|
handoff_reason,
|
|
)
|
|
actor.state = state
|
|
await asyncio.to_thread(
|
|
self._set_state,
|
|
actor.registration.voice_session_id,
|
|
state,
|
|
handoff_reason,
|
|
{"media_uuid": actor.registration.media_uuid},
|
|
)
|
|
|
|
async def _close_session(self, session_id: str, *, reason: str) -> None:
|
|
actor = self._actors.get(session_id)
|
|
if actor is None:
|
|
return
|
|
await self._cleanup_actor(actor, reason=reason, error=False)
|
|
|
|
async def _cleanup_actor(self, actor: MediaActor, *, reason: str, error: bool) -> None:
|
|
if actor.closed:
|
|
return
|
|
actor.closed = True
|
|
logger.warning(
|
|
"audiosocket.cleanup session_id=%s reason=%s error=%s first_pcm=%s",
|
|
actor.registration.voice_session_id,
|
|
reason,
|
|
error,
|
|
actor.first_pcm_logged,
|
|
)
|
|
self._actors.pop(actor.registration.voice_session_id, None)
|
|
if actor.keepalive_task is not None:
|
|
actor.keepalive_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await actor.keepalive_task
|
|
if actor.handoff_task is not None:
|
|
actor.handoff_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await actor.handoff_task
|
|
if actor.worker_task is not None:
|
|
actor.worker_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await actor.worker_task
|
|
actor.writer.close()
|
|
with contextlib.suppress(Exception):
|
|
await actor.writer.wait_closed()
|
|
|
|
if error:
|
|
await asyncio.to_thread(
|
|
self._handle_media_error,
|
|
actor.registration.voice_session_id,
|
|
reason,
|
|
{"media_uuid": actor.registration.media_uuid},
|
|
)
|
|
return
|
|
await asyncio.to_thread(self._mark_media_ended, actor.registration.voice_session_id, reason)
|