diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/README.md b/ai_agents/agents/ten_packages/extension/conversation_recorder/README.md index 89c0339e57..342ace58bc 100644 --- a/ai_agents/agents/ten_packages/extension/conversation_recorder/README.md +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/README.md @@ -31,6 +31,7 @@ boto3 # Only needed for S3 storage | `filename` | string | auto-generated | Custom filename (works with all storage modes) | | `start_trigger` | string | `"on_user_joined"` | When to start recording: `"on_user_joined"` or `"on_start"` | | `sample_rate` | int | `24000` | Audio sample rate in Hz | +| `source_prebuffer_ms` | int | `120` | Per-source recording jitter buffer before audio is drained to the WAV | ### GCS Storage Properties diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/audio_mixer.py b/ai_agents/agents/ten_packages/extension/conversation_recorder/audio_mixer.py index 7af4f100d0..27d2f423c9 100644 --- a/ai_agents/agents/ten_packages/extension/conversation_recorder/audio_mixer.py +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/audio_mixer.py @@ -1,20 +1,37 @@ import numpy as np import threading +from dataclasses import dataclass, field from typing import Dict, Deque from collections import deque +@dataclass +class _SourceBuffer: + samples: Deque[float] = field(default_factory=deque) + ready: bool = False + + class AudioMixer: - def __init__(self, sample_rate=24000, channels=1, chunk_duration_ms=40): + def __init__( + self, + sample_rate=24000, + channels=1, + chunk_duration_ms=40, + source_prebuffer_ms=120, + ): self.sample_rate = sample_rate self.channels = channels # Calculate chunk size in samples. # e.g., 24000Hz * 0.04s = 960 samples self.chunk_size = int(self.sample_rate * (chunk_duration_ms / 1000.0)) + self.prebuffer_size = max( + self.chunk_size, + int(self.sample_rate * (source_prebuffer_ms / 1000.0)), + ) - # Buffers: map string(stream_id) -> Deque[float] (samples) + # Buffers: map string(stream_id) -> source playout state. # We use deque for efficient pop from left. - self.buffers: Dict[str, Deque[float]] = {} + self.buffers: Dict[str, _SourceBuffer] = {} self.lock = threading.Lock() def _resample( @@ -65,40 +82,63 @@ def push_audio( with self.lock: if source_id not in self.buffers: - self.buffers[source_id] = deque() + self.buffers[source_id] = _SourceBuffer() # Extend deque with samples - self.buffers[source_id].extend(audio_array) + self.buffers[source_id].samples.extend(audio_array) - def mix_next_chunk(self) -> bytes: + def mix_next_chunk(self, drain: bool = False) -> bytes: """ Extracts `chunk_size` samples from all buffers, mixes them, and returns bytes. - If a buffer has insufficient data, it contributes 0 (silence) for the missing part. + A source must have a small prebuffer before playout starts. This avoids + writing bursty model audio as alternating audio/silence fragments. + If a ready source underruns, it waits to rebuffer instead of consuming + a partial chunk. `drain=True` consumes partial tails during shutdown. If ALL buffers are empty, returns empty bytes (indicating no data to write). """ with self.lock: # Check if any buffer has data - has_data = any(len(buf) > 0 for buf in self.buffers.values()) + has_data = any( + len(source.samples) > 0 for source in self.buffers.values() + ) if not has_data: return b"" # Initialize mixer buffer mixed_chunk = np.zeros(self.chunk_size, dtype=np.float32) + consumed_any = False # Mix each source - for buf in self.buffers.values(): - # Extract up to chunk_size samples - count = min(len(buf), self.chunk_size) - if count > 0: - # Create temporary array from deque slice - # Iterating deque is fast enough for 960 items? - # Ideally we'd chunk this better, but for python MVP this is readable. - # Optimization: slice deque to list then array. - samples = [buf.popleft() for _ in range(count)] - samples_arr = np.array(samples, dtype=np.float32) - - # Add to mix - mixed_chunk[:count] += samples_arr + for source in self.buffers.values(): + available = len(source.samples) + if available == 0: + continue + + if drain: + count = min(available, self.chunk_size) + else: + if not source.ready: + if available < self.prebuffer_size: + continue + source.ready = True + + if available < self.chunk_size: + source.ready = False + continue + + count = self.chunk_size + + # Create temporary array from deque slice. This is small + # (usually 960 samples) and keeps the logic simple. + samples = [source.samples.popleft() for _ in range(count)] + samples_arr = np.array(samples, dtype=np.float32) + + # Add to mix + mixed_chunk[:count] += samples_arr + consumed_any = True + + if not consumed_any: + return b"" # Clip and convert to int16 mixed_chunk = np.clip(mixed_chunk, -32768, 32767) diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/extension.py b/ai_agents/agents/ten_packages/extension/conversation_recorder/extension.py index 7fd06a0bf2..8cb9fd941f 100644 --- a/ai_agents/agents/ten_packages/extension/conversation_recorder/extension.py +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/extension.py @@ -10,10 +10,11 @@ import json import signal import atexit +import urllib.request +import os from .audio_mixer import AudioMixer from .storage import StorageFactory - # Global registry of active recorders for signal handling _active_recorders = [] @@ -31,6 +32,22 @@ def _signal_handler(signum, _frame): raise SystemExit(128 + signum) +def _send_webhook_request(ten_env, url: str, payload: dict): + try: + data = json.dumps(payload).encode('utf-8') + req = urllib.request.Request( + url, + data=data, + headers={'Content-Type': 'application/json'}, + method='POST' + ) + with urllib.request.urlopen(req, timeout=5) as response: + res_body = response.read().decode('utf-8') + ten_env.log_info(f"Webhook response status: {response.status}, body: {res_body}") + except Exception as e: + ten_env.log_warn(f"Failed to send webhook: {e}") + + class ConversationRecorderExtension(AsyncExtension): def __init__(self, name: str): super().__init__(name) @@ -44,6 +61,11 @@ def __init__(self, name: str): self._flush_counter = 0 self._signals_registered = False + def _get_loop(self) -> asyncio.AbstractEventLoop: + if self.loop is None: + self.loop = asyncio.get_running_loop() + return self.loop + async def on_init(self, ten_env: AsyncTenEnv) -> None: ten_env.log_info("ConversationRecorderExtension on_init") config_json, _ = await ten_env.get_property_to_json() @@ -52,8 +74,12 @@ async def on_init(self, ten_env: AsyncTenEnv) -> None: # Get sample rate from config, default to 24000Hz (Gemini output rate) sample_rate = self.config.get("sample_rate", 24000) + source_prebuffer_ms = self.config.get("source_prebuffer_ms", 120) - self.mixer = AudioMixer(sample_rate=sample_rate) + self.mixer = AudioMixer( + sample_rate=sample_rate, + source_prebuffer_ms=source_prebuffer_ms, + ) self.storage = StorageFactory.create_storage( self.config.get("storage_type", "local"), self.config ) @@ -154,10 +180,11 @@ async def start_recording(self, ten_env: AsyncTenEnv): ten_env.log_info("Starting recording session...") self.is_recording = True + loop = self._get_loop() # Open storage in executor to avoid blocking if self.storage: - await self.loop.run_in_executor(None, self.storage.open) + await loop.run_in_executor(None, self.storage.open) if hasattr(self.storage, "actual_file_path"): ten_env.log_info( f"Recording to file: {self.storage.actual_file_path}" @@ -171,16 +198,73 @@ async def stop_recording(self, ten_env: AsyncTenEnv): ten_env.log_info("Stopping recording session...") self.is_recording = False + loop = self._get_loop() if self.recording_task: await self.recording_task self.recording_task = None if self.storage: + await self._drain_mixer(ten_env) file_path = getattr(self.storage, "actual_file_path", None) - await self.loop.run_in_executor(None, self.storage.close) + try: + await loop.run_in_executor(None, self.storage.close) + except Exception as err: + ten_env.log_error( + f"Failed to save recording at {file_path}: {err}" + ) + raise + file_path = getattr(self.storage, "actual_file_path", file_path) ten_env.log_info(f"Recording saved to: {file_path}") + webhook_url = self.config.get("webhook_url") or os.getenv( + "CONVERSATION_RECORD_WEBHOOK_URL" + ) + if webhook_url: + import time + + channel_name = self.config.get("channel", "") + payload = { + "channel_name": channel_name, + "status": "uploaded", + "file_path": file_path, + "timestamp": int(time.time() * 1000), + } + ten_env.log_info( + f"Sending recording upload webhook to {webhook_url} " + f"with payload {payload}" + ) + try: + await loop.run_in_executor( + None, + _send_webhook_request, + ten_env, + webhook_url, + payload, + ) + except Exception as webhook_err: + ten_env.log_warn( + f"Failed to await webhook request: {webhook_err}" + ) + + async def _drain_mixer(self, ten_env: AsyncTenEnv): + if not getattr(self, "mixer", None) or not self.storage: + return + + loop = self._get_loop() + while True: + mixed_bytes = self.mixer.mix_next_chunk(drain=True) + if not mixed_bytes: + break + try: + await loop.run_in_executor( + None, self.storage.write, mixed_bytes + ) + except Exception as e: + ten_env.log_error(f"Error draining recording mixer: {e}") + raise + async def _recording_loop(self, ten_env: AsyncTenEnv): + loop = self._get_loop() while self.is_recording: try: # Sleep approx one chunk duration (40ms) @@ -191,7 +275,7 @@ async def _recording_loop(self, ten_env: AsyncTenEnv): if mixed_bytes and self.storage: # Write in thread pool - await self.loop.run_in_executor( + await loop.run_in_executor( None, self.storage.write, mixed_bytes ) @@ -200,9 +284,7 @@ async def _recording_loop(self, ten_env: AsyncTenEnv): if self._flush_counter >= 25: self._flush_counter = 0 if self.storage: - await self.loop.run_in_executor( - None, self.storage.flush - ) + await loop.run_in_executor(None, self.storage.flush) except Exception as e: ten_env.log_error(f"Error in recording loop: {e}") diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/manifest.json b/ai_agents/agents/ten_packages/extension/conversation_recorder/manifest.json index b641172654..d43b354040 100644 --- a/ai_agents/agents/ten_packages/extension/conversation_recorder/manifest.json +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/manifest.json @@ -26,6 +26,12 @@ "api": { "property": { "properties": { + "webhook_url": { + "type": "string" + }, + "channel": { + "type": "string" + }, "storage_type": { "type": "string" }, @@ -41,6 +47,9 @@ "sample_rate": { "type": "int32" }, + "source_prebuffer_ms": { + "type": "int32" + }, "gcp_bucket_name": { "type": "string" }, @@ -88,4 +97,4 @@ "requirements.txt" ] } -} \ No newline at end of file +} diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/storage.py b/ai_agents/agents/ten_packages/extension/conversation_recorder/storage.py index c4d1e1a09b..ac3893619d 100644 --- a/ai_agents/agents/ten_packages/extension/conversation_recorder/storage.py +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/storage.py @@ -190,29 +190,26 @@ def close(self): # Upload to GCS if local_path and os.path.exists(local_path): - try: - _, bucket = self._get_gcs_client() - blob_name = self.upload_prefix.rstrip("/") - if blob_name: - blob_name = ( - f"{blob_name}/{os.path.basename(local_path)}" - ) - else: - blob_name = os.path.basename(local_path) - - blob = bucket.blob(blob_name) - blob.upload_from_filename(local_path) - self.actual_file_path = ( - f"gs://{self.bucket_name}/{blob_name}" + _, bucket = self._get_gcs_client() + blob_name = self.upload_prefix.rstrip("/") + if blob_name: + blob_name = ( + f"{blob_name}/{os.path.basename(local_path)}" ) - finally: - # Cleanup temp file - try: - os.remove(local_path) - if self.temp_dir: - os.rmdir(self.temp_dir) - except OSError: - pass + else: + blob_name = os.path.basename(local_path) + + blob = bucket.blob(blob_name) + blob.upload_from_filename(local_path) + self.actual_file_path = f"gs://{self.bucket_name}/{blob_name}" + + # Preserve the local recording if the upload raises. + try: + os.remove(local_path) + if self.temp_dir: + os.rmdir(self.temp_dir) + except OSError: + pass self.local_storage = None @@ -350,30 +347,29 @@ def close(self): # Upload to S3 if local_path and os.path.exists(local_path): + s3_client = self._get_s3_client() + key = self.upload_prefix.rstrip("/") + if key: + key = f"{key}/{os.path.basename(local_path)}" + else: + key = os.path.basename(local_path) + + s3_client.upload_file(local_path, self.bucket_name, key) + + # Build the actual file path URL + if self.endpoint_url: + self.actual_file_path = ( + f"{self.endpoint_url}/{self.bucket_name}/{key}" + ) + else: + self.actual_file_path = f"s3://{self.bucket_name}/{key}" + + # Preserve the local recording if the upload raises. try: - s3_client = self._get_s3_client() - key = self.upload_prefix.rstrip("/") - if key: - key = f"{key}/{os.path.basename(local_path)}" - else: - key = os.path.basename(local_path) - - s3_client.upload_file(local_path, self.bucket_name, key) - - # Build the actual file path URL - if self.endpoint_url: - self.actual_file_path = ( - f"{self.endpoint_url}/{self.bucket_name}/{key}" - ) - else: - self.actual_file_path = f"s3://{self.bucket_name}/{key}" - finally: - # Cleanup temp file - try: - os.remove(local_path) - if self.temp_dir: - os.rmdir(self.temp_dir) - except OSError: - pass + os.remove(local_path) + if self.temp_dir: + os.rmdir(self.temp_dir) + except OSError: + pass self.local_storage = None diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/bin/start b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/bin/start new file mode 100755 index 0000000000..9a685c404d --- /dev/null +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/bin/start @@ -0,0 +1,9 @@ +#!/bin/bash + +set -e + +cd "$(dirname "${BASH_SOURCE[0]}")/../.." + +export PYTHONPATH=../../:.ten/app:.ten/app/ten_packages/system/ten_runtime_python/lib:.ten/app/ten_packages/system/ten_runtime_python/interface:$PYTHONPATH + +pytest -s tests/ "$@" diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_audio_mixer.py b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_audio_mixer.py new file mode 100644 index 0000000000..b9714b5b77 --- /dev/null +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_audio_mixer.py @@ -0,0 +1,83 @@ +import importlib.util +from pathlib import Path + +import numpy as np + + +audio_mixer_spec = importlib.util.spec_from_file_location( + "conversation_recorder_audio_mixer", + Path(__file__).parent.parent / "audio_mixer.py", +) +audio_mixer_module = importlib.util.module_from_spec(audio_mixer_spec) +audio_mixer_spec.loader.exec_module(audio_mixer_module) +AudioMixer = audio_mixer_module.AudioMixer + + +def pcm(samples): + return np.array(samples, dtype=np.int16).tobytes() + + +def samples(audio): + return np.frombuffer(audio, dtype=np.int16) + + +def test_mixer_prebuffers_source_before_consuming_partial_audio(): + mixer = AudioMixer( + sample_rate=1000, + chunk_duration_ms=40, + source_prebuffer_ms=120, + ) + + mixer.push_audio("assistant", pcm([100] * 60), 1000) + + assert mixer.mix_next_chunk() == b"" + assert len(mixer.buffers["assistant"].samples) == 60 + + mixer.push_audio("assistant", pcm([100] * 60), 1000) + mixed = samples(mixer.mix_next_chunk()) + + assert len(mixed) == 40 + assert np.all(mixed == 100) + assert len(mixer.buffers["assistant"].samples) == 80 + + +def test_mixer_waits_to_rebuffer_after_source_underrun(): + mixer = AudioMixer( + sample_rate=1000, + chunk_duration_ms=40, + source_prebuffer_ms=120, + ) + + mixer.push_audio("assistant", pcm([100] * 120), 1000) + + assert len(samples(mixer.mix_next_chunk())) == 40 + assert len(samples(mixer.mix_next_chunk())) == 40 + assert len(samples(mixer.mix_next_chunk())) == 40 + + mixer.push_audio("assistant", pcm([100] * 20), 1000) + + assert mixer.mix_next_chunk() == b"" + assert len(mixer.buffers["assistant"].samples) == 20 + + mixer.push_audio("assistant", pcm([100] * 100), 1000) + + assert len(samples(mixer.mix_next_chunk())) == 40 + + +def test_mixer_drain_consumes_partial_tail_with_trailing_silence(): + mixer = AudioMixer( + sample_rate=1000, + chunk_duration_ms=40, + source_prebuffer_ms=120, + ) + + mixer.push_audio("assistant", pcm([100] * 20), 1000) + + assert mixer.mix_next_chunk() == b"" + + mixed = samples(mixer.mix_next_chunk(drain=True)) + + assert len(mixed) == 40 + assert np.all(mixed[:20] == 100) + assert np.all(mixed[20:] == 0) + assert mixer.mix_next_chunk(drain=True) == b"" diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_extension.py b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_extension.py new file mode 100644 index 0000000000..c9ca790d62 --- /dev/null +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_extension.py @@ -0,0 +1,49 @@ +import asyncio + +from conversation_recorder.extension import ConversationRecorderExtension + + +class FakeTenEnv: + def log_info(self, _message: str) -> None: + pass + + def log_error(self, _message: str) -> None: + pass + + +class FakeStorage: + actual_file_path = "/tmp/test.wav" + + def __init__(self) -> None: + self.opened = False + self.closed = False + + def open(self) -> None: + self.opened = True + + def close(self) -> None: + self.closed = True + + +def test_start_recording_before_on_start_initializes_event_loop() -> None: + async def run_scenario() -> None: + extension = ConversationRecorderExtension.__new__( + ConversationRecorderExtension + ) + storage = FakeStorage() + extension.is_recording = False + extension.recording_task = None + extension.loop = None + extension.storage = storage + env = FakeTenEnv() + + assert extension.loop is None + + await extension.start_recording(env) + await extension.stop_recording(env) + + assert extension.loop is asyncio.get_running_loop() + assert storage.opened is True + assert storage.closed is True + + asyncio.run(run_scenario()) diff --git a/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_storage.py b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_storage.py new file mode 100644 index 0000000000..50fa18443e --- /dev/null +++ b/ai_agents/agents/ten_packages/extension/conversation_recorder/tests/test_storage.py @@ -0,0 +1,79 @@ +import importlib.util +import os +import shutil +from pathlib import Path + +import pytest + + +storage_spec = importlib.util.spec_from_file_location( + "conversation_recorder_storage", + Path(__file__).parent.parent / "storage.py", +) +storage_module = importlib.util.module_from_spec(storage_spec) +storage_spec.loader.exec_module(storage_module) +GCSStorage = storage_module.GCSStorage + + +class FakeBlob: + def __init__(self, error=None): + self.error = error + self.uploaded_path = None + + def upload_from_filename(self, path): + self.uploaded_path = path + if self.error: + raise self.error + + +class FakeBucket: + def __init__(self, blob): + self.test_blob = blob + self.requested_name = None + + def blob(self, name): + self.requested_name = name + return self.test_blob + + +def create_storage(bucket): + storage = GCSStorage( + bucket_name="recordings-bucket", + filename="audio.wav", + upload_prefix="sessions/test/", + ) + storage._get_gcs_client = lambda: (None, bucket) + storage.open() + storage.write(b"\x00\x00") + return storage + + +def test_gcs_close_reports_remote_path_and_removes_uploaded_temp_file(): + blob = FakeBlob() + bucket = FakeBucket(blob) + storage = create_storage(bucket) + local_path = storage.actual_file_path + + storage.close() + + assert bucket.requested_name == "sessions/test/audio.wav" + assert blob.uploaded_path == local_path + assert storage.actual_file_path == ( + "gs://recordings-bucket/sessions/test/audio.wav" + ) + assert not os.path.exists(local_path) + + +def test_gcs_close_keeps_local_file_when_upload_fails(): + bucket = FakeBucket(FakeBlob(RuntimeError("upload rejected"))) + storage = create_storage(bucket) + local_path = storage.actual_file_path + + try: + with pytest.raises(RuntimeError, match="upload rejected"): + storage.close() + + assert storage.actual_file_path == local_path + assert os.path.exists(local_path) + finally: + shutil.rmtree(storage.temp_dir)